<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package go_kafka_client

import (
	"time"
	"math"
	"fmt"
	"sync"
	metrics "github.com/rcrowley/go-metrics"
)

type WorkerManager struct {
	Id string
	Config *ConsumerConfig
	Strategy          WorkerStrategy
	FailureHook       FailedCallback
	FailedAttemptHook FailedAttemptCallback
	Workers           []*Worker
	AvailableWorkers  chan *Worker
	CurrentBatch map[TaskId]*Task //TODO inspect for race conditions
	InputChannel      chan [] *Message
	TopicPartition TopicAndPartition
	LargestOffset int64
	lastCommittedOffset int64
	FailCounter *FailureCounter
	batchProcessed    chan bool
	stopLock          sync.Mutex
	managerStop       chan bool
	processingStop    chan bool
	commitStop 		chan bool

	activeWorkersCounter metrics.Counter
	pendingTasksCounter metrics.Counter
	batchDurationTimer metrics.Timer
	idleTimer metrics.Timer
}

func NewWorkerManager(id string, config *ConsumerConfig, topicPartition TopicAndPartition,
					wmsIdleTimer metrics.Timer, batchDurationTimer metrics.Timer, activeWorkersCounter metrics.Counter,
					pendingWMsTasksCounter metrics.Counter) *WorkerManager {
	workers := make([]*Worker, config.NumWorkers)
	availableWorkers := make(chan *Worker, config.NumWorkers)
	for i := 0; i < config.NumWorkers; i++ {
		workers[i] = &Worker{
			OutputChannel: make(chan WorkerResult),
			TaskTimeout: config.WorkerTaskTimeout,
		}
		availableWorkers <- workers[i]
	}

	return &WorkerManager {
		Id: id,
		Config: config,
		Strategy: config.Strategy,
		FailureHook: config.WorkerFailureCallback,
		FailedAttemptHook: config.WorkerFailedAttemptCallback,
		AvailableWorkers: availableWorkers,
		Workers: workers,
		InputChannel: make(chan [] *Message),
		CurrentBatch: make(map[TaskId]*Task),
		TopicPartition: topicPartition,
		LargestOffset: 0,
		FailCounter: NewFailureCounter(config.WorkerRetryThreshold, config.WorkerThresholdTimeWindow),
		batchProcessed: make(chan bool),
		managerStop: make(chan bool),
		processingStop: make(chan bool),
		commitStop: make(chan bool),
		activeWorkersCounter: activeWorkersCounter,
		pendingTasksCounter: pendingWMsTasksCounter,
		batchDurationTimer: batchDurationTimer,
		idleTimer: wmsIdleTimer,
	}
}

func (wm *WorkerManager) String() string {
	return wm.Id
}

func (wm *WorkerManager) Start() {
	go wm.processBatch()
	go wm.commitBatch()
	for {
		startIdle := time.Now()
		select {
		case batch := <-wm.InputChannel: {
			wm.idleTimer.Update(time.Since(startIdle))
			Debug(wm, "WorkerManager got batch")
			wm.batchDurationTimer.Time(func() {
				wm.startBatch(batch)
			})
			Debug(wm, "WorkerManager got batch processed")
		}
		case <-wm.managerStop: return
		}
	}
}

func (wm *WorkerManager) Stop() chan bool {
	finished := make(chan bool)
	go func() {
		Debugf(wm, "Trying to stop workerManager")
		inLock(&wm.stopLock, func() {
				Debug(wm, "Stopping manager")
				wm.managerStop <- true
				Debug(wm, "Stopping processor")
				wm.processingStop <- true
				Debug(wm, "Successful manager stop")
				Debug(wm, "Stopping committer")
				wm.commitStop <- true
				Debug(wm, "Successful committer stop")
				finished <- true
				Debug(wm, "Leaving manager stop")
			})
		Debugf(wm, "Stopped workerManager")
	}()

	return finished
}

func (wm *WorkerManager) startBatch(batch []*Message) {
	inLock(&wm.stopLock, func() {
			for _, message := range batch {
				topicPartition := TopicAndPartition{ message.Topic, message.Partition }
				wm.CurrentBatch[TaskId{ topicPartition, message.Offset }] = &Task{Msg: message,}
			}
			wm.pendingTasksCounter.Inc(int64(len(wm.CurrentBatch)))
			for _, task := range wm.CurrentBatch {
				worker := <-wm.AvailableWorkers
				wm.activeWorkersCounter.Inc(1)
				wm.pendingTasksCounter.Dec(1)
				worker.Start(task, wm.Strategy)
			}

			<-wm.batchProcessed
		})
}

func (wm *WorkerManager) commitBatch() {
	for {
		select {
		case <-wm.commitStop: {
			wm.commitOffset()
			return
		}
		//TODO I don't like that. think think think!
		case <-time.After(wm.Config.OffsetCommitInterval): {
			wm.commitOffset()
		}
		}
	}
}

func (wm *WorkerManager) commitOffset() {
	largestOffset := wm.LargestOffset
	Tracef(wm, "Inside commit offset with largest %d and last %d", largestOffset, wm.lastCommittedOffset)
	if largestOffset <= wm.lastCommittedOffset { return }

	success := false
	for i := 0; i <= wm.Config.OffsetsCommitMaxRetries; i++ {
		err := wm.Config.Coordinator.CommitOffset(wm.Config.Groupid, &wm.TopicPartition, largestOffset)
		if err == nil {
			success = true
			Debugf(wm, "Successfully committed offset %d for %s", largestOffset, wm.TopicPartition)
			break
		} else {
			Infof(wm, "Failed to commit offset %d for %s. Retying...", largestOffset, &wm.TopicPartition)
		}
	}

	if !success {
		Errorf(wm, "Failed to commit offset %d for %s after %d retries", largestOffset, &wm.TopicPartition, wm.Config.OffsetsCommitMaxRetries)
		//TODO: what to do next?
	} else {
		wm.lastCommittedOffset = largestOffset
	}
}

func (wm *WorkerManager) IsBatchProcessed() bool {
	return len(wm.CurrentBatch) == 0
}

func (wm *WorkerManager) processBatch() {
	outputChannels := make([]*chan WorkerResult, wm.Config.NumWorkers)
	for i, worker := range wm.Workers {
		outputChannels[i] = &worker.OutputChannel
	}

	resultsChannel := make(chan WorkerResult)
	for {
		stopRedirecting := redirectChannelsTo(outputChannels, resultsChannel)
		select {
		case result := <-resultsChannel: {
			go func() {
				stopRedirecting <- true
			}()

			task := wm.CurrentBatch[result.Id()]
			if result.Success() {
				wm.taskIsDone(result)
			} else {
				if _, ok := result.(*TimedOutResult); ok {
					task.Callee.OutputChannel = make(chan WorkerResult)
				}

				Warnf(wm, "Worker task %s has failed", result.Id())
				task.Retries++
				if task.Retries > wm.Config.MaxWorkerRetries {
					Errorf(wm, "Worker task %s has failed after %d retries", result.Id(), wm.Config.MaxWorkerRetries)

					var decision FailedDecision
					if wm.FailCounter.Failed() {
						decision = wm.FailureHook(wm)
					} else {
						decision = wm.FailedAttemptHook(task, result)
					}
					switch decision {
					case CommitOffsetAndContinue: {
						wm.taskIsDone(result)
					}
					case DoNotCommitOffsetAndContinue: {
						wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
						delete(wm.CurrentBatch, result.Id())
					}
					case CommitOffsetAndStop: {
						wm.taskIsDone(result)
						wm.stopBatch()
					}
					case DoNotCommitOffsetAndStop: {
						wm.stopBatch()
					}
					}
				} else {
					Warnf(wm, "Retrying worker task %s %dth time", result.Id(), task.Retries)
					time.Sleep(wm.Config.WorkerBackoff)
					go task.Callee.Start(task, wm.Strategy)
				}
			}

			if wm.IsBatchProcessed() {
				Debug(wm, "Sending batch processed")
				wm.batchProcessed <- true
				Debug(wm, "Received batch processed")
			}
		}
		case <-wm.processingStop: {
			go func() {
				stopRedirecting <- true
			}()
			return
		}
		}

	}
}

func (wm *WorkerManager) stopBatch() {
	wm.CurrentBatch = make(map[TaskId]*Task)
	for _, worker := range wm.Workers {
		worker.OutputChannel = make(chan WorkerResult)
	}
}

func (wm *WorkerManager) taskIsDone(result WorkerResult) {
	Tracef(wm, "Task is done: %d", result.Id().Offset)
	wm.LargestOffset = int64(math.Max(float64(wm.LargestOffset), float64(result.Id().Offset)))
	wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
	wm.activeWorkersCounter.Dec(1)
	delete(wm.CurrentBatch, result.Id())
}

type Worker struct {
	OutputChannel chan WorkerResult
	TaskTimeout   time.Duration
	Closed        bool
}

func (w *Worker) String() string {
	return "worker"
}

func (w *Worker) Start(task *Task, strategy WorkerStrategy) {
	task.Callee = w
	go func() {
		resultChannel := make(chan WorkerResult)
		go func() {resultChannel <- strategy(w, task.Msg, task.Id())}()
		select {
		case result := <-resultChannel: {
			w.OutputChannel <- result
		}
		case <-time.After(w.TaskTimeout): {
			w.OutputChannel <- &TimedOutResult{task.Id()}
		}
		}
	}()
}

type WorkerStrategy func(*Worker, *Message, TaskId) WorkerResult

type FailedCallback func(*WorkerManager) FailedDecision

type FailedAttemptCallback func(*Task, WorkerResult) FailedDecision

type FailureCounter struct {
	count           int32
	failed          bool
	countLock       sync.Mutex
	failedThreshold int32
}

func NewFailureCounter(FailedThreshold int32, WorkerThresholdTimeWindow time.Duration) *FailureCounter {
	counter := &FailureCounter {
		failedThreshold : FailedThreshold,
	}
	go func() {
		for {
			select {
			case <-time.After(WorkerThresholdTimeWindow): {
				if counter.count >= FailedThreshold {
					counter.failed = true
					return
				} else {
					inLock(&counter.countLock, func() {
						counter.count = 0
					})
				}
			}
			}
		}
	}()
	return counter
}

func (f *FailureCounter) Failed() bool {
	inLock(&f.countLock, func() { f.count++ })
	return f.count >= f.failedThreshold || f.failed
}

type Task struct {
	Msg *Message
	Retries int
	Callee *Worker
}

func (t *Task) Id() TaskId {
	return TaskId{TopicAndPartition{t.Msg.Topic, t.Msg.Partition}, t.Msg.Offset}
}

type WorkerResult interface {
	Id() TaskId
	Success() bool
}

type SuccessfulResult struct {
	id TaskId
}

func NewSuccessfulResult(id TaskId) *SuccessfulResult {
	return &SuccessfulResult{id}
}

func (sr *SuccessfulResult) String() string {
	return fmt.Sprintf("{Success: %s}", sr.Id())
}

func (wr *SuccessfulResult) Id() TaskId {
	return wr.id
}

func (wr *SuccessfulResult) Success() bool {
	return true
}

type ProcessingFailedResult struct {
	id TaskId
}

func NewProcessingFailedResult(id TaskId) *ProcessingFailedResult {
	return &ProcessingFailedResult{id}
}

func (sr *ProcessingFailedResult) String() string {
	return fmt.Sprintf("{Failed: %s}", sr.Id())
}

func (wr *ProcessingFailedResult) Id() TaskId {
	return wr.id
}

func (wr *ProcessingFailedResult) Success() bool {
	return false
}

type TimedOutResult struct {
	id TaskId
}

func (sr *TimedOutResult) String() string {
	return fmt.Sprintf("{Timed out: %s}", sr.Id())
}

func (wr *TimedOutResult) Id() TaskId {
	return wr.id
}

func (wr *TimedOutResult) Success() bool {
	return false
}

type TaskId struct {
	TopicPartition TopicAndPartition
	Offset         int64
}

func (tid TaskId) String() string {
	return fmt.Sprintf("%s, Offset: %d", &tid.TopicPartition, tid.Offset)
}

type FailedDecision int32

const (
	CommitOffsetAndContinue FailedDecision = iota
	DoNotCommitOffsetAndContinue
	CommitOffsetAndStop
	DoNotCommitOffsetAndStop
)
<commit_msg>Workers offset fix<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package go_kafka_client

import (
	"time"
	"math"
	"fmt"
	"sync"
	metrics "github.com/rcrowley/go-metrics"
)

type WorkerManager struct {
	Id string
	Config *ConsumerConfig
	Strategy          WorkerStrategy
	FailureHook       FailedCallback
	FailedAttemptHook FailedAttemptCallback
	Workers           []*Worker
	AvailableWorkers  chan *Worker
	CurrentBatch map[TaskId]*Task //TODO inspect for race conditions
	InputChannel      chan [] *Message
	TopicPartition TopicAndPartition
	LargestOffset int64
	lastCommittedOffset int64
	FailCounter *FailureCounter
	batchProcessed    chan bool
	stopLock          sync.Mutex
	managerStop       chan bool
	processingStop    chan bool
	commitStop 		chan bool

	activeWorkersCounter metrics.Counter
	pendingTasksCounter metrics.Counter
	batchDurationTimer metrics.Timer
	idleTimer metrics.Timer
}

func NewWorkerManager(id string, config *ConsumerConfig, topicPartition TopicAndPartition,
					wmsIdleTimer metrics.Timer, batchDurationTimer metrics.Timer, activeWorkersCounter metrics.Counter,
					pendingWMsTasksCounter metrics.Counter) *WorkerManager {
	workers := make([]*Worker, config.NumWorkers)
	availableWorkers := make(chan *Worker, config.NumWorkers)
	for i := 0; i < config.NumWorkers; i++ {
		workers[i] = &Worker{
			OutputChannel: make(chan WorkerResult),
			TaskTimeout: config.WorkerTaskTimeout,
		}
		availableWorkers <- workers[i]
	}

	return &WorkerManager {
		Id: id,
		Config: config,
		Strategy: config.Strategy,
		FailureHook: config.WorkerFailureCallback,
		FailedAttemptHook: config.WorkerFailedAttemptCallback,
		AvailableWorkers: availableWorkers,
		Workers: workers,
		InputChannel: make(chan [] *Message),
		CurrentBatch: make(map[TaskId]*Task),
		TopicPartition: topicPartition,
		LargestOffset: InvalidOffset,
		FailCounter: NewFailureCounter(config.WorkerRetryThreshold, config.WorkerThresholdTimeWindow),
		batchProcessed: make(chan bool),
		managerStop: make(chan bool),
		processingStop: make(chan bool),
		commitStop: make(chan bool),
		activeWorkersCounter: activeWorkersCounter,
		pendingTasksCounter: pendingWMsTasksCounter,
		batchDurationTimer: batchDurationTimer,
		idleTimer: wmsIdleTimer,
	}
}

func (wm *WorkerManager) String() string {
	return wm.Id
}

func (wm *WorkerManager) Start() {
	go wm.processBatch()
	go wm.commitBatch()
	for {
		startIdle := time.Now()
		select {
		case batch := <-wm.InputChannel: {
			wm.idleTimer.Update(time.Since(startIdle))
			Debug(wm, "WorkerManager got batch")
			wm.batchDurationTimer.Time(func() {
				wm.startBatch(batch)
			})
			Debug(wm, "WorkerManager got batch processed")
		}
		case <-wm.managerStop: return
		}
	}
}

func (wm *WorkerManager) Stop() chan bool {
	finished := make(chan bool)
	go func() {
		Debugf(wm, "Trying to stop workerManager")
		inLock(&wm.stopLock, func() {
				Debug(wm, "Stopping manager")
				wm.managerStop <- true
				Debug(wm, "Stopping processor")
				wm.processingStop <- true
				Debug(wm, "Successful manager stop")
				Debug(wm, "Stopping committer")
				wm.commitStop <- true
				Debug(wm, "Successful committer stop")
				finished <- true
				Debug(wm, "Leaving manager stop")
			})
		Debugf(wm, "Stopped workerManager")
	}()

	return finished
}

func (wm *WorkerManager) startBatch(batch []*Message) {
	inLock(&wm.stopLock, func() {
			for _, message := range batch {
				topicPartition := TopicAndPartition{ message.Topic, message.Partition }
				wm.CurrentBatch[TaskId{ topicPartition, message.Offset }] = &Task{Msg: message,}
			}
			wm.pendingTasksCounter.Inc(int64(len(wm.CurrentBatch)))
			for _, task := range wm.CurrentBatch {
				worker := <-wm.AvailableWorkers
				wm.activeWorkersCounter.Inc(1)
				wm.pendingTasksCounter.Dec(1)
				worker.Start(task, wm.Strategy)
			}

			<-wm.batchProcessed
		})
}

func (wm *WorkerManager) commitBatch() {
	for {
		select {
		case <-wm.commitStop: {
			wm.commitOffset()
			return
		}
		//TODO I don't like that. think think think!
		case <-time.After(wm.Config.OffsetCommitInterval): {
			wm.commitOffset()
		}
		}
	}
}

func (wm *WorkerManager) commitOffset() {
	largestOffset := wm.LargestOffset
	Tracef(wm, "Inside commit offset with largest %d and last %d", largestOffset, wm.lastCommittedOffset)
	if largestOffset <= wm.lastCommittedOffset || isOffsetInvalid(largestOffset) { return }

	success := false
	for i := 0; i <= wm.Config.OffsetsCommitMaxRetries; i++ {
		err := wm.Config.Coordinator.CommitOffset(wm.Config.Groupid, &wm.TopicPartition, largestOffset)
		if err == nil {
			success = true
			Debugf(wm, "Successfully committed offset %d for %s", largestOffset, wm.TopicPartition)
			break
		} else {
			Infof(wm, "Failed to commit offset %d for %s. Retying...", largestOffset, &wm.TopicPartition)
		}
	}

	if !success {
		Errorf(wm, "Failed to commit offset %d for %s after %d retries", largestOffset, &wm.TopicPartition, wm.Config.OffsetsCommitMaxRetries)
		//TODO: what to do next?
	} else {
		wm.lastCommittedOffset = largestOffset
	}
}

func (wm *WorkerManager) IsBatchProcessed() bool {
	return len(wm.CurrentBatch) == 0
}

func (wm *WorkerManager) processBatch() {
	outputChannels := make([]*chan WorkerResult, wm.Config.NumWorkers)
	for i, worker := range wm.Workers {
		outputChannels[i] = &worker.OutputChannel
	}

	resultsChannel := make(chan WorkerResult)
	for {
		stopRedirecting := redirectChannelsTo(outputChannels, resultsChannel)
		select {
		case result := <-resultsChannel: {
			go func() {
				stopRedirecting <- true
			}()

			task := wm.CurrentBatch[result.Id()]
			if result.Success() {
				wm.taskIsDone(result)
			} else {
				if _, ok := result.(*TimedOutResult); ok {
					task.Callee.OutputChannel = make(chan WorkerResult)
				}

				Warnf(wm, "Worker task %s has failed", result.Id())
				task.Retries++
				if task.Retries > wm.Config.MaxWorkerRetries {
					Errorf(wm, "Worker task %s has failed after %d retries", result.Id(), wm.Config.MaxWorkerRetries)

					var decision FailedDecision
					if wm.FailCounter.Failed() {
						decision = wm.FailureHook(wm)
					} else {
						decision = wm.FailedAttemptHook(task, result)
					}
					switch decision {
					case CommitOffsetAndContinue: {
						wm.taskIsDone(result)
					}
					case DoNotCommitOffsetAndContinue: {
						wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
						delete(wm.CurrentBatch, result.Id())
					}
					case CommitOffsetAndStop: {
						wm.taskIsDone(result)
						wm.stopBatch()
					}
					case DoNotCommitOffsetAndStop: {
						wm.stopBatch()
					}
					}
				} else {
					Warnf(wm, "Retrying worker task %s %dth time", result.Id(), task.Retries)
					time.Sleep(wm.Config.WorkerBackoff)
					go task.Callee.Start(task, wm.Strategy)
				}
			}

			if wm.IsBatchProcessed() {
				Debug(wm, "Sending batch processed")
				wm.batchProcessed <- true
				Debug(wm, "Received batch processed")
			}
		}
		case <-wm.processingStop: {
			go func() {
				stopRedirecting <- true
			}()
			return
		}
		}

	}
}

func (wm *WorkerManager) stopBatch() {
	wm.CurrentBatch = make(map[TaskId]*Task)
	for _, worker := range wm.Workers {
		worker.OutputChannel = make(chan WorkerResult)
	}
}

func (wm *WorkerManager) taskIsDone(result WorkerResult) {
	Tracef(wm, "Task is done: %d", result.Id().Offset)
	wm.LargestOffset = int64(math.Max(float64(wm.LargestOffset), float64(result.Id().Offset)))
	wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
	wm.activeWorkersCounter.Dec(1)
	delete(wm.CurrentBatch, result.Id())
}

type Worker struct {
	OutputChannel chan WorkerResult
	TaskTimeout   time.Duration
	Closed        bool
}

func (w *Worker) String() string {
	return "worker"
}

func (w *Worker) Start(task *Task, strategy WorkerStrategy) {
	task.Callee = w
	go func() {
		resultChannel := make(chan WorkerResult)
		go func() {resultChannel <- strategy(w, task.Msg, task.Id())}()
		select {
		case result := <-resultChannel: {
			w.OutputChannel <- result
		}
		case <-time.After(w.TaskTimeout): {
			w.OutputChannel <- &TimedOutResult{task.Id()}
		}
		}
	}()
}

type WorkerStrategy func(*Worker, *Message, TaskId) WorkerResult

type FailedCallback func(*WorkerManager) FailedDecision

type FailedAttemptCallback func(*Task, WorkerResult) FailedDecision

type FailureCounter struct {
	count           int32
	failed          bool
	countLock       sync.Mutex
	failedThreshold int32
}

func NewFailureCounter(FailedThreshold int32, WorkerThresholdTimeWindow time.Duration) *FailureCounter {
	counter := &FailureCounter {
		failedThreshold : FailedThreshold,
	}
	go func() {
		for {
			select {
			case <-time.After(WorkerThresholdTimeWindow): {
				if counter.count >= FailedThreshold {
					counter.failed = true
					return
				} else {
					inLock(&counter.countLock, func() {
						counter.count = 0
					})
				}
			}
			}
		}
	}()
	return counter
}

func (f *FailureCounter) Failed() bool {
	inLock(&f.countLock, func() { f.count++ })
	return f.count >= f.failedThreshold || f.failed
}

type Task struct {
	Msg *Message
	Retries int
	Callee *Worker
}

func (t *Task) Id() TaskId {
	return TaskId{TopicAndPartition{t.Msg.Topic, t.Msg.Partition}, t.Msg.Offset}
}

type WorkerResult interface {
	Id() TaskId
	Success() bool
}

type SuccessfulResult struct {
	id TaskId
}

func NewSuccessfulResult(id TaskId) *SuccessfulResult {
	return &SuccessfulResult{id}
}

func (sr *SuccessfulResult) String() string {
	return fmt.Sprintf("{Success: %s}", sr.Id())
}

func (wr *SuccessfulResult) Id() TaskId {
	return wr.id
}

func (wr *SuccessfulResult) Success() bool {
	return true
}

type ProcessingFailedResult struct {
	id TaskId
}

func NewProcessingFailedResult(id TaskId) *ProcessingFailedResult {
	return &ProcessingFailedResult{id}
}

func (sr *ProcessingFailedResult) String() string {
	return fmt.Sprintf("{Failed: %s}", sr.Id())
}

func (wr *ProcessingFailedResult) Id() TaskId {
	return wr.id
}

func (wr *ProcessingFailedResult) Success() bool {
	return false
}

type TimedOutResult struct {
	id TaskId
}

func (sr *TimedOutResult) String() string {
	return fmt.Sprintf("{Timed out: %s}", sr.Id())
}

func (wr *TimedOutResult) Id() TaskId {
	return wr.id
}

func (wr *TimedOutResult) Success() bool {
	return false
}

type TaskId struct {
	TopicPartition TopicAndPartition
	Offset         int64
}

func (tid TaskId) String() string {
	return fmt.Sprintf("%s, Offset: %d", &tid.TopicPartition, tid.Offset)
}

type FailedDecision int32

const (
	CommitOffsetAndContinue FailedDecision = iota
	DoNotCommitOffsetAndContinue
	CommitOffsetAndStop
	DoNotCommitOffsetAndStop
)
<|endoftext|>


<commit_before>package psdock

import (
	"bytes"
	"code.google.com/p/go.crypto/ssh/terminal"
	"errors"
	"github.com/kr/pty"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"syscall"
	"time"
)

type Process struct {
	Cmd           *exec.Cmd
	Conf          *Config
	Pty           *os.File
	Term          *terminal.Terminal
	Status        int
	StatusChannel chan ProcessStatus
	oldTermState  *terminal.State
}

//NewProcess creates a new struct of type *Process and returns its address
func NewProcess(conf *Config) *Process {
	var cmd *exec.Cmd
	if len(conf.Args) > 0 {
		cmd = exec.Command(conf.Command, strings.Split(conf.Args, " ")...)
	} else {
		cmd = exec.Command(conf.Command)
	}
	newStatusChannel := make(chan ProcessStatus, 1)
	return &Process{Cmd: cmd, Conf: conf, StatusChannel: newStatusChannel}
}

//SetEnvVars sets the environment variables for the launched process
func (p *Process) SetEnvVars() {
	if len(p.Conf.EnvVars) == 0 {
		//If we do not want to pass any env var to the child process, we still have
		//to write something in p.Cmd.Env so that it is non-empty
		p.Cmd.Env = append(p.Cmd.Env, " ")
		return
	}
	for _, envVar := range strings.Split(p.Conf.EnvVars, ",") {
		p.Cmd.Env = append(p.Cmd.Env, envVar)
	}
}

//SetUser tries to change the current user to newUsername
func (p *Process) SetUser() error {
	currentUser, err := user.Current()
	if err != nil {
		return errors.New("Can't determine the current user !\n" + err.Error())
	}

	if p.Conf.UserName == currentUser.Username {
		return nil
	}

	newUser, err := user.Lookup(p.Conf.UserName)
	if err != nil {
		return errors.New("Can't find the user" + p.Conf.UserName + "!\n" + err.Error())
	}

	newUserUID, err := strconv.Atoi(newUser.Uid)
	if err != nil {
		return errors.New("Can't determine the new user UID !\n" + err.Error())
	}

	if err := syscall.Setuid(newUserUID); err != nil {
		return errors.New("Can't change the user !\n" + err.Error())
	}

	return nil
}

func (p *Process) Terminate(maxTryCount int) error {
	if maxTryCount > 0 {
		if err := syscall.Kill(p.Cmd.Process.Pid, syscall.SIGTERM); err == nil {
			return nil
		}
	} else {
		return syscall.Kill(p.Cmd.Process.Pid, syscall.SIGKILL)
	}
	time.Sleep(100 * time.Millisecond)
	return p.Terminate(maxTryCount - 1)
}

func (p *Process) isStarted() bool {
	return p.Cmd.Process != nil
}

func (p *Process) isRunning() bool {
	return p.isStarted() && p.hasBoundPort()
}

func (p *Process) hasBoundPort() bool {
	if p.Conf.BindPort == 0 {
		return true
	}

	//We execute netstat -an | grep bindPort to find if bindPort is open
	netstCmd := exec.Command("netstat", "-an")
	grepCmd := exec.Command("grep", string(p.Conf.BindPort))
	netstOut, _ := netstCmd.Output()
	grepCmd.Stdin = bytes.NewBuffer(netstOut)
	grepOut, _ := grepCmd.Output()

	return len(grepOut) > 0
}

func (p *Process) redirectStdin() error {
	var err error
	p.oldTermState, err = terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return errors.New("Can't redirect stdin:" + err.Error())
	}
	newTerminal := terminal.NewTerminal(os.Stdin, "")
	cb := func(s string, i int, r rune) (string, int, bool) {
		car := []byte{byte(r)}
		newTerminal.Write(car)
		return s, i, false
	}
	newTerminal.AutoCompleteCallback = cb
	go io.Copy(p.Pty, os.Stdin)
	return nil
}

func (p *Process) restoreStdin() error {
	err := terminal.Restore(int(os.Stdin.Fd()), p.oldTermState)
	return err
}

func (p *Process) redirectStdout() error {
	go io.Copy(os.Stdout, p.Pty)
	return nil
}

func (p *Process) NotifyStatusChanged() error {
	if p.Conf.WebHook == "" {
		return nil
	}
	statusStr := ""
	if p.Status == PROCESS_STARTED {
		statusStr = "started"
	} else if p.Status == PROCESS_RUNNING {
		statusStr = "running"
	} else {
		statusStr = "stopped"
	}
	body := `{
							"ps":
								{ "status":` + statusStr + `}
						}`

	req, err := http.NewRequest("PUT", p.Conf.WebHook, bytes.NewBufferString(body))
	if err != nil {
		return errors.New("Failed to construct the HTTP request" + err.Error())
	}

	req.Header.Add("Content-Type", "application/json")
	client := &http.Client{}

	resp, err := client.Do(req)
	defer resp.Body.Close()
	if err != nil {
		return errors.New("Was not able to trigger the hook!\n" + err.Error())
	}
	return nil
}

func (p *Process) Start() error {
	//We start the process
	var startErr error
	p.Pty, startErr = pty.Start(p.Cmd)
	if startErr != nil {
		return startErr
	}

	go func() {
		var err error
		if err = p.redirectStdin(); err != nil {
			p.StatusChannel <- ProcessStatus{Status: -1, Err: err}
		}
		if err = p.redirectStdout(); err != nil {
			p.StatusChannel <- ProcessStatus{Status: -1, Err: err}
		}

		for !p.isStarted() {
			time.Sleep(100 * time.Millisecond)
		}
		p.Status = PROCESS_STARTED
		if err = p.NotifyStatusChanged(); err != nil {
			log.Println(err)
		}
		p.StatusChannel <- ProcessStatus{Status: p.Status, Err: nil}

		for p.isRunning() == false {
			time.Sleep(100 * time.Millisecond)
		}
		p.Status = PROCESS_RUNNING
		if err = p.NotifyStatusChanged(); err != nil {
			log.Println(err)
		}
		p.StatusChannel <- ProcessStatus{Status: p.Status, Err: nil}

		err = p.Cmd.Wait()
		if err != nil {
			p.restoreStdin()
			p.Status = PROCESS_STOPPED
			p.NotifyStatusChanged()
			log.Fatal(err)
		}

		//p has stopped
		p.restoreStdin()
		p.Status = PROCESS_STOPPED
		if err = p.NotifyStatusChanged(); err != nil {
			log.Println(err)
		}
		p.StatusChannel <- ProcessStatus{Status: p.Status, Err: nil}
	}()

	return nil
}
<commit_msg>Corrected env-vars-export behaviour<commit_after>package psdock

import (
	"bytes"
	"code.google.com/p/go.crypto/ssh/terminal"
	"errors"
	"github.com/kr/pty"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"syscall"
	"time"
)

type Process struct {
	Cmd           *exec.Cmd
	Conf          *Config
	Pty           *os.File
	Term          *terminal.Terminal
	Status        int
	StatusChannel chan ProcessStatus
	oldTermState  *terminal.State
}

//NewProcess creates a new struct of type *Process and returns its address
func NewProcess(conf *Config) *Process {
	var cmd *exec.Cmd
	if len(conf.Args) > 0 {
		cmd = exec.Command(conf.Command, strings.Split(conf.Args, " ")...)
	} else {
		cmd = exec.Command(conf.Command)
	}
	newStatusChannel := make(chan ProcessStatus, 1)
	return &Process{Cmd: cmd, Conf: conf, StatusChannel: newStatusChannel}
}

//SetEnvVars sets the environment variables for the launched process
//If p.Conf.EnvVars is empty, we pass all the current env vars to the child
func (p *Process) SetEnvVars() {
	if len(p.Conf.EnvVars) == 0 {
		return
	}
	for _, envVar := range strings.Split(p.Conf.EnvVars, ",") {
		p.Cmd.Env = append(p.Cmd.Env, envVar)
	}
}

//SetUser tries to change the current user to newUsername
func (p *Process) SetUser() error {
	currentUser, err := user.Current()
	if err != nil {
		return errors.New("Can't determine the current user !\n" + err.Error())
	}

	if p.Conf.UserName == currentUser.Username {
		return nil
	}

	newUser, err := user.Lookup(p.Conf.UserName)
	if err != nil {
		return errors.New("Can't find the user" + p.Conf.UserName + "!\n" + err.Error())
	}

	newUserUID, err := strconv.Atoi(newUser.Uid)
	if err != nil {
		return errors.New("Can't determine the new user UID !\n" + err.Error())
	}

	if err := syscall.Setuid(newUserUID); err != nil {
		return errors.New("Can't change the user !\n" + err.Error())
	}

	return nil
}

func (p *Process) Terminate(maxTryCount int) error {
	if maxTryCount > 0 {
		if err := syscall.Kill(p.Cmd.Process.Pid, syscall.SIGTERM); err == nil {
			return nil
		}
	} else {
		return syscall.Kill(p.Cmd.Process.Pid, syscall.SIGKILL)
	}
	time.Sleep(100 * time.Millisecond)
	return p.Terminate(maxTryCount - 1)
}

func (p *Process) isStarted() bool {
	return p.Cmd.Process != nil
}

func (p *Process) isRunning() bool {
	return p.isStarted() && p.hasBoundPort()
}

func (p *Process) hasBoundPort() bool {
	if p.Conf.BindPort == 0 {
		return true
	}

	//We execute netstat -an | grep bindPort to find if bindPort is open
	netstCmd := exec.Command("netstat", "-an")
	grepCmd := exec.Command("grep", string(p.Conf.BindPort))
	netstOut, _ := netstCmd.Output()
	grepCmd.Stdin = bytes.NewBuffer(netstOut)
	grepOut, _ := grepCmd.Output()

	return len(grepOut) > 0
}

func (p *Process) redirectStdin() error {
	var err error
	p.oldTermState, err = terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return errors.New("Can't redirect stdin:" + err.Error())
	}
	newTerminal := terminal.NewTerminal(os.Stdin, "")
	cb := func(s string, i int, r rune) (string, int, bool) {
		car := []byte{byte(r)}
		newTerminal.Write(car)
		return s, i, false
	}
	newTerminal.AutoCompleteCallback = cb
	go io.Copy(p.Pty, os.Stdin)
	return nil
}

func (p *Process) restoreStdin() error {
	err := terminal.Restore(int(os.Stdin.Fd()), p.oldTermState)
	return err
}

func (p *Process) redirectStdout() error {
	go io.Copy(os.Stdout, p.Pty)
	return nil
}

func (p *Process) NotifyStatusChanged() error {
	if p.Conf.WebHook == "" {
		return nil
	}
	statusStr := ""
	if p.Status == PROCESS_STARTED {
		statusStr = "started"
	} else if p.Status == PROCESS_RUNNING {
		statusStr = "running"
	} else {
		statusStr = "stopped"
	}
	body := `{
							"ps":
								{ "status":` + statusStr + `}
						}`

	req, err := http.NewRequest("PUT", p.Conf.WebHook, bytes.NewBufferString(body))
	if err != nil {
		return errors.New("Failed to construct the HTTP request" + err.Error())
	}

	req.Header.Add("Content-Type", "application/json")
	client := &http.Client{}

	resp, err := client.Do(req)
	defer resp.Body.Close()
	if err != nil {
		return errors.New("Was not able to trigger the hook!\n" + err.Error())
	}
	return nil
}

func (p *Process) Start() error {
	//We start the process
	var startErr error
	p.Pty, startErr = pty.Start(p.Cmd)
	if startErr != nil {
		return startErr
	}

	go func() {
		var err error
		if err = p.redirectStdin(); err != nil {
			p.StatusChannel <- ProcessStatus{Status: -1, Err: err}
		}
		if err = p.redirectStdout(); err != nil {
			p.StatusChannel <- ProcessStatus{Status: -1, Err: err}
		}

		for !p.isStarted() {
			time.Sleep(100 * time.Millisecond)
		}
		p.Status = PROCESS_STARTED
		if err = p.NotifyStatusChanged(); err != nil {
			log.Println(err)
		}
		p.StatusChannel <- ProcessStatus{Status: p.Status, Err: nil}

		for p.isRunning() == false {
			time.Sleep(100 * time.Millisecond)
		}
		p.Status = PROCESS_RUNNING
		if err = p.NotifyStatusChanged(); err != nil {
			log.Println(err)
		}
		p.StatusChannel <- ProcessStatus{Status: p.Status, Err: nil}

		err = p.Cmd.Wait()
		if err != nil {
			p.restoreStdin()
			p.Status = PROCESS_STOPPED
			p.NotifyStatusChanged()
			log.Fatal(err)
		}

		//p has stopped
		p.restoreStdin()
		p.Status = PROCESS_STOPPED
		if err = p.NotifyStatusChanged(); err != nil {
			log.Println(err)
		}
		p.StatusChannel <- ProcessStatus{Status: p.Status, Err: nil}
	}()

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2019 by Michael Dvorkin and contributors. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	"encoding/json"
	"io/ioutil"
	"sort"
	"strings"

	"github.com/Knetic/govaluate"
)

const defaultGainColor = "green"
const defaultLossColor = "red"
const defaultTagColor = "yellow"

// Profile manages Mop program settings as defined by user (ex. list of
// stock tickers). The settings are serialized using JSON and saved in
// the ~/.moprc file.
type Profile struct {
	Tickers          []string                       // List of stock tickers to display.
	MarketRefresh    int                            // Time interval to refresh market data.
	QuotesRefresh    int                            // Time interval to refresh stock quotes.
	SortColumn       int                            // Column number by which we sort stock quotes.
	Ascending        bool                           // True when sort order is ascending.
	Grouped          bool                           // True when stocks are grouped by advancing/declining.
	Filter           string                         // Filter in human form
	TickerColors     struct {                       // Ticker colors
		Gain             string
		Loss             string
		Tag              string
	}
	filterExpression *govaluate.EvaluableExpression // The filter as a govaluate expression
	selectedColumn   int                            // Stores selected column number when the column editor is active.
	filename         string                         // Path to the file in which the configuration is stored
}

func IsSupportedColor(colorName string) bool {
	switch colorName {
		case
			"black",
			"red",
			"green",
			"yellow",
			"blue",
			"magenta",
			"cyan",
			"white":
			return true
	}
	return false
}

// Creates the profile and attempts to load the settings from ~/.moprc file.
// If the file is not there it gets created with default values.
func NewProfile(filename string) *Profile {
	profile := &Profile{filename: filename}
	data, err := ioutil.ReadFile(filename)
	if err != nil { // Set default values:
		profile.MarketRefresh = 12 // Market data gets fetched every 12s (5 times per minute).
		profile.QuotesRefresh = 5  // Stock quotes get updated every 5s (12 times per minute).
		profile.Grouped = false    // Stock quotes are *not* grouped by advancing/declining.
		profile.Tickers = []string{`AAPL`, `C`, `GOOG`, `IBM`, `KO`, `ORCL`, `V`}
		profile.SortColumn = 0   // Stock quotes are sorted by ticker name.
		profile.Ascending = true // A to Z.
		profile.Filter = ""
		profile.TickerColors.Gain = defaultGainColor
		profile.TickerColors.Loss = defaultLossColor
		profile.TickerColors.Tag = defaultTagColor
		profile.Save()
	} else {
		json.Unmarshal(data, profile)

		profile.TickerColors.Gain = strings.ToLower(profile.TickerColors.Gain)
		if !IsSupportedColor(profile.TickerColors.Gain) {
			profile.TickerColors.Gain = defaultGainColor
		}

		profile.TickerColors.Loss = strings.ToLower(profile.TickerColors.Loss)
		if !IsSupportedColor(profile.TickerColors.Loss) {
			profile.TickerColors.Loss = defaultLossColor
		}

		profile.TickerColors.Tag = strings.ToLower(profile.TickerColors.Tag)
		if !IsSupportedColor(profile.TickerColors.Tag) {
			profile.TickerColors.Tag = defaultTagColor
		}

		profile.SetFilter(profile.Filter)
	}
	profile.selectedColumn = -1

	return profile
}

// Save serializes settings using JSON and saves them in ~/.moprc file.
func (profile *Profile) Save() error {
	data, err := json.Marshal(profile)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(profile.filename, data, 0644)
}

// AddTickers updates the list of existing tickers to add the new ones making
// sure there are no duplicates.
func (profile *Profile) AddTickers(tickers []string) (added int, err error) {
	added, err = 0, nil
	existing := make(map[string]bool)

	// Build a hash of existing tickers so we could look it up quickly.
	for _, ticker := range profile.Tickers {
		existing[ticker] = true
	}

	// Iterate over the list of new tickers excluding the ones that
	// already exist.
	for _, ticker := range tickers {
		if _, found := existing[ticker]; !found {
			profile.Tickers = append(profile.Tickers, ticker)
			added++
		}
	}

	if added > 0 {
		sort.Strings(profile.Tickers)
		err = profile.Save()
	}

	return
}

// RemoveTickers removes requested stock tickers from the list we track.
func (profile *Profile) RemoveTickers(tickers []string) (removed int, err error) {
	removed, err = 0, nil
	for _, ticker := range tickers {
		for i, existing := range profile.Tickers {
			if ticker == existing {
				// Requested ticker is there: remove i-th slice item.
				profile.Tickers = append(profile.Tickers[:i], profile.Tickers[i+1:]...)
				removed++
			}
		}
	}

	if removed > 0 {
		err = profile.Save()
	}

	return
}

// Reorder gets called by the column editor to either reverse sorting order
// for the current column, or to pick another sort column.
func (profile *Profile) Reorder() error {
	if profile.selectedColumn == profile.SortColumn {
		profile.Ascending = !profile.Ascending // Reverse sort order.
	} else {
		profile.SortColumn = profile.selectedColumn // Pick new sort column.
	}
	return profile.Save()
}

// Regroup flips the flag that controls whether the stock quotes are grouped
// by advancing/declining issues.
func (profile *Profile) Regroup() error {
	profile.Grouped = !profile.Grouped
	return profile.Save()
}

// SetFilter creates a govaluate.EvaluableExpression.
func (profile *Profile) SetFilter(filter string) {
	if len(filter) > 0 {
		var err error
		profile.filterExpression, err = govaluate.NewEvaluableExpression(filter)

		if err != nil {
			panic(err)
		}

	} else if len(filter) == 0 && profile.filterExpression != nil {
		profile.filterExpression = nil
	}

	profile.Filter = filter
}
<commit_msg>Save the `.moprc` file in a cleaner way.<commit_after>// Copyright (c) 2013-2019 by Michael Dvorkin and contributors. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	"encoding/json"
	"io/ioutil"
	"sort"
	"strings"

	"github.com/Knetic/govaluate"
)

const defaultGainColor = "green"
const defaultLossColor = "red"
const defaultTagColor = "yellow"

// Profile manages Mop program settings as defined by user (ex. list of
// stock tickers). The settings are serialized using JSON and saved in
// the ~/.moprc file.
type Profile struct {
	Tickers          []string                       // List of stock tickers to display.
	MarketRefresh    int                            // Time interval to refresh market data.
	QuotesRefresh    int                            // Time interval to refresh stock quotes.
	SortColumn       int                            // Column number by which we sort stock quotes.
	Ascending        bool                           // True when sort order is ascending.
	Grouped          bool                           // True when stocks are grouped by advancing/declining.
	Filter           string                         // Filter in human form
	TickerColors     struct {                       // Ticker colors
		Gain             string
		Loss             string
		Tag              string
	}
	filterExpression *govaluate.EvaluableExpression // The filter as a govaluate expression
	selectedColumn   int                            // Stores selected column number when the column editor is active.
	filename         string                         // Path to the file in which the configuration is stored
}

func IsSupportedColor(colorName string) bool {
	switch colorName {
		case
			"black",
			"red",
			"green",
			"yellow",
			"blue",
			"magenta",
			"cyan",
			"white":
			return true
	}
	return false
}

// Creates the profile and attempts to load the settings from ~/.moprc file.
// If the file is not there it gets created with default values.
func NewProfile(filename string) *Profile {
	profile := &Profile{filename: filename}
	data, err := ioutil.ReadFile(filename)
	if err != nil { // Set default values:
		profile.MarketRefresh = 12 // Market data gets fetched every 12s (5 times per minute).
		profile.QuotesRefresh = 5  // Stock quotes get updated every 5s (12 times per minute).
		profile.Grouped = false    // Stock quotes are *not* grouped by advancing/declining.
		profile.Tickers = []string{`AAPL`, `C`, `GOOG`, `IBM`, `KO`, `ORCL`, `V`}
		profile.SortColumn = 0   // Stock quotes are sorted by ticker name.
		profile.Ascending = true // A to Z.
		profile.Filter = ""
		profile.TickerColors.Gain = defaultGainColor
		profile.TickerColors.Loss = defaultLossColor
		profile.TickerColors.Tag = defaultTagColor
		profile.Save()
	} else {
		json.Unmarshal(data, profile)

		profile.TickerColors.Gain = strings.ToLower(profile.TickerColors.Gain)
		if !IsSupportedColor(profile.TickerColors.Gain) {
			profile.TickerColors.Gain = defaultGainColor
		}

		profile.TickerColors.Loss = strings.ToLower(profile.TickerColors.Loss)
		if !IsSupportedColor(profile.TickerColors.Loss) {
			profile.TickerColors.Loss = defaultLossColor
		}

		profile.TickerColors.Tag = strings.ToLower(profile.TickerColors.Tag)
		if !IsSupportedColor(profile.TickerColors.Tag) {
			profile.TickerColors.Tag = defaultTagColor
		}

		profile.SetFilter(profile.Filter)
	}
	profile.selectedColumn = -1

	return profile
}

// Save serializes settings using JSON and saves them in ~/.moprc file.
func (profile *Profile) Save() error {
	data, err := json.MarshalIndent(profile, "", "    ")
	if err != nil {
		return err
	}

	return ioutil.WriteFile(profile.filename, data, 0644)
}

// AddTickers updates the list of existing tickers to add the new ones making
// sure there are no duplicates.
func (profile *Profile) AddTickers(tickers []string) (added int, err error) {
	added, err = 0, nil
	existing := make(map[string]bool)

	// Build a hash of existing tickers so we could look it up quickly.
	for _, ticker := range profile.Tickers {
		existing[ticker] = true
	}

	// Iterate over the list of new tickers excluding the ones that
	// already exist.
	for _, ticker := range tickers {
		if _, found := existing[ticker]; !found {
			profile.Tickers = append(profile.Tickers, ticker)
			added++
		}
	}

	if added > 0 {
		sort.Strings(profile.Tickers)
		err = profile.Save()
	}

	return
}

// RemoveTickers removes requested stock tickers from the list we track.
func (profile *Profile) RemoveTickers(tickers []string) (removed int, err error) {
	removed, err = 0, nil
	for _, ticker := range tickers {
		for i, existing := range profile.Tickers {
			if ticker == existing {
				// Requested ticker is there: remove i-th slice item.
				profile.Tickers = append(profile.Tickers[:i], profile.Tickers[i+1:]...)
				removed++
			}
		}
	}

	if removed > 0 {
		err = profile.Save()
	}

	return
}

// Reorder gets called by the column editor to either reverse sorting order
// for the current column, or to pick another sort column.
func (profile *Profile) Reorder() error {
	if profile.selectedColumn == profile.SortColumn {
		profile.Ascending = !profile.Ascending // Reverse sort order.
	} else {
		profile.SortColumn = profile.selectedColumn // Pick new sort column.
	}
	return profile.Save()
}

// Regroup flips the flag that controls whether the stock quotes are grouped
// by advancing/declining issues.
func (profile *Profile) Regroup() error {
	profile.Grouped = !profile.Grouped
	return profile.Save()
}

// SetFilter creates a govaluate.EvaluableExpression.
func (profile *Profile) SetFilter(filter string) {
	if len(filter) > 0 {
		var err error
		profile.filterExpression, err = govaluate.NewEvaluableExpression(filter)

		if err != nil {
			panic(err)
		}

	} else if len(filter) == 0 && profile.filterExpression != nil {
		profile.filterExpression = nil
	}

	profile.Filter = filter
}
<|endoftext|>


<commit_before>package render

import (
	"github.com/shurcooL/github_flavored_markdown"

	"html/template"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

type Document struct {
	InputFile string
	BaseName  string
	Title     string
	RawBody   string
	HTMLBody  template.HTML
	FullHTML  string
	Date      time.Time
	Updated   time.Time
	Params    map[string]string
}

func NewDocument(inputFile string) *Document {
	basename := path.Base(inputFile)
	doc := &Document{
		InputFile: inputFile,
		BaseName:  basename,
		Date:      time.Now(),
		Updated:   time.Now(),
		Params:    make(map[string]string),
	}
	return doc
}

func (d *Document) ParseDocument() {
	data, err := ioutil.ReadFile(d.InputFile)
	if err != nil {
		log.Fatalf("Couldn't read file: '%s': %s", d.InputFile, err)
	}
	startContentLine := d.parseHeader(data)
	d.parseContents(data, startContentLine)
}

func (d *Document) parseHeader(data []byte) int {
	lines := strings.Split(string(data), "\n")

	headerState := 0
	startMdContentLine := 0
	for i, line := range lines {
		line = strings.TrimSpace(line)

		// in the header section: from --- to ---
		if headerState == 1 {
			colonIndex := strings.Index(line, ":")
			if colonIndex > 0 {
				key := strings.TrimSpace(line[:colonIndex])
				value := strings.TrimSpace(line[colonIndex+1:])
				// remove surrounding quotes
				value = strings.Trim(value, "\"")
				switch key {
				case "Title":
					d.Title = value
				case "Date":
					d.Date, _ =
						time.Parse("2006-01-02", value)
				case "Updated":
					d.Updated, _ =
						time.Parse("2006-01-02", value)
				default:
					d.Params[key] = value
				}
			}
		} else if headerState >= 2 {
			startMdContentLine = i
			break
		}

		if strings.HasPrefix(line, "---") {
			headerState += 1
		}
	}

	return startMdContentLine
}

func (d *Document) parseContents(data []byte, startLine int) {
	lines := strings.Split(string(data), "\n")
	contents := strings.Join(lines[startLine:], "\n")
	d.RawBody = contents
}

func (d *Document) RenderHTML(outDir string) {
	text := []byte(d.RawBody)
	d.HTMLBody = template.HTML(github_flavored_markdown.Markdown(text))
	base := strings.TrimSuffix(d.InputFile, filepath.Ext(d.InputFile))
	htmlFile := strings.Join([]string{base, "html"}, ".")
	htmlPath := filepath.Join(outDir, htmlFile)

	f, err := os.Create(htmlPath)
	if err != nil {
		log.Fatalf("Couldn't open '%s' : %s\n", htmlPath, err)
	}
	defer f.Close()

	baseTemplate := "views/base.html"
	t, err := template.ParseFiles(baseTemplate)
	if err != nil {
		log.Fatalf("Template '%s' parse error: %s\n", baseTemplate, err)
	}
	err = t.Execute(f, d)
	if err != nil {
		log.Fatalf("Template render error: %s\n", err)
	}

	// alternative approach:
	// http://stackoverflow.com/questions/23124008/how-can-i-render-markdown-to-a-golang-templatehtml-or-tmpl-with-blackfriday
}
<commit_msg>+ correct to where the views are now placed<commit_after>package render

import (
	"github.com/shurcooL/github_flavored_markdown"

	"html/template"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

type Document struct {
	InputFile string
	BaseName  string
	Title     string
	RawBody   string
	HTMLBody  template.HTML
	FullHTML  string
	Date      time.Time
	Updated   time.Time
	Params    map[string]string
}

func NewDocument(inputFile string) *Document {
	basename := path.Base(inputFile)
	doc := &Document{
		InputFile: inputFile,
		BaseName:  basename,
		Date:      time.Now(),
		Updated:   time.Now(),
		Params:    make(map[string]string),
	}
	return doc
}

func (d *Document) ParseDocument() {
	data, err := ioutil.ReadFile(d.InputFile)
	if err != nil {
		log.Fatalf("Couldn't read file: '%s': %s", d.InputFile, err)
	}
	startContentLine := d.parseHeader(data)
	d.parseContents(data, startContentLine)
}

func (d *Document) parseHeader(data []byte) int {
	lines := strings.Split(string(data), "\n")

	headerState := 0
	startMdContentLine := 0
	for i, line := range lines {
		line = strings.TrimSpace(line)

		// in the header section: from --- to ---
		if headerState == 1 {
			colonIndex := strings.Index(line, ":")
			if colonIndex > 0 {
				key := strings.TrimSpace(line[:colonIndex])
				value := strings.TrimSpace(line[colonIndex+1:])
				// remove surrounding quotes
				value = strings.Trim(value, "\"")
				switch key {
				case "Title":
					d.Title = value
				case "Date":
					d.Date, _ =
						time.Parse("2006-01-02", value)
				case "Updated":
					d.Updated, _ =
						time.Parse("2006-01-02", value)
				default:
					d.Params[key] = value
				}
			}
		} else if headerState >= 2 {
			startMdContentLine = i
			break
		}

		if strings.HasPrefix(line, "---") {
			headerState += 1
		}
	}

	return startMdContentLine
}

func (d *Document) parseContents(data []byte, startLine int) {
	lines := strings.Split(string(data), "\n")
	contents := strings.Join(lines[startLine:], "\n")
	d.RawBody = contents
}

func (d *Document) RenderHTML(outDir string) {
	text := []byte(d.RawBody)
	d.HTMLBody = template.HTML(github_flavored_markdown.Markdown(text))
	base := strings.TrimSuffix(d.InputFile, filepath.Ext(d.InputFile))
	htmlFile := strings.Join([]string{base, "html"}, ".")
	htmlPath := filepath.Join(outDir, htmlFile)

	f, err := os.Create(htmlPath)
	if err != nil {
		log.Fatalf("Couldn't open '%s' : %s\n", htmlPath, err)
	}
	defer f.Close()

	baseTemplate := "assets/views/base.html"
	t, err := template.ParseFiles(baseTemplate)
	if err != nil {
		log.Fatalf("Template '%s' parse error: %s\n", baseTemplate, err)
	}
	err = t.Execute(f, d)
	if err != nil {
		log.Fatalf("Template render error: %s\n", err)
	}

	// alternative approach:
	// http://stackoverflow.com/questions/23124008/how-can-i-render-markdown-to-a-golang-templatehtml-or-tmpl-with-blackfriday
}
<|endoftext|>


<commit_before>package jira

import (
	"fmt"
)

// ProjectService handles projects for the JIRA instance / API.
//
// JIRA API docs: https://docs.atlassian.com/jira/REST/latest/#api/2/project
type ProjectService struct {
	client *Client
}

// ProjectList represent a list of Projects
type ProjectList []struct {
	Expand          string          `json:"expand"`
	Self            string          `json:"self"`
	ID              string          `json:"id"`
	Key             string          `json:"key"`
	Name            string          `json:"name"`
	AvatarUrls      AvatarUrls      `json:"avatarUrls"`
	ProjectTypeKey  string          `json:"projectTypeKey"`
	ProjectCategory ProjectCategory `json:"projectCategory,omitempty"`
}

// ProjectCategory represents a single project category
type ProjectCategory struct {
	Self        string `json:"self"`
	ID          string `json:"id"`
	Name        string `json:"name"`
	Description string `json:"description"`
}

// Project represents a JIRA Project.
type Project struct {
	Expand       string             `json:"expand,omitempty"`
	Self         string             `json:"self,omitempty"`
	ID           string             `json:"id,omitempty"`
	Key          string             `json:"key,omitempty"`
	Description  string             `json:"description,omitempty"`
	Lead         User               `json:"lead,omitempty"`
	Components   []ProjectComponent `json:"components,omitempty"`
	IssueTypes   []IssueType        `json:"issueTypes,omitempty"`
	URL          string             `json:"url,omitempty"`
	Email        string             `json:"email,omitempty"`
	AssigneeType string             `json:"assigneeType,omitempty"`
	Versions     []Version      `json:"versions,omitempty"`
	Name         string             `json:"name,omitempty"`
	Roles        struct {
		Developers string `json:"Developers,omitempty"`
	} `json:"roles,omitempty"`
	AvatarUrls      AvatarUrls      `json:"avatarUrls,omitempty"`
	ProjectCategory ProjectCategory `json:"projectCategory,omitempty"`
}

// Version represents a single release version of a project
type Version struct {
	Self            string `json:"self"`
	ID              string `json:"id"`
	Name            string `json:"name"`
	Archived        bool   `json:"archived"`
	Released        bool   `json:"released"`
	ReleaseDate     string `json:"releaseDate"`
	UserReleaseDate string `json:"userReleaseDate"`
	ProjectID       int    `json:"projectId"`  // Unlike other IDs, this is returned as a number
}

// ProjectComponent represents a single component of a project
type ProjectComponent struct {
	Self                string `json:"self"`
	ID                  string `json:"id"`
	Name                string `json:"name"`
	Description         string `json:"description"`
	Lead                User   `json:"lead"`
	AssigneeType        string `json:"assigneeType"`
	Assignee            User   `json:"assignee"`
	RealAssigneeType    string `json:"realAssigneeType"`
	RealAssignee        User   `json:"realAssignee"`
	IsAssigneeTypeValid bool   `json:"isAssigneeTypeValid"`
	Project             string `json:"project"`
	ProjectID           int    `json:"projectId"`
}

// GetList gets all projects form JIRA
//
// JIRA API docs: https://docs.atlassian.com/jira/REST/latest/#api/2/project-getAllProjects
func (s *ProjectService) GetList() (*ProjectList, *Response, error) {
	apiEndpoint := "rest/api/2/project"
	req, err := s.client.NewRequest("GET", apiEndpoint, nil)
	if err != nil {
		return nil, nil, err
	}

	projectList := new(ProjectList)
	resp, err := s.client.Do(req, projectList)
	if err != nil {
		return nil, resp, err
	}
	return projectList, resp, nil
}

// Get returns a full representation of the project for the given issue key.
// JIRA will attempt to identify the project by the projectIdOrKey path parameter.
// This can be an project id, or an project key.
//
// JIRA API docs: https://docs.atlassian.com/jira/REST/latest/#api/2/project-getProject
func (s *ProjectService) Get(projectID string) (*Project, *Response, error) {
	apiEndpoint := fmt.Sprintf("/rest/api/2/project/%s", projectID)
	req, err := s.client.NewRequest("GET", apiEndpoint, nil)
	if err != nil {
		return nil, nil, err
	}

	project := new(Project)
	resp, err := s.client.Do(req, project)
	if err != nil {
		return nil, resp, err
	}
	return project, resp, nil
}
<commit_msg>Code formatting<commit_after>package jira

import (
	"fmt"
)

// ProjectService handles projects for the JIRA instance / API.
//
// JIRA API docs: https://docs.atlassian.com/jira/REST/latest/#api/2/project
type ProjectService struct {
	client *Client
}

// ProjectList represent a list of Projects
type ProjectList []struct {
	Expand          string          `json:"expand"`
	Self            string          `json:"self"`
	ID              string          `json:"id"`
	Key             string          `json:"key"`
	Name            string          `json:"name"`
	AvatarUrls      AvatarUrls      `json:"avatarUrls"`
	ProjectTypeKey  string          `json:"projectTypeKey"`
	ProjectCategory ProjectCategory `json:"projectCategory,omitempty"`
}

// ProjectCategory represents a single project category
type ProjectCategory struct {
	Self        string `json:"self"`
	ID          string `json:"id"`
	Name        string `json:"name"`
	Description string `json:"description"`
}

// Project represents a JIRA Project.
type Project struct {
	Expand       string             `json:"expand,omitempty"`
	Self         string             `json:"self,omitempty"`
	ID           string             `json:"id,omitempty"`
	Key          string             `json:"key,omitempty"`
	Description  string             `json:"description,omitempty"`
	Lead         User               `json:"lead,omitempty"`
	Components   []ProjectComponent `json:"components,omitempty"`
	IssueTypes   []IssueType        `json:"issueTypes,omitempty"`
	URL          string             `json:"url,omitempty"`
	Email        string             `json:"email,omitempty"`
	AssigneeType string             `json:"assigneeType,omitempty"`
	Versions     []Version          `json:"versions,omitempty"`
	Name         string             `json:"name,omitempty"`
	Roles        struct {
		Developers string `json:"Developers,omitempty"`
	} `json:"roles,omitempty"`
	AvatarUrls      AvatarUrls      `json:"avatarUrls,omitempty"`
	ProjectCategory ProjectCategory `json:"projectCategory,omitempty"`
}

// Version represents a single release version of a project
type Version struct {
	Self            string `json:"self"`
	ID              string `json:"id"`
	Name            string `json:"name"`
	Archived        bool   `json:"archived"`
	Released        bool   `json:"released"`
	ReleaseDate     string `json:"releaseDate"`
	UserReleaseDate string `json:"userReleaseDate"`
	ProjectID       int    `json:"projectId"`  // Unlike other IDs, this is returned as a number
}

// ProjectComponent represents a single component of a project
type ProjectComponent struct {
	Self                string `json:"self"`
	ID                  string `json:"id"`
	Name                string `json:"name"`
	Description         string `json:"description"`
	Lead                User   `json:"lead"`
	AssigneeType        string `json:"assigneeType"`
	Assignee            User   `json:"assignee"`
	RealAssigneeType    string `json:"realAssigneeType"`
	RealAssignee        User   `json:"realAssignee"`
	IsAssigneeTypeValid bool   `json:"isAssigneeTypeValid"`
	Project             string `json:"project"`
	ProjectID           int    `json:"projectId"`
}

// GetList gets all projects form JIRA
//
// JIRA API docs: https://docs.atlassian.com/jira/REST/latest/#api/2/project-getAllProjects
func (s *ProjectService) GetList() (*ProjectList, *Response, error) {
	apiEndpoint := "rest/api/2/project"
	req, err := s.client.NewRequest("GET", apiEndpoint, nil)
	if err != nil {
		return nil, nil, err
	}

	projectList := new(ProjectList)
	resp, err := s.client.Do(req, projectList)
	if err != nil {
		return nil, resp, err
	}
	return projectList, resp, nil
}

// Get returns a full representation of the project for the given issue key.
// JIRA will attempt to identify the project by the projectIdOrKey path parameter.
// This can be an project id, or an project key.
//
// JIRA API docs: https://docs.atlassian.com/jira/REST/latest/#api/2/project-getProject
func (s *ProjectService) Get(projectID string) (*Project, *Response, error) {
	apiEndpoint := fmt.Sprintf("/rest/api/2/project/%s", projectID)
	req, err := s.client.NewRequest("GET", apiEndpoint, nil)
	if err != nil {
		return nil, nil, err
	}

	project := new(Project)
	resp, err := s.client.Do(req, project)
	if err != nil {
		return nil, resp, err
	}
	return project, resp, nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"labix.org/v2/mgo"
	"log"
	. "replay"
	"runtime"
	"sync/atomic"
	"time"
)

func panicOnError(err error) {
	if err != nil {
		panic(err)
	}
}

var (
	opsFilename string
	url         string
	workers     int
	maxOps      int
	numSkipOps     int
	socketTimeout	int64
	style		string
	sampleRate	float64
)

const (
	// Set one minute timeout on mongo socket connections (nanoseconds) by default
	DEFAULT_MGO_SOCKET_TIMEOUT=60000000000
)

func init() {
	flag.StringVar(&opsFilename, "ops_filename", "",
		"The file for the serialized ops, generated by the record scripts.")
	flag.StringVar(&url, "url", "",
		"The database server's url, in the format of <host>[:<port>]")
	flag.StringVar(&style, "style", "",
		"How to replay the the ops. You can choose: \n"+
			"	stress: repaly ops at fast as possible\n"+
			"	real: repaly ops in accordance to ops' timestamps")
	flag.IntVar(&workers, "workers", 10,
		"Number of workers that sends ops to database.")
	flag.IntVar(&maxOps, "maxOps", 0,
		"[Optional] Maximal amount of ops to be replayed from the "+
			"ops_filename file. By setting it to `0`, replayer will "+
			"replay all the ops.")
	flag.IntVar(&numSkipOps, "numSkipOps", 0,
		"[Optional] Skip first N ops. Useful for when the total ops in ops_filename" +
		" exceeds available memory and you're running in stress mode.")
	flag.Int64Var(&socketTimeout, "socketTimeout", DEFAULT_MGO_SOCKET_TIMEOUT, "Mongo socket timeout in nanoseconds. Defaults to 60 seconds.")
	flag.Float64Var(&sampleRate, "sample_rate", 0.0, "sample ops for latency")
}

func parseFlags() error {
	flag.Parse()
	if style != "stress" && style != "real" {
		return errors.New("Cannot recognize the style: " + style)
	}
	if workers <= 0 {
		return errors.New("`workers` should be a positive number")
	}
	if maxOps == 0 {
		maxOps = 4294967295
	}
	return nil
}

func main() {
	// Will enable system threads to make sure all cpus can be well utilized.
	runtime.GOMAXPROCS(100)
	err := parseFlags()
	panicOnError(err)

	// Prepare to dispatch ops
	var reader OpsReader
	var opsChan chan *Op
	if style == "stress" {
		err, reader = NewFileByLineOpsReader(opsFilename)
		panicOnError(err)
		if numSkipOps > 0 {
			err = reader.SkipOps(numSkipOps)
			panicOnError(err)
		}
		opsChan = NewBestEffortOpsDispatcher(reader, maxOps)
	} else {
		// TODO NewCyclicOpsReader: do we really want to make it cyclic?
		reader = NewCyclicOpsReader(func() OpsReader {
			err, reader := NewFileByLineOpsReader(opsFilename)
			panicOnError(err)
			return reader
		})
		if numSkipOps > 0 {
			err = reader.SkipOps(numSkipOps)
			panicOnError(err)
		}
		opsChan = NewByTimeOpsDispatcher(reader, maxOps)
	}

	latencyChan := make(chan Latency, workers)

	// Set up workers to do the job
	exit := make(chan int)
	opsExecuted := int64(0)
	fetch := func(id int, statsCollector IStatsCollector) {
		log.Printf("Worker #%d report for duty\n", id)

		session, err := mgo.Dial(url)
		panicOnError(err)
		session.SetSocketTimeout(time.Duration(socketTimeout))

		defer session.Close()
		exec := OpsExecutorWithStats(session, statsCollector)
		for {
			op := <-opsChan
			if op == nil {
				break
			}
			exec.Execute(op)
			atomic.AddInt64(&opsExecuted, 1)
		}
		exit <- 1
		log.Printf("Worker #%d done!\n", id)
	}
	statsCollectorList := make([]*StatsCollector, workers)
	for i := 0; i < workers; i++ {
		statsCollectorList[i] = NewStatsCollector()
		statsCollectorList[i].SampleLatencies(sampleRate, latencyChan)
		go fetch(i, statsCollectorList[i])
	}

	// Periodically report execution status
	go func() {
		statsAnalyzer := NewStatsAnalyzer(statsCollectorList, &opsExecuted,
			latencyChan, int(sampleRate*float64(maxOps)))
		toFloat := func(nano int64) float64 {
			return float64(nano) / float64(1e6)
		}
		report := func() {
			status := statsAnalyzer.GetStatus()
			log.Printf("Executed %d ops, %.2f ops/sec", opsExecuted,
				status.OpsPerSec)
			for _, opType := range AllOpTypes {
				allTime := status.AllTimeLatencies[opType]
				sinceLast := status.SinceLastLatencies[opType]
				log.Printf("  Op type: %s, count: %d, ops/sec: %.2f",
					opType, status.Counts[opType], status.TypeOpsSec[opType])
				template := "   %s: P50: %.2fms, P70: %.2fms, P90: %.2fms, " +
					"P95 %.2fms, P99 %.2fms, Max %.2fms\n"
				log.Printf(template, "Total", toFloat(allTime[P50]),
					toFloat(allTime[P70]), toFloat(allTime[P90]),
					toFloat(allTime[P95]), toFloat(allTime[P99]),
					toFloat(allTime[P100]))
				log.Printf(template, "Last ", toFloat(sinceLast[P50]),
					toFloat(sinceLast[P70]), toFloat(sinceLast[P90]),
					toFloat(sinceLast[P95]), toFloat(sinceLast[P99]),
					toFloat(sinceLast[P100]))
			}
		}
		defer report()

		for opsExecuted < int64(maxOps) {
			time.Sleep(5 * time.Second)
			report()
		}
	}()

	// Wait for workers
	received := 0
	for received < workers {
		<-exit
		received += 1
	}
}
<commit_msg>Add optional verbose option to print op errors and other info for future troubleshooting.<commit_after>package main

import (
	"errors"
	"flag"
	"labix.org/v2/mgo"
	"log"
	. "replay"
	"runtime"
	"sync/atomic"
	"time"
)

func panicOnError(err error) {
	if err != nil {
		panic(err)
	}
}

var (
	opsFilename string
	url         string
	workers     int
	maxOps      int
	numSkipOps     int
	socketTimeout	int64
	style		string
	sampleRate	float64
	verbose		bool
)

const (
	// Set one minute timeout on mongo socket connections (nanoseconds) by default
	DEFAULT_MGO_SOCKET_TIMEOUT=60000000000
)

func init() {
	flag.StringVar(&opsFilename, "ops_filename", "",
		"The file for the serialized ops, generated by the record scripts.")
	flag.StringVar(&url, "url", "",
		"The database server's url, in the format of <host>[:<port>]")
	flag.StringVar(&style, "style", "",
		"How to replay the the ops. You can choose: \n"+
			"	stress: repaly ops at fast as possible\n"+
			"	real: repaly ops in accordance to ops' timestamps")
	flag.IntVar(&workers, "workers", 10,
		"Number of workers that sends ops to database.")
	flag.IntVar(&maxOps, "maxOps", 0,
		"[Optional] Maximal amount of ops to be replayed from the "+
			"ops_filename file. By setting it to `0`, replayer will "+
			"replay all the ops.")
	flag.IntVar(&numSkipOps, "numSkipOps", 0,
		"[Optional] Skip first N ops. Useful for when the total ops in ops_filename" +
		" exceeds available memory and you're running in stress mode.")
	flag.Int64Var(&socketTimeout, "socketTimeout", DEFAULT_MGO_SOCKET_TIMEOUT, "Mongo socket timeout in nanoseconds. Defaults to 60 seconds.")
	flag.Float64Var(&sampleRate, "sample_rate", 0.0, "sample ops for latency")
	// TODO: define a real error logger
	flag.BoolVar(&verbose, "verbose", false, "[Optional] Print op errors and other verbose information to stdout.")
}

func parseFlags() error {
	flag.Parse()
	if style != "stress" && style != "real" {
		return errors.New("Cannot recognize the style: " + style)
	}
	if workers <= 0 {
		return errors.New("`workers` should be a positive number")
	}
	if maxOps == 0 {
		maxOps = 4294967295
	}
	return nil
}

func main() {
	// Will enable system threads to make sure all cpus can be well utilized.
	runtime.GOMAXPROCS(100)
	err := parseFlags()
	panicOnError(err)

	// Prepare to dispatch ops
	var reader OpsReader
	var opsChan chan *Op
	if style == "stress" {
		err, reader = NewFileByLineOpsReader(opsFilename)
		panicOnError(err)
		if numSkipOps > 0 {
			err = reader.SkipOps(numSkipOps)
			panicOnError(err)
		}
		opsChan = NewBestEffortOpsDispatcher(reader, maxOps)
	} else {
		// TODO NewCyclicOpsReader: do we really want to make it cyclic?
		reader = NewCyclicOpsReader(func() OpsReader {
			err, reader := NewFileByLineOpsReader(opsFilename)
			panicOnError(err)
			return reader
		})
		if numSkipOps > 0 {
			err = reader.SkipOps(numSkipOps)
			panicOnError(err)
		}
		opsChan = NewByTimeOpsDispatcher(reader, maxOps)
	}

	latencyChan := make(chan Latency, workers)

	// Set up workers to do the job
	exit := make(chan int)
	opsExecuted := int64(0)
	fetch := func(id int, statsCollector IStatsCollector) {
		log.Printf("Worker #%d report for duty\n", id)

		session, err := mgo.Dial(url)
		panicOnError(err)
		session.SetSocketTimeout(time.Duration(socketTimeout))

		defer session.Close()
		exec := OpsExecutorWithStats(session, statsCollector)
		for {
			op := <-opsChan
			if op == nil {
				break
			}
			err := exec.Execute(op)

			if verbose == true && err != nil {
				log.Println(err.Error())
			}
			atomic.AddInt64(&opsExecuted, 1)
		}
		exit <- 1
		log.Printf("Worker #%d done!\n", id)
	}
	statsCollectorList := make([]*StatsCollector, workers)
	for i := 0; i < workers; i++ {
		statsCollectorList[i] = NewStatsCollector()
		statsCollectorList[i].SampleLatencies(sampleRate, latencyChan)
		go fetch(i, statsCollectorList[i])
	}

	// Periodically report execution status
	go func() {
		statsAnalyzer := NewStatsAnalyzer(statsCollectorList, &opsExecuted,
			latencyChan, int(sampleRate*float64(maxOps)))
		toFloat := func(nano int64) float64 {
			return float64(nano) / float64(1e6)
		}
		report := func() {
			status := statsAnalyzer.GetStatus()
			log.Printf("Executed %d ops, %.2f ops/sec", opsExecuted,
				status.OpsPerSec)
			for _, opType := range AllOpTypes {
				allTime := status.AllTimeLatencies[opType]
				sinceLast := status.SinceLastLatencies[opType]
				log.Printf("  Op type: %s, count: %d, ops/sec: %.2f",
					opType, status.Counts[opType], status.TypeOpsSec[opType])
				template := "   %s: P50: %.2fms, P70: %.2fms, P90: %.2fms, " +
					"P95 %.2fms, P99 %.2fms, Max %.2fms\n"
				log.Printf(template, "Total", toFloat(allTime[P50]),
					toFloat(allTime[P70]), toFloat(allTime[P90]),
					toFloat(allTime[P95]), toFloat(allTime[P99]),
					toFloat(allTime[P100]))
				log.Printf(template, "Last ", toFloat(sinceLast[P50]),
					toFloat(sinceLast[P70]), toFloat(sinceLast[P90]),
					toFloat(sinceLast[P95]), toFloat(sinceLast[P99]),
					toFloat(sinceLast[P100]))
			}
		}
		defer report()

		for opsExecuted < int64(maxOps) {
			time.Sleep(5 * time.Second)
			report()
		}
	}()

	// Wait for workers
	received := 0
	for received < workers {
		<-exit
		received += 1
	}
}
<|endoftext|>


<commit_before>package zenodb

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudfoundry/gosigar"
	"github.com/dustin/go-humanize"
	"github.com/getlantern/goexpr/geo"
	"github.com/getlantern/goexpr/isp"
	geredis "github.com/getlantern/goexpr/redis"
	"github.com/getlantern/golog"
	"github.com/getlantern/vtime"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb/common"
	"github.com/getlantern/zenodb/planner"
	"github.com/getlantern/zenodb/sql"
	"github.com/rickar/props"
	"github.com/shirou/gopsutil/process"
	"gopkg.in/redis.v5"
)

var (
	log = golog.LoggerFor("zenodb")

	systemRAM float64
)

func init() {
	mem := sigar.Mem{}
	err := mem.Get()
	if err != nil {
		panic(fmt.Sprintf("Unable to get system memory: %v", err))
	}
	systemRAM = float64(mem.Total)
}

// DBOpts provides options for configuring the database.
type DBOpts struct {
	// Dir points at the directory that contains the data files.
	Dir string
	// SchemaFile points at a YAML schema file that configures the tables and
	// views in the database.
	SchemaFile string
	// AliasesFile points at a file that contains expression aliases in the form
	// name=template(%v, %v), with one alias per line.
	AliasesFile string
	// EnableGeo enables geolocation functions
	EnableGeo bool
	// ISPProvider configures a provider of ISP lookups. Specify this to allow the
	// use of ISP functions.
	ISPProvider isp.Provider
	// RedisClient provides a connection to redis which enables the use of Redis
	// expressions like HGET.
	RedisClient *redis.Client
	// RedisCacheSize controls the size of redis hash caches
	RedisCacheSize int
	// VirtualTime, if true, tells zenodb to use a virtual clock that advances
	// based on the timestamps of Points received via inserts.
	VirtualTime bool
	// WALSyncInterval governs how frequently to sync the WAL to disk. 0 means
	// it syncs after every write (which is not great for performance).
	WALSyncInterval time.Duration
	// MaxWALSize limits how much WAL data to keep (in bytes)
	MaxWALSize int
	// WALCompressionSize specifies the size beyond which to compress WAL segments
	WALCompressionSize int
	// MaxMemoryRatio caps the maximum memory of this process. When the system
	// comes under memory pressure, it will start flushing table memstores.
	MaxMemoryRatio float64
	// Passthrough flags this node as a passthrough (won't store data in tables,
	// just WAL). Passthrough nodes will also outsource queries to specific
	// partition handlers. Requires that NumPartitions be specified.
	Passthrough bool
	// NumPartitions identifies how many partitions to split data from
	// passthrough nodes.
	NumPartitions int
	// Partition identies the partition owned by this follower
	Partition int
	// MaxFollowAge limits how far back to go when follower pulls data from
	// leader
	MaxFollowAge time.Duration
	// Follow is a function that allows a follower to request following a stream
	// from a passthrough node.
	Follow                     func(f func() *common.Follow, cb func(data []byte, newOffset wal.Offset) error)
	RegisterRemoteQueryHandler func(partition int, query planner.QueryClusterFN)
}

// DB is a zenodb database.
type DB struct {
	opts                 *DBOpts
	clock                vtime.Clock
	tables               map[string]*table
	orderedTables        []*table
	streams              map[string]*wal.WAL
	newStreamSubscriber  map[string]chan *tableWithOffset
	tablesMutex          sync.RWMutex
	isSorting            bool
	nextTableToSort      int
	memory               uint64
	flushMutex           sync.Mutex
	followerJoined       chan *follower
	processFollowersOnce sync.Once
	remoteQueryHandlers  map[int]chan planner.QueryClusterFN
	closed               bool
}

// NewDB creates a database using the given options.
func NewDB(opts *DBOpts) (*DB, error) {
	var err error
	db := &DB{
		opts:                opts,
		clock:               vtime.RealClock,
		tables:              make(map[string]*table),
		streams:             make(map[string]*wal.WAL),
		newStreamSubscriber: make(map[string]chan *tableWithOffset),
		followerJoined:      make(chan *follower, opts.NumPartitions),
		remoteQueryHandlers: make(map[int]chan planner.QueryClusterFN),
	}
	if opts.VirtualTime {
		db.clock = vtime.NewVirtualClock(time.Time{})
	}
	if opts.MaxWALSize <= 0 {
		opts.MaxWALSize = 10 * 1024768 // 10 MB
	}
	if opts.WALCompressionSize <= 0 {
		opts.WALCompressionSize = opts.MaxWALSize / 10
	}

	// Create db dir
	err = os.MkdirAll(opts.Dir, 0755)
	if err != nil && !os.IsExist(err) {
		return nil, fmt.Errorf("Unable to create db dir at %v: %v", opts.Dir, err)
	}

	if opts.EnableGeo {
		log.Debug("Enabling geolocation functions")
		err = geo.Init(filepath.Join(opts.Dir, "geoip.dat"))
		if err != nil {
			return nil, fmt.Errorf("Unable to initialize geo: %v", err)
		}
	}

	if opts.ISPProvider != nil {
		log.Debugf("Setting ISP provider to %v", opts.ISPProvider)
		isp.SetProvider(opts.ISPProvider)
	}

	if opts.AliasesFile != "" {
		registerAliases(opts.AliasesFile)
	}

	if opts.RedisClient != nil && opts.RedisCacheSize > 0 {
		log.Debug("Enabling redis expressions")
		geredis.Configure(opts.RedisClient, opts.RedisCacheSize)
	}

	if opts.SchemaFile != "" {
		err = db.pollForSchema(opts.SchemaFile)
		if err != nil {
			return nil, fmt.Errorf("Unable to apply schema: %v", err)
		}
	}
	log.Debugf("Dir: %v    SchemaFile: %v", opts.Dir, opts.SchemaFile)

	if db.opts.RegisterRemoteQueryHandler != nil {
		go db.opts.RegisterRemoteQueryHandler(db.opts.Partition, db.queryForRemote)
	}

	if db.opts.MaxMemoryRatio > 0 {
		log.Debugf("Limiting maximum memstore memory to %v", humanize.Bytes(db.maxMemoryBytes()))
	}
	go db.trackMemStats()

	return db, err
}

func (db *DB) Close() {
	log.Debug("Closing")
	db.tablesMutex.Lock()
	for name, stream := range db.streams {
		log.Debugf("Closing stream %v", name)
		stream.Close()
		delete(db.streams, name)
	}
	db.tablesMutex.Unlock()
}

func registerAliases(aliasesFile string) {
	log.Debugf("Registering aliases from file at %v", aliasesFile)

	file, err := os.Open(aliasesFile)
	if err != nil {
		log.Errorf("Unable to open aliases file at %v: %v", aliasesFile, err)
		return
	}
	defer file.Close()

	p, err := props.Read(file)
	if err != nil {
		log.Errorf("Unable to read aliases file at %v: %v", aliasesFile, err)
		return
	}

	for _, alias := range p.Names() {
		template := strings.TrimSpace(p.Get(alias))
		alias = strings.TrimSpace(alias)
		sql.RegisterAlias(alias, template)
		log.Debugf("Registered alias %v = %v", alias, template)
	}
}

// TableStats returns the TableStats for the named table.
func (db *DB) TableStats(table string) TableStats {
	t := db.getTable(table)
	if t == nil {
		return TableStats{}
	}
	t.statsMutex.RLock()
	defer t.statsMutex.RUnlock()
	return t.stats
}

// AllTableStats returns all TableStats for all tables, keyed to the table
// names.
func (db *DB) AllTableStats() map[string]TableStats {
	m := make(map[string]TableStats)
	tables := make(map[string]*table, 0)
	db.tablesMutex.RLock()
	for name, t := range db.tables {
		tables[name] = t
	}
	db.tablesMutex.RUnlock()
	for name, t := range tables {
		t.statsMutex.RLock()
		m[name] = t.stats
		t.statsMutex.RUnlock()
	}
	return m
}

// PrintTableStats prints the stats for the named table to a string.
func (db *DB) PrintTableStats(table string) string {
	stats := db.TableStats(table)
	now := db.clock.Now()
	return fmt.Sprintf("%v (%v)\tFiltered: %v    Queued: %v    Inserted: %v    Dropped: %v    Expired: %v",
		table,
		now.In(time.UTC),
		humanize.Comma(stats.FilteredPoints),
		humanize.Comma(stats.QueuedPoints),
		humanize.Comma(stats.InsertedPoints),
		humanize.Comma(stats.DroppedPoints),
		humanize.Comma(stats.ExpiredValues))
}

func (db *DB) getTable(table string) *table {
	db.tablesMutex.RLock()
	t := db.tables[strings.ToLower(table)]
	db.tablesMutex.RUnlock()
	return t
}

func (db *DB) now(table string) time.Time {
	return db.clock.Now()
}

func (db *DB) capWALAge(wal *wal.WAL) {
	for {
		time.Sleep(1 * time.Minute)
		err := wal.TruncateToSize(int64(db.opts.MaxWALSize))
		if err != nil {
			log.Errorf("Error truncating WAL: %v", err)
		}
		err = wal.CompressBeforeSize(int64(db.opts.WALCompressionSize))
		if err != nil {
			log.Errorf("Error compressing WAL: %v", err)
		}
	}
}

func (db *DB) trackMemStats() {
	for {
		db.updateMemStats()
		time.Sleep(2 * time.Second)
	}
}

func (db *DB) updateMemStats() {
	p, err := process.NewProcess(int32(os.Getpid()))
	if err != nil {
		log.Errorf("Unable to get process info: %v", err)
		return
	}
	mi, err := p.MemoryInfo()
	if err != nil {
		log.Errorf("Unable to get memory info for process: %v", err)
		return
	}
	memstats := &runtime.MemStats{}
	runtime.ReadMemStats(memstats)
	atomic.StoreUint64(&db.memory, memstats.Alloc)
	log.Debugf("Memory InUse: %v    Alloc: %v    Sys: %v     RSS: %v", humanize.Bytes(memstats.HeapInuse), humanize.Bytes(memstats.Alloc), humanize.Bytes(memstats.Sys), humanize.Bytes(mi.RSS))
}

func (db *DB) capMemStoreSize() {
	if db.opts.MaxMemoryRatio <= 0 {
		return
	}

	db.tablesMutex.RLock()
	sizes := make(byCurrentSize, 0, len(db.tables))
	for _, table := range db.tables {
		if !table.Virtual {
			sizes = append(sizes, &memStoreSize{table, table.memStoreSize()})
		}
	}
	db.tablesMutex.RUnlock()

	db.flushMutex.Lock()
	actual := atomic.LoadUint64(&db.memory)
	allowed := db.maxMemoryBytes()
	if actual > allowed {
		// Force flushing on the table with the largest memstore
		sort.Sort(sizes)
		log.Debugf("Memory usage of %v exceeds allowed %v, forcing flush on %v", humanize.Bytes(actual), humanize.Bytes(allowed), sizes[0].t.Name)
		sizes[0].t.forceFlush()
		db.updateMemStats()
		log.Debugf("Done forcing flush on %v", sizes[0].t.Name)
	}
	db.flushMutex.Unlock()
}

func (db *DB) maxMemoryBytes() uint64 {
	return uint64(systemRAM * db.opts.MaxMemoryRatio)
}

type memStoreSize struct {
	t    *table
	size int
}

type byCurrentSize []*memStoreSize

func (a byCurrentSize) Len() int           { return len(a) }
func (a byCurrentSize) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byCurrentSize) Less(i, j int) bool { return a[i].size > a[j].size }
<commit_msg>Forcing GC before forcing flush<commit_after>package zenodb

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"runtime/debug"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudfoundry/gosigar"
	"github.com/dustin/go-humanize"
	"github.com/getlantern/goexpr/geo"
	"github.com/getlantern/goexpr/isp"
	geredis "github.com/getlantern/goexpr/redis"
	"github.com/getlantern/golog"
	"github.com/getlantern/vtime"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb/common"
	"github.com/getlantern/zenodb/planner"
	"github.com/getlantern/zenodb/sql"
	"github.com/rickar/props"
	"github.com/shirou/gopsutil/process"
	"gopkg.in/redis.v5"
)

var (
	log = golog.LoggerFor("zenodb")

	systemRAM float64
)

func init() {
	mem := sigar.Mem{}
	err := mem.Get()
	if err != nil {
		panic(fmt.Sprintf("Unable to get system memory: %v", err))
	}
	systemRAM = float64(mem.Total)
}

// DBOpts provides options for configuring the database.
type DBOpts struct {
	// Dir points at the directory that contains the data files.
	Dir string
	// SchemaFile points at a YAML schema file that configures the tables and
	// views in the database.
	SchemaFile string
	// AliasesFile points at a file that contains expression aliases in the form
	// name=template(%v, %v), with one alias per line.
	AliasesFile string
	// EnableGeo enables geolocation functions
	EnableGeo bool
	// ISPProvider configures a provider of ISP lookups. Specify this to allow the
	// use of ISP functions.
	ISPProvider isp.Provider
	// RedisClient provides a connection to redis which enables the use of Redis
	// expressions like HGET.
	RedisClient *redis.Client
	// RedisCacheSize controls the size of redis hash caches
	RedisCacheSize int
	// VirtualTime, if true, tells zenodb to use a virtual clock that advances
	// based on the timestamps of Points received via inserts.
	VirtualTime bool
	// WALSyncInterval governs how frequently to sync the WAL to disk. 0 means
	// it syncs after every write (which is not great for performance).
	WALSyncInterval time.Duration
	// MaxWALSize limits how much WAL data to keep (in bytes)
	MaxWALSize int
	// WALCompressionSize specifies the size beyond which to compress WAL segments
	WALCompressionSize int
	// MaxMemoryRatio caps the maximum memory of this process. When the system
	// comes under memory pressure, it will start flushing table memstores.
	MaxMemoryRatio float64
	// Passthrough flags this node as a passthrough (won't store data in tables,
	// just WAL). Passthrough nodes will also outsource queries to specific
	// partition handlers. Requires that NumPartitions be specified.
	Passthrough bool
	// NumPartitions identifies how many partitions to split data from
	// passthrough nodes.
	NumPartitions int
	// Partition identies the partition owned by this follower
	Partition int
	// MaxFollowAge limits how far back to go when follower pulls data from
	// leader
	MaxFollowAge time.Duration
	// Follow is a function that allows a follower to request following a stream
	// from a passthrough node.
	Follow                     func(f func() *common.Follow, cb func(data []byte, newOffset wal.Offset) error)
	RegisterRemoteQueryHandler func(partition int, query planner.QueryClusterFN)
}

// DB is a zenodb database.
type DB struct {
	opts                 *DBOpts
	clock                vtime.Clock
	tables               map[string]*table
	orderedTables        []*table
	streams              map[string]*wal.WAL
	newStreamSubscriber  map[string]chan *tableWithOffset
	tablesMutex          sync.RWMutex
	isSorting            bool
	nextTableToSort      int
	memory               uint64
	flushMutex           sync.Mutex
	followerJoined       chan *follower
	processFollowersOnce sync.Once
	remoteQueryHandlers  map[int]chan planner.QueryClusterFN
	closed               bool
}

// NewDB creates a database using the given options.
func NewDB(opts *DBOpts) (*DB, error) {
	var err error
	db := &DB{
		opts:                opts,
		clock:               vtime.RealClock,
		tables:              make(map[string]*table),
		streams:             make(map[string]*wal.WAL),
		newStreamSubscriber: make(map[string]chan *tableWithOffset),
		followerJoined:      make(chan *follower, opts.NumPartitions),
		remoteQueryHandlers: make(map[int]chan planner.QueryClusterFN),
	}
	if opts.VirtualTime {
		db.clock = vtime.NewVirtualClock(time.Time{})
	}
	if opts.MaxWALSize <= 0 {
		opts.MaxWALSize = 10 * 1024768 // 10 MB
	}
	if opts.WALCompressionSize <= 0 {
		opts.WALCompressionSize = opts.MaxWALSize / 10
	}

	// Create db dir
	err = os.MkdirAll(opts.Dir, 0755)
	if err != nil && !os.IsExist(err) {
		return nil, fmt.Errorf("Unable to create db dir at %v: %v", opts.Dir, err)
	}

	if opts.EnableGeo {
		log.Debug("Enabling geolocation functions")
		err = geo.Init(filepath.Join(opts.Dir, "geoip.dat"))
		if err != nil {
			return nil, fmt.Errorf("Unable to initialize geo: %v", err)
		}
	}

	if opts.ISPProvider != nil {
		log.Debugf("Setting ISP provider to %v", opts.ISPProvider)
		isp.SetProvider(opts.ISPProvider)
	}

	if opts.AliasesFile != "" {
		registerAliases(opts.AliasesFile)
	}

	if opts.RedisClient != nil && opts.RedisCacheSize > 0 {
		log.Debug("Enabling redis expressions")
		geredis.Configure(opts.RedisClient, opts.RedisCacheSize)
	}

	if opts.SchemaFile != "" {
		err = db.pollForSchema(opts.SchemaFile)
		if err != nil {
			return nil, fmt.Errorf("Unable to apply schema: %v", err)
		}
	}
	log.Debugf("Dir: %v    SchemaFile: %v", opts.Dir, opts.SchemaFile)

	if db.opts.RegisterRemoteQueryHandler != nil {
		go db.opts.RegisterRemoteQueryHandler(db.opts.Partition, db.queryForRemote)
	}

	if db.opts.MaxMemoryRatio > 0 {
		log.Debugf("Limiting maximum memstore memory to %v", humanize.Bytes(db.maxMemoryBytes()))
	}
	go db.trackMemStats()

	return db, err
}

func (db *DB) Close() {
	log.Debug("Closing")
	db.tablesMutex.Lock()
	for name, stream := range db.streams {
		log.Debugf("Closing stream %v", name)
		stream.Close()
		delete(db.streams, name)
	}
	db.tablesMutex.Unlock()
}

func registerAliases(aliasesFile string) {
	log.Debugf("Registering aliases from file at %v", aliasesFile)

	file, err := os.Open(aliasesFile)
	if err != nil {
		log.Errorf("Unable to open aliases file at %v: %v", aliasesFile, err)
		return
	}
	defer file.Close()

	p, err := props.Read(file)
	if err != nil {
		log.Errorf("Unable to read aliases file at %v: %v", aliasesFile, err)
		return
	}

	for _, alias := range p.Names() {
		template := strings.TrimSpace(p.Get(alias))
		alias = strings.TrimSpace(alias)
		sql.RegisterAlias(alias, template)
		log.Debugf("Registered alias %v = %v", alias, template)
	}
}

// TableStats returns the TableStats for the named table.
func (db *DB) TableStats(table string) TableStats {
	t := db.getTable(table)
	if t == nil {
		return TableStats{}
	}
	t.statsMutex.RLock()
	defer t.statsMutex.RUnlock()
	return t.stats
}

// AllTableStats returns all TableStats for all tables, keyed to the table
// names.
func (db *DB) AllTableStats() map[string]TableStats {
	m := make(map[string]TableStats)
	tables := make(map[string]*table, 0)
	db.tablesMutex.RLock()
	for name, t := range db.tables {
		tables[name] = t
	}
	db.tablesMutex.RUnlock()
	for name, t := range tables {
		t.statsMutex.RLock()
		m[name] = t.stats
		t.statsMutex.RUnlock()
	}
	return m
}

// PrintTableStats prints the stats for the named table to a string.
func (db *DB) PrintTableStats(table string) string {
	stats := db.TableStats(table)
	now := db.clock.Now()
	return fmt.Sprintf("%v (%v)\tFiltered: %v    Queued: %v    Inserted: %v    Dropped: %v    Expired: %v",
		table,
		now.In(time.UTC),
		humanize.Comma(stats.FilteredPoints),
		humanize.Comma(stats.QueuedPoints),
		humanize.Comma(stats.InsertedPoints),
		humanize.Comma(stats.DroppedPoints),
		humanize.Comma(stats.ExpiredValues))
}

func (db *DB) getTable(table string) *table {
	db.tablesMutex.RLock()
	t := db.tables[strings.ToLower(table)]
	db.tablesMutex.RUnlock()
	return t
}

func (db *DB) now(table string) time.Time {
	return db.clock.Now()
}

func (db *DB) capWALAge(wal *wal.WAL) {
	for {
		time.Sleep(1 * time.Minute)
		err := wal.TruncateToSize(int64(db.opts.MaxWALSize))
		if err != nil {
			log.Errorf("Error truncating WAL: %v", err)
		}
		err = wal.CompressBeforeSize(int64(db.opts.WALCompressionSize))
		if err != nil {
			log.Errorf("Error compressing WAL: %v", err)
		}
	}
}

func (db *DB) trackMemStats() {
	for {
		db.updateMemStats()
		time.Sleep(2 * time.Second)
	}
}

func (db *DB) updateMemStats() {
	p, err := process.NewProcess(int32(os.Getpid()))
	if err != nil {
		log.Errorf("Unable to get process info: %v", err)
		return
	}
	mi, err := p.MemoryInfo()
	if err != nil {
		log.Errorf("Unable to get memory info for process: %v", err)
		return
	}
	memstats := &runtime.MemStats{}
	runtime.ReadMemStats(memstats)
	atomic.StoreUint64(&db.memory, memstats.Alloc)
	log.Debugf("Memory InUse: %v    Alloc: %v    Sys: %v     RSS: %v", humanize.Bytes(memstats.HeapInuse), humanize.Bytes(memstats.Alloc), humanize.Bytes(memstats.Sys), humanize.Bytes(mi.RSS))
}

func (db *DB) capMemStoreSize() {
	if db.opts.MaxMemoryRatio <= 0 {
		return
	}

	db.tablesMutex.RLock()
	sizes := make(byCurrentSize, 0, len(db.tables))
	for _, table := range db.tables {
		if !table.Virtual {
			sizes = append(sizes, &memStoreSize{table, table.memStoreSize()})
		}
	}
	db.tablesMutex.RUnlock()

	db.flushMutex.Lock()
	actual := atomic.LoadUint64(&db.memory)
	allowed := db.maxMemoryBytes()
	if actual > allowed {
		// First try to regain memory with GC
		log.Debugf("Memory usage of %v exceeds allowed %v, forcing GC", humanize.Bytes(actual), humanize.Bytes(allowed))
		debug.FreeOSMemory()
		db.updateMemStats()
	}
	actual = atomic.LoadUint64(&db.memory)
	if actual > allowed {
		// Force flushing on the table with the largest memstore
		sort.Sort(sizes)
		log.Debugf("Memory usage of %v exceeds allowed %v even after GC, forcing flush on %v", humanize.Bytes(actual), humanize.Bytes(allowed), sizes[0].t.Name)
		sizes[0].t.forceFlush()
		db.updateMemStats()
		log.Debugf("Done forcing flush on %v", sizes[0].t.Name)
	}
	db.flushMutex.Unlock()
}

func (db *DB) maxMemoryBytes() uint64 {
	return uint64(systemRAM * db.opts.MaxMemoryRatio)
}

type memStoreSize struct {
	t    *table
	size int
}

type byCurrentSize []*memStoreSize

func (a byCurrentSize) Len() int           { return len(a) }
func (a byCurrentSize) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byCurrentSize) Less(i, j int) bool { return a[i].size > a[j].size }
<|endoftext|>