<commit_before>package client

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"

	"github.com/flynn/go-discoverd"
	"github.com/flynn/rpcplus"
	"github.com/flynn/sampi/types"
)

func New() (*Client, error) {
	services, err := discoverd.Services("flynn-sampi", discoverd.DefaultTimeout)
	if err != nil {
		return nil, err
	}
	if len(services) == 0 {
		return nil, errors.New("sampi: no servers found")
	}
	c, err := rpcplus.DialHTTP("tcp", services[0].Addr)
	return &Client{c}, err
}

type Client struct {
	c *rpcplus.Client
}

func (c *Client) State() (map[string]sampi.Host, error) {
	var state map[string]sampi.Host
	err := c.c.Call("Scheduler.State", struct{}{}, &state)
	return state, err
}

func (c *Client) Schedule(req *sampi.ScheduleReq) (*sampi.ScheduleRes, error) {
	var res sampi.ScheduleRes
	err := c.c.Call("Scheduler.Schedule", req, &res)
	return &res, err
}

func (c *Client) RegisterHost(host *sampi.Host, stream chan *sampi.Job) *error {
	return &c.c.StreamGo("Scheduler.RegisterHost", host, stream).Error
}

func (c *Client) RemoveJobs(jobIDs []string) error {
	return c.c.Call("Scheduler.RemoveJobs", jobIDs, &struct{}{})
}

func RandomJobID(prefix string) string { return prefix + randomID() }

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<commit_msg>host/sampi: using explicit leader api<commit_after>package client

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
	"time"

	"github.com/flynn/go-discoverd"
	"github.com/flynn/rpcplus"
	"github.com/flynn/sampi/types"
)

func New() (*Client, error) {
	services, err := discoverd.NewServiceSet("flynn-sampi")
	if err != nil {
		return nil, err
	}
	select {
	case <-services.Watch(true, true):
	case <-time.After(time.Second):
		return nil, errors.New("sampi: no servers found")
	}
	c, err := rpcplus.DialHTTP("tcp", services.Leader().Addr)
	return &Client{c}, err
}

type Client struct {
	c *rpcplus.Client
}

func (c *Client) State() (map[string]sampi.Host, error) {
	var state map[string]sampi.Host
	err := c.c.Call("Scheduler.State", struct{}{}, &state)
	return state, err
}

func (c *Client) Schedule(req *sampi.ScheduleReq) (*sampi.ScheduleRes, error) {
	var res sampi.ScheduleRes
	err := c.c.Call("Scheduler.Schedule", req, &res)
	return &res, err
}

func (c *Client) RegisterHost(host *sampi.Host, stream chan *sampi.Job) *error {
	return &c.c.StreamGo("Scheduler.RegisterHost", host, stream).Error
}

func (c *Client) RemoveJobs(jobIDs []string) error {
	return c.c.Call("Scheduler.RemoveJobs", jobIDs, &struct{}{})
}

func RandomJobID(prefix string) string { return prefix + randomID() }

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed 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 indexer

//Max number of vbuckets supported in the system
const MAX_NUM_VBUCKETS = 1024

//Supervisor's channel capacity to buffer requests
//from workers
const WORKER_MSG_QUEUE_LEN = 100000

// Default cluster address
const DEFAULT_CLUSTER_ENDPOINT = "127.0.0.1:9000"

const LOCALHOST = "127.0.0.1"

//Maintenance Topic Name
const MAINT_TOPIC = "MAINT_STREAM_TOPIC"

//Catchup Topic Name
const CATCHUP_TOPIC = "CATCHUP_STREAM_TOPIC"

//Initial Stream Topic Name
const INIT_TOPIC = "INIT_STREAM_TOPIC"

//Default Pool Name
const DEFAULT_POOL = "default"

//Default Number of Workers started by a stream reader
//to processed incoming mutation. Max can be upto the
//number of vbuckets and minimum must be equal to the
//number of vbuckets
const DEFAULT_NUM_STREAM_READER_WORKERS = 8

//Buffer for each of stream reader worker to queue
//up mutations before processing
const MAX_STREAM_READER_WORKER_BUFFER = 1000

//Max number of snapshot to be retained per index.
//Older snapshots are deleted.
const MAX_SNAPSHOTS_PER_INDEX = 5

//Slab Manager Specific constants
const DEFAULT_START_CHUNK_SIZE = 256
const DEFAULT_SLAB_SIZE = DEFAULT_START_CHUNK_SIZE * 1024
const DEFAULT_MAX_SLAB_MEMORY = DEFAULT_SLAB_SIZE * 1024

//Internal Buffer Size for Each Slice to store incoming
//requests
const SLICE_COMMAND_BUFFER_SIZE = 10000

//Time in milliseconds for a slice to poll for
//any outstanding writes before commit
const SLICE_COMMIT_POLL_INTERVAL = 20

//Max Length of Secondary Key
const MAX_SEC_KEY_LEN = 4096

//Buffer Length for encoded Sec Key
const MAX_SEC_KEY_BUFFER_LEN = MAX_SEC_KEY_LEN * 3

const INDEXER_ID_KEY = "IndexerId"
<commit_msg>indexer: increase worker message buffer capacity<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed 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 indexer

//Max number of vbuckets supported in the system
const MAX_NUM_VBUCKETS = 1024

//Supervisor's channel capacity to buffer requests
//from workers
const WORKER_MSG_QUEUE_LEN = 1000000

// Default cluster address
const DEFAULT_CLUSTER_ENDPOINT = "127.0.0.1:9000"

const LOCALHOST = "127.0.0.1"

//Maintenance Topic Name
const MAINT_TOPIC = "MAINT_STREAM_TOPIC"

//Catchup Topic Name
const CATCHUP_TOPIC = "CATCHUP_STREAM_TOPIC"

//Initial Stream Topic Name
const INIT_TOPIC = "INIT_STREAM_TOPIC"

//Default Pool Name
const DEFAULT_POOL = "default"

//Default Number of Workers started by a stream reader
//to processed incoming mutation. Max can be upto the
//number of vbuckets and minimum must be equal to the
//number of vbuckets
const DEFAULT_NUM_STREAM_READER_WORKERS = 8

//Buffer for each of stream reader worker to queue
//up mutations before processing
const MAX_STREAM_READER_WORKER_BUFFER = 1000

//Max number of snapshot to be retained per index.
//Older snapshots are deleted.
const MAX_SNAPSHOTS_PER_INDEX = 5

//Slab Manager Specific constants
const DEFAULT_START_CHUNK_SIZE = 256
const DEFAULT_SLAB_SIZE = DEFAULT_START_CHUNK_SIZE * 1024
const DEFAULT_MAX_SLAB_MEMORY = DEFAULT_SLAB_SIZE * 1024

//Internal Buffer Size for Each Slice to store incoming
//requests
const SLICE_COMMAND_BUFFER_SIZE = 10000

//Time in milliseconds for a slice to poll for
//any outstanding writes before commit
const SLICE_COMMIT_POLL_INTERVAL = 20

//Max Length of Secondary Key
const MAX_SEC_KEY_LEN = 4096

//Buffer Length for encoded Sec Key
const MAX_SEC_KEY_BUFFER_LEN = MAX_SEC_KEY_LEN * 3

const INDEXER_ID_KEY = "IndexerId"
<|endoftext|>


<commit_before>package scheduler

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/nyaxt/otaru/logger"
	"github.com/nyaxt/otaru/util"
)

var mylog = logger.Registry().Category("scheduler")

type Result interface {
	Err() error
}

type ErrorResult struct {
	Error error
}

func (er ErrorResult) Err() error    { return er.Error }
func (ErrorResult) ImplName() string { return "ErrorResult" }

func (er ErrorResult) MarshalJSON() ([]byte, error) {
	type jsonRep struct {
		ErrorString string `json:"error_str,omitempty"`
	}

	var errstr string
	if er.Error != nil {
		errstr = er.Error.Error()
	}

	return json.Marshal(jsonRep{ErrorString: errstr})
}

type Task interface {
	Run(ctx context.Context) Result
}

func describeTask(t Task) string {
	type stringifier interface {
		String() string
	}
	if sf, ok := t.(stringifier); ok {
		return sf.String()
	}
	return util.TryGetImplName(t)
}

type State int32

const (
	JobScheduled = iota
	JobStarted   = iota
	JobFinished  = iota
	JobAborted   = iota
)

func (st State) String() string {
	switch st {
	case JobScheduled:
		return "JobScheduled"
	case JobStarted:
		return "JobStarted"
	case JobFinished:
		return "JobFinished"
	case JobAborted:
		return "JobAborted"
	default:
		return "UnknownJobState"
	}
}

var ZeroTime time.Time

type job struct {
	ID
	State
	// Issuer string

	CreatedAt   time.Time
	ScheduledAt time.Time
	StartedAt   time.Time
	FinishedAt  time.Time

	Task
	Result
	DoneCallback

	mu         sync.Mutex
	cancelfn   context.CancelFunc
	scheduledC chan struct{}
}

func (j *job) String() string {
	return fmt.Sprintf("job{ID: %d, CreatedAt: %v, Task: %s}",
		j.ID,
		j.CreatedAt,
		describeTask(j.Task),
	)
}

type JobView struct {
	ID    `json:"id"`
	State `json:"state"`
	// Issuer string

	TaskDesc string `json:"task_desc"`

	CreatedAt   time.Time `json:"created_at"`
	ScheduledAt time.Time `json:"scheduled_at"`
	StartedAt   time.Time `json:"started_at"`
	FinishedAt  time.Time `json:"finished_at"`

	Result `json:"result"`
}

type DoneCallback func(*JobView)

func (j *job) ViewWithLock() *JobView {
	return &JobView{
		ID:          j.ID,
		State:       j.State,
		TaskDesc:    describeTask(j.Task),
		CreatedAt:   j.CreatedAt,
		ScheduledAt: j.ScheduledAt,
		StartedAt:   j.StartedAt,
		FinishedAt:  j.FinishedAt,
		Result:      j.Result,
	}
}

func (j *job) View() *JobView {
	j.mu.Lock()
	defer j.mu.Unlock()
	return j.ViewWithLock()
}

type jobQuery struct {
	ID
	resultC chan []*JobView
}

type abortReq struct {
	ID
	doneC chan struct{}
}

type Scheduler struct {
	numRunners  int
	numWaitJobs int

	idGen

	scheduleC     chan *job
	queryC        chan *jobQuery
	abortC        chan *abortReq
	runC          chan *job
	forceGCC      chan struct{}
	joinScheduleC chan struct{}
	joinRunnerC   chan struct{}

	ZombiePeriod time.Duration
}

func NewScheduler() *Scheduler {
	s := &Scheduler{
		numRunners:    4, // FIXME
		numWaitJobs:   0,
		idGen:         idGen{0},
		scheduleC:     make(chan *job, 1),
		queryC:        make(chan *jobQuery, 1),
		abortC:        make(chan *abortReq, 1),
		runC:          make(chan *job, 8),
		forceGCC:      make(chan struct{}),
		joinScheduleC: make(chan struct{}),
		joinRunnerC:   make(chan struct{}),

		ZombiePeriod: 30 * time.Second,
	}

	go s.schedulerMain()
	for i := 0; i < s.numRunners; i++ {
		go s.runnerMain()
	}

	return s
}

type Stats struct {
	NumRunners  int `json:"num_runners"`
	NumWaitJobs int `json:"num_wait_jobs"`
}

func (s *Scheduler) GetStats() *Stats {
	return &Stats{
		NumRunners:  s.numRunners,
		NumWaitJobs: s.numWaitJobs,
	}
}

func (s *Scheduler) RunAt(task Task, at time.Time, cb DoneCallback) ID {
	id := s.idGen.genID()
	j := &job{
		ID:           id,
		CreatedAt:    time.Now(),
		ScheduledAt:  at,
		Task:         task,
		DoneCallback: cb,
		scheduledC:   make(chan struct{}),
	}

	s.scheduleC <- j
	<-j.scheduledC

	return id
}

func (s *Scheduler) RunImmediately(task Task, cb DoneCallback) ID {
	return s.RunAt(task, ZeroTime, cb)
}

func (s *Scheduler) RunImmediatelyBlock(task Task) *JobView {
	doneC := make(chan *JobView)
	s.RunAt(task, ZeroTime, func(v *JobView) { doneC <- v; close(doneC) })
	return <-doneC
}

func (s *Scheduler) Query(id ID) *JobView {
	q := &jobQuery{
		ID:      id,
		resultC: make(chan []*JobView),
	}
	s.queryC <- q
	rs := <-q.resultC
	if rs == nil {
		return nil
	}
	return rs[0]
}

func (s *Scheduler) QueryAll() []*JobView {
	q := &jobQuery{
		ID:      allJobs,
		resultC: make(chan []*JobView),
	}
	s.queryC <- q
	return <-q.resultC
}

func (s *Scheduler) abortInternal(id ID) {
	req := &abortReq{
		ID:    id,
		doneC: make(chan struct{}),
	}
	s.abortC <- req
	<-req.doneC
}

func (s *Scheduler) Abort(id ID) {
	if id == allJobs {
		// allJobs should only be used internally
		return
	}

	s.abortInternal(id)
}

func (s *Scheduler) stop() {
	logger.Infof(mylog, "scheduler stop start")
	close(s.scheduleC)
	close(s.queryC)
	close(s.abortC)

	<-s.joinScheduleC
	for i := 0; i < s.numRunners; i++ {
		<-s.joinRunnerC
	}
	close(s.forceGCC)
	logger.Infof(mylog, "scheduler stop done")
}

func (s *Scheduler) RunAllAndStop() { s.stop() }

func (s *Scheduler) AbortAllAndStop() {
	s.abortInternal(allJobs)
	s.stop()
}

func abortJob(j *job) {
	j.mu.Lock()
	switch j.State {
	case JobScheduled:
		j.State = JobAborted
		if j.DoneCallback != nil {
			go j.DoneCallback(j.ViewWithLock())
		}
	case JobStarted:
		j.cancelfn()
	case JobFinished:
		// Job has already finished. Too late.
	case JobAborted:
		// Job is already aborted. Nothing to do.
	}
	j.mu.Unlock()
}

func (s *Scheduler) ForceGC() { s.forceGCC <- struct{}{} }

const gcPeriod = 10 * time.Second

func (s *Scheduler) schedulerMain() {
	var tickC <-chan time.Time
	tickC = make(chan time.Time)
	var nextWakeUp time.Time
	var timer *time.Timer
	renewTimer := func() {
		d := nextWakeUp.Sub(time.Now())
		if timer == nil {
			timer = time.NewTimer(d)
		} else {
			timer.Stop()
			timer.Reset(d)
		}
		tickC = timer.C
	}

	waitJobs := make([]*job, 0)
	jobs := make(map[ID]*job)

	gcTickC := time.Tick(gcPeriod)
	doGC := func() {
		logger.Debugf(mylog, "JobGC start.")
		gcstart := time.Now()
		threshold := gcstart.Add(-s.ZombiePeriod)

		oldJobIDs := make([]ID, 0)
		for id, j := range jobs {
			j.mu.Lock()
			if j.State == JobAborted || j.State == JobFinished {
				if j.FinishedAt.Before(threshold) {
					logger.Debugf(mylog, "GC-ing job %v: %v since aborted/finished.", j, gcstart.Sub(j.FinishedAt))
					oldJobIDs = append(oldJobIDs, id)
				}
			}
			j.mu.Unlock()
		}
		for _, id := range oldJobIDs {
			delete(jobs, id)
		}

		logger.Debugf(mylog, "JobGC end. Took %v. Deleted %d entries", time.Since(gcstart), len(oldJobIDs))
	}

	defer func() {
		close(s.runC)
		s.joinScheduleC <- struct{}{}
	}()

	scheduleC := s.scheduleC
	for {
		if scheduleC == nil && s.numWaitJobs == 0 {
			return
		}

		select {
		case j, more := <-scheduleC:
			if !more {
				// stop polling on scheduleC
				scheduleC = nil
				continue
			}

			if _, ok := jobs[j.ID]; ok {
				logger.Criticalf(mylog, "job ID %v is already taken. received duplicate: %v", j.ID, j)
				if j.scheduledC != nil {
					close(j.scheduledC)
				}
				continue
			}
			jobs[j.ID] = j

			now := time.Now()
			if j.ScheduledAt.Before(now) {
				s.runC <- j
			} else {
				if nextWakeUp.IsZero() || nextWakeUp.After(j.ScheduledAt) {
					nextWakeUp = j.ScheduledAt
					renewTimer()
				}

				waitJobs = append(waitJobs, j)
				s.numWaitJobs = len(waitJobs)
			}
			if j.scheduledC != nil {
				close(j.scheduledC)
			}

		case q := <-s.queryC:
			if q == nil {
				continue
			}
			if q.ID == allJobs {
				jvs := make([]*JobView, 0, len(jobs))
				for _, j := range jobs {
					jvs = append(jvs, j.View())
				}
				q.resultC <- jvs
				continue
			}

			id := q.ID
			j, ok := jobs[id]
			if !ok {
				q.resultC <- nil
			} else {
				q.resultC <- []*JobView{j.View()}
			}

		case req := <-s.abortC:
			if req == nil {
				continue
			}

			if req.ID == allJobs {
				for _, j := range jobs {
					abortJob(j)
				}

				req.doneC <- struct{}{}
				continue
			}

			j, ok := jobs[req.ID]
			if !ok {
				logger.Warningf(mylog, "Abort target job ID %d doesn't exist.", req.ID)
				req.doneC <- struct{}{}
				continue
			}
			abortJob(j)
			req.doneC <- struct{}{}

		case <-tickC:
			stillWaitJobs := make([]*job, 0, len(waitJobs))
			now := time.Now()
			var zero time.Time
			nextWakeUp = zero
			for _, j := range waitJobs {
				if j.ScheduledAt.Before(now) {
					s.runC <- j
				} else {
					if nextWakeUp.IsZero() || nextWakeUp.After(j.ScheduledAt) {
						nextWakeUp = j.ScheduledAt
					}

					stillWaitJobs = append(stillWaitJobs, j)
				}
			}
			waitJobs = stillWaitJobs
			s.numWaitJobs = len(waitJobs)

			renewTimer()

		case <-gcTickC:
			doGC()
		case <-s.forceGCC:
			doGC()
		}
	}
}

func (s *Scheduler) runnerMain() {
	for j := range s.runC {
		j.mu.Lock()
		if j.State == JobAborted {
			j.mu.Unlock()
			continue
		}
		if j.State != JobScheduled {
			logger.Infof(mylog, "Skipping job not in scheduled state: %v", j)

			j.mu.Unlock()
			continue
		}

		task := j.Task
		ctx, cancelfn := context.WithCancel(context.Background())
		j.cancelfn = cancelfn
		j.StartedAt = time.Now()
		j.State = JobStarted

		j.mu.Unlock()
		logger.Debugf(mylog, "About to run job %v", j)
		result := task.Run(ctx)
		if result != nil {
			if err := result.Err(); err != nil {
				logger.Warningf(mylog, "Job %v failed with error: %v", j, err)
			}
		}
		finishedAt := time.Now()
		j.mu.Lock()

		j.Result = result
		j.FinishedAt = finishedAt
		if ctx.Err() != nil {
			j.State = JobAborted
		} else {
			j.State = JobFinished
		}
		if j.DoneCallback != nil {
			go j.DoneCallback(j.ViewWithLock())
		}
		j.mu.Unlock()
	}
	s.joinRunnerC <- struct{}{}
}
<commit_msg>don't count aborted jobs as waitJobs. Fix #36<commit_after>package scheduler

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/nyaxt/otaru/logger"
	"github.com/nyaxt/otaru/util"
)

var mylog = logger.Registry().Category("scheduler")

type Result interface {
	Err() error
}

type ErrorResult struct {
	Error error
}

func (er ErrorResult) Err() error    { return er.Error }
func (ErrorResult) ImplName() string { return "ErrorResult" }

func (er ErrorResult) MarshalJSON() ([]byte, error) {
	type jsonRep struct {
		ErrorString string `json:"error_str,omitempty"`
	}

	var errstr string
	if er.Error != nil {
		errstr = er.Error.Error()
	}

	return json.Marshal(jsonRep{ErrorString: errstr})
}

type Task interface {
	Run(ctx context.Context) Result
}

func describeTask(t Task) string {
	type stringifier interface {
		String() string
	}
	if sf, ok := t.(stringifier); ok {
		return sf.String()
	}
	return util.TryGetImplName(t)
}

type State int32

const (
	JobScheduled = iota
	JobStarted   = iota
	JobFinished  = iota
	JobAborted   = iota
)

func (st State) String() string {
	switch st {
	case JobScheduled:
		return "JobScheduled"
	case JobStarted:
		return "JobStarted"
	case JobFinished:
		return "JobFinished"
	case JobAborted:
		return "JobAborted"
	default:
		return "UnknownJobState"
	}
}

var ZeroTime time.Time

type job struct {
	ID
	State
	// Issuer string

	CreatedAt   time.Time
	ScheduledAt time.Time
	StartedAt   time.Time
	FinishedAt  time.Time

	Task
	Result
	DoneCallback

	mu         sync.Mutex
	cancelfn   context.CancelFunc
	scheduledC chan struct{}
}

func (j *job) String() string {
	return fmt.Sprintf("job{ID: %d, CreatedAt: %v, Task: %s, State: %v}",
		j.ID,
		j.CreatedAt,
		describeTask(j.Task),
		j.State,
	)
}

type JobView struct {
	ID    `json:"id"`
	State `json:"state"`
	// Issuer string

	TaskDesc string `json:"task_desc"`

	CreatedAt   time.Time `json:"created_at"`
	ScheduledAt time.Time `json:"scheduled_at"`
	StartedAt   time.Time `json:"started_at"`
	FinishedAt  time.Time `json:"finished_at"`

	Result `json:"result"`
}

type DoneCallback func(*JobView)

func (j *job) ViewWithLock() *JobView {
	return &JobView{
		ID:          j.ID,
		State:       j.State,
		TaskDesc:    describeTask(j.Task),
		CreatedAt:   j.CreatedAt,
		ScheduledAt: j.ScheduledAt,
		StartedAt:   j.StartedAt,
		FinishedAt:  j.FinishedAt,
		Result:      j.Result,
	}
}

func (j *job) View() *JobView {
	j.mu.Lock()
	defer j.mu.Unlock()
	return j.ViewWithLock()
}

type jobQuery struct {
	ID
	resultC chan []*JobView
}

type abortReq struct {
	ID
	doneC chan struct{}
}

type Scheduler struct {
	numRunners  int
	numWaitJobs int

	idGen

	scheduleC     chan *job
	queryC        chan *jobQuery
	abortC        chan *abortReq
	runC          chan *job
	forceGCC      chan struct{}
	joinScheduleC chan struct{}
	joinRunnerC   chan struct{}

	ZombiePeriod time.Duration
}

func NewScheduler() *Scheduler {
	s := &Scheduler{
		numRunners:    4, // FIXME
		numWaitJobs:   0,
		idGen:         idGen{0},
		scheduleC:     make(chan *job, 1),
		queryC:        make(chan *jobQuery, 1),
		abortC:        make(chan *abortReq, 1),
		runC:          make(chan *job, 8),
		forceGCC:      make(chan struct{}),
		joinScheduleC: make(chan struct{}),
		joinRunnerC:   make(chan struct{}),

		ZombiePeriod: 30 * time.Second,
	}

	go s.schedulerMain()
	for i := 0; i < s.numRunners; i++ {
		go s.runnerMain()
	}

	return s
}

type Stats struct {
	NumRunners  int `json:"num_runners"`
	NumWaitJobs int `json:"num_wait_jobs"`
}

func (s *Scheduler) GetStats() *Stats {
	return &Stats{
		NumRunners:  s.numRunners,
		NumWaitJobs: s.numWaitJobs,
	}
}

func (s *Scheduler) RunAt(task Task, at time.Time, cb DoneCallback) ID {
	id := s.idGen.genID()
	j := &job{
		ID:           id,
		CreatedAt:    time.Now(),
		ScheduledAt:  at,
		Task:         task,
		DoneCallback: cb,
		scheduledC:   make(chan struct{}),
	}

	s.scheduleC <- j
	<-j.scheduledC

	return id
}

func (s *Scheduler) RunImmediately(task Task, cb DoneCallback) ID {
	return s.RunAt(task, ZeroTime, cb)
}

func (s *Scheduler) RunImmediatelyBlock(task Task) *JobView {
	doneC := make(chan *JobView)
	s.RunAt(task, ZeroTime, func(v *JobView) { doneC <- v; close(doneC) })
	return <-doneC
}

func (s *Scheduler) Query(id ID) *JobView {
	q := &jobQuery{
		ID:      id,
		resultC: make(chan []*JobView),
	}
	s.queryC <- q
	rs := <-q.resultC
	if rs == nil {
		return nil
	}
	return rs[0]
}

func (s *Scheduler) QueryAll() []*JobView {
	q := &jobQuery{
		ID:      allJobs,
		resultC: make(chan []*JobView),
	}
	s.queryC <- q
	return <-q.resultC
}

func (s *Scheduler) abortInternal(id ID) {
	req := &abortReq{
		ID:    id,
		doneC: make(chan struct{}),
	}
	s.abortC <- req
	<-req.doneC
}

func (s *Scheduler) Abort(id ID) {
	if id == allJobs {
		// allJobs should only be used internally
		return
	}

	s.abortInternal(id)
}

func (s *Scheduler) stop() {
	logger.Infof(mylog, "scheduler stop start")
	close(s.scheduleC)
	close(s.queryC)
	close(s.abortC)

	<-s.joinScheduleC
	for i := 0; i < s.numRunners; i++ {
		<-s.joinRunnerC
	}
	close(s.forceGCC)
	logger.Infof(mylog, "scheduler stop done")
}

func (s *Scheduler) RunAllAndStop() { s.stop() }

func (s *Scheduler) AbortAllAndStop() {
	s.abortInternal(allJobs)
	s.stop()
}

func abortJob(j *job) {
	j.mu.Lock()
	switch j.State {
	case JobScheduled:
		j.State = JobAborted
		if j.DoneCallback != nil {
			go j.DoneCallback(j.ViewWithLock())
		}
	case JobStarted:
		j.cancelfn()
	case JobFinished:
		// Job has already finished. Too late.
	case JobAborted:
		// Job is already aborted. Nothing to do.
	}
	j.mu.Unlock()
}

func (s *Scheduler) ForceGC() { s.forceGCC <- struct{}{} }

const gcPeriod = 10 * time.Second

func (s *Scheduler) schedulerMain() {
	var tickC <-chan time.Time
	tickC = make(chan time.Time)
	var nextWakeUp time.Time
	var timer *time.Timer
	renewTimer := func() {
		d := nextWakeUp.Sub(time.Now())
		if timer == nil {
			timer = time.NewTimer(d)
		} else {
			timer.Stop()
			timer.Reset(d)
		}
		tickC = timer.C
	}

	waitJobs := make([]*job, 0)
	jobs := make(map[ID]*job)

	gcTickC := time.Tick(gcPeriod)
	doGC := func() {
		logger.Debugf(mylog, "JobGC start.")
		gcstart := time.Now()
		threshold := gcstart.Add(-s.ZombiePeriod)

		oldJobIDs := make([]ID, 0)
		for id, j := range jobs {
			j.mu.Lock()
			if j.State == JobAborted || j.State == JobFinished {
				if j.FinishedAt.Before(threshold) {
					logger.Debugf(mylog, "GC-ing job %v: %v since aborted/finished.", j, gcstart.Sub(j.FinishedAt))
					oldJobIDs = append(oldJobIDs, id)
				}
			}
			j.mu.Unlock()
		}
		for _, id := range oldJobIDs {
			delete(jobs, id)
		}

		logger.Debugf(mylog, "JobGC end. Took %v. Deleted %d entries", time.Since(gcstart), len(oldJobIDs))
	}

	defer func() {
		close(s.runC)
		s.joinScheduleC <- struct{}{}
	}()

	scheduleC := s.scheduleC
	for {
		if scheduleC == nil && s.numWaitJobs == 0 {
			return
		}

		select {
		case j, more := <-scheduleC:
			if !more {
				// stop polling on scheduleC
				scheduleC = nil

				// kick scheduler
				nextWakeUp = time.Now()
				renewTimer()
				continue
			}

			if _, ok := jobs[j.ID]; ok {
				logger.Criticalf(mylog, "job ID %v is already taken. received duplicate: %v", j.ID, j)
				if j.scheduledC != nil {
					close(j.scheduledC)
				}
				continue
			}
			jobs[j.ID] = j

			now := time.Now()
			if j.ScheduledAt.Before(now) {
				s.runC <- j
			} else {
				if nextWakeUp.IsZero() || nextWakeUp.After(j.ScheduledAt) {
					nextWakeUp = j.ScheduledAt
					renewTimer()
				}

				waitJobs = append(waitJobs, j)
				s.numWaitJobs = len(waitJobs)
			}
			if j.scheduledC != nil {
				close(j.scheduledC)
			}

		case q := <-s.queryC:
			if q == nil {
				continue
			}
			if q.ID == allJobs {
				jvs := make([]*JobView, 0, len(jobs))
				for _, j := range jobs {
					jvs = append(jvs, j.View())
				}
				q.resultC <- jvs
				continue
			}

			id := q.ID
			j, ok := jobs[id]
			if !ok {
				q.resultC <- nil
			} else {
				q.resultC <- []*JobView{j.View()}
			}

		case req := <-s.abortC:
			if req == nil {
				continue
			}

			if req.ID == allJobs {
				for _, j := range jobs {
					abortJob(j)
				}

				req.doneC <- struct{}{}
				continue
			}

			j, ok := jobs[req.ID]
			if !ok {
				logger.Warningf(mylog, "Abort target job ID %d doesn't exist.", req.ID)
				req.doneC <- struct{}{}
				continue
			}
			abortJob(j)
			req.doneC <- struct{}{}

		case <-tickC:
			stillWaitJobs := make([]*job, 0, len(waitJobs))
			now := time.Now()
			var zero time.Time
			nextWakeUp = zero
			for _, j := range waitJobs {
				if j.State != JobScheduled {
					continue
				}
				if j.ScheduledAt.Before(now) {
					s.runC <- j
				} else {
					if nextWakeUp.IsZero() || nextWakeUp.After(j.ScheduledAt) {
						nextWakeUp = j.ScheduledAt
					}

					stillWaitJobs = append(stillWaitJobs, j)
				}
			}
			waitJobs = stillWaitJobs
			s.numWaitJobs = len(waitJobs)

			renewTimer()

		case <-gcTickC:
			doGC()
		case <-s.forceGCC:
			doGC()
		}
	}
}

func (s *Scheduler) runnerMain() {
	for j := range s.runC {
		j.mu.Lock()
		if j.State == JobAborted {
			j.mu.Unlock()
			continue
		}
		if j.State != JobScheduled {
			logger.Infof(mylog, "Skipping job not in scheduled state: %v", j)

			j.mu.Unlock()
			continue
		}

		task := j.Task
		ctx, cancelfn := context.WithCancel(context.Background())
		j.cancelfn = cancelfn
		j.StartedAt = time.Now()
		j.State = JobStarted

		j.mu.Unlock()
		logger.Debugf(mylog, "About to run job %v", j)
		result := task.Run(ctx)
		if result != nil {
			if err := result.Err(); err != nil {
				logger.Warningf(mylog, "Job %v failed with error: %v", j, err)
			}
		}
		finishedAt := time.Now()
		j.mu.Lock()

		j.Result = result
		j.FinishedAt = finishedAt
		if ctx.Err() != nil {
			j.State = JobAborted
		} else {
			j.State = JobFinished
		}
		if j.DoneCallback != nil {
			go j.DoneCallback(j.ViewWithLock())
		}
		j.mu.Unlock()
	}
	s.joinRunnerC <- struct{}{}
}
<|endoftext|>


<commit_before>//  Copyright (c) 2017 Couchbase, Inc.
//
// Licensed 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 scorch

import (
	"fmt"
	"sync/atomic"

	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/index/scorch/segment"
)

type segmentIntroduction struct {
	id        uint64
	data      segment.Segment
	obsoletes map[uint64]*roaring.Bitmap
	ids       []string
	internal  map[string][]byte

	applied   chan error
	persisted chan error
}

type persistIntroduction struct {
	persisted map[uint64]segment.Segment
	applied   notificationChan
}

type epochWatcher struct {
	epoch    uint64
	notifyCh notificationChan
}

type snapshotReversion struct {
	snapshot  *IndexSnapshot
	applied   chan error
	persisted chan error
}

func (s *Scorch) mainLoop() {
	var epochWatchers []*epochWatcher
OUTER:
	for {
		atomic.AddUint64(&s.stats.TotIntroduceLoop, 1)

		select {
		case <-s.closeCh:
			break OUTER

		case epochWatcher := <-s.introducerNotifier:
			epochWatchers = append(epochWatchers, epochWatcher)

		case nextMerge := <-s.merges:
			s.introduceMerge(nextMerge)

		case next := <-s.introductions:
			err := s.introduceSegment(next)
			if err != nil {
				continue OUTER
			}

		case persist := <-s.persists:
			s.introducePersist(persist)

		case revertTo := <-s.revertToSnapshots:
			err := s.revertToSnapshot(revertTo)
			if err != nil {
				continue OUTER
			}
		}

		var epochCurr uint64
		s.rootLock.RLock()
		if s.root != nil {
			epochCurr = s.root.epoch
		}
		s.rootLock.RUnlock()
		var epochWatchersNext []*epochWatcher
		for _, w := range epochWatchers {
			if w.epoch < epochCurr {
				close(w.notifyCh)
			} else {
				epochWatchersNext = append(epochWatchersNext, w)
			}
		}
		epochWatchers = epochWatchersNext
	}

	s.asyncTasks.Done()
}

func (s *Scorch) introduceSegment(next *segmentIntroduction) error {
	atomic.AddUint64(&s.stats.TotIntroduceSegmentBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroduceSegmentEnd, 1)

	s.rootLock.RLock()
	root := s.root
	s.rootLock.RUnlock()

	nsegs := len(root.segment)

	// prepare new index snapshot
	newSnapshot := &IndexSnapshot{
		parent:   s,
		segment:  make([]*SegmentSnapshot, 0, nsegs+1),
		offsets:  make([]uint64, 0, nsegs+1),
		internal: make(map[string][]byte, len(root.internal)),
		refs:     1,
	}

	// iterate through current segments
	var running uint64
	for i := range root.segment {
		// see if optimistic work included this segment
		delta, ok := next.obsoletes[root.segment[i].id]
		if !ok {
			var err error
			delta, err = root.segment[i].segment.DocNumbers(next.ids)
			if err != nil {
				next.applied <- fmt.Errorf("error computing doc numbers: %v", err)
				close(next.applied)
				_ = newSnapshot.DecRef()
				return err
			}
		}

		newss := &SegmentSnapshot{
			id:         root.segment[i].id,
			segment:    root.segment[i].segment,
			cachedDocs: root.segment[i].cachedDocs,
		}

		// apply new obsoletions
		if root.segment[i].deleted == nil {
			newss.deleted = delta
		} else {
			newss.deleted = roaring.Or(root.segment[i].deleted, delta)
		}

		// check for live size before copying
		if newss.LiveSize() > 0 {
			newSnapshot.segment = append(newSnapshot.segment, newss)
			root.segment[i].segment.AddRef()
			newSnapshot.offsets = append(newSnapshot.offsets, running)
			running += newss.segment.Count()
		}
	}

	// append new segment, if any, to end of the new index snapshot
	if next.data != nil {
		newSegmentSnapshot := &SegmentSnapshot{
			id:         next.id,
			segment:    next.data, // take ownership of next.data's ref-count
			cachedDocs: &cachedDocs{cache: nil},
		}
		newSnapshot.segment = append(newSnapshot.segment, newSegmentSnapshot)
		newSnapshot.offsets = append(newSnapshot.offsets, running)

		// increment numItemsIntroduced which tracks the number of items
		// queued for persistence.
		atomic.AddUint64(&s.stats.TotIntroducedItems, newSegmentSnapshot.Count())
		atomic.AddUint64(&s.stats.TotIntroducedSegmentsBatch, 1)
	}
	// copy old values
	for key, oldVal := range root.internal {
		newSnapshot.internal[key] = oldVal
	}
	// set new values and apply deletes
	for key, newVal := range next.internal {
		if newVal != nil {
			newSnapshot.internal[key] = newVal
		} else {
			delete(newSnapshot.internal, key)
		}
	}

	newSnapshot.updateSize()
	s.rootLock.Lock()
	if next.persisted != nil {
		s.rootPersisted = append(s.rootPersisted, next.persisted)
	}
	// swap in new index snapshot
	newSnapshot.epoch = s.nextSnapshotEpoch
	s.nextSnapshotEpoch++
	rootPrev := s.root
	s.root = newSnapshot
	// release lock
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	close(next.applied)

	return nil
}

func (s *Scorch) introducePersist(persist *persistIntroduction) {
	atomic.AddUint64(&s.stats.TotIntroducePersistBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroducePersistEnd, 1)

	s.rootLock.Lock()
	root := s.root
	nextSnapshotEpoch := s.nextSnapshotEpoch
	s.nextSnapshotEpoch++
	s.rootLock.Unlock()

	newIndexSnapshot := &IndexSnapshot{
		parent:   s,
		epoch:    nextSnapshotEpoch,
		segment:  make([]*SegmentSnapshot, len(root.segment)),
		offsets:  make([]uint64, len(root.offsets)),
		internal: make(map[string][]byte, len(root.internal)),
		refs:     1,
	}

	for i, segmentSnapshot := range root.segment {
		// see if this segment has been replaced
		if replacement, ok := persist.persisted[segmentSnapshot.id]; ok {
			newSegmentSnapshot := &SegmentSnapshot{
				id:         segmentSnapshot.id,
				segment:    replacement,
				deleted:    segmentSnapshot.deleted,
				cachedDocs: segmentSnapshot.cachedDocs,
			}
			newIndexSnapshot.segment[i] = newSegmentSnapshot
			delete(persist.persisted, segmentSnapshot.id)

			// update items persisted incase of a new segment snapshot
			atomic.AddUint64(&s.stats.TotPersistedItems, newSegmentSnapshot.Count())
			atomic.AddUint64(&s.stats.TotPersistedSegments, 1)
		} else {
			newIndexSnapshot.segment[i] = root.segment[i]
			newIndexSnapshot.segment[i].segment.AddRef()
		}
		newIndexSnapshot.offsets[i] = root.offsets[i]
	}

	for k, v := range root.internal {
		newIndexSnapshot.internal[k] = v
	}

	newIndexSnapshot.updateSize()
	s.rootLock.Lock()
	rootPrev := s.root
	s.root = newIndexSnapshot
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	close(persist.applied)
}

func (s *Scorch) introduceMerge(nextMerge *segmentMerge) {
	atomic.AddUint64(&s.stats.TotIntroduceMergeBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroduceMergeEnd, 1)

	s.rootLock.RLock()
	root := s.root
	s.rootLock.RUnlock()

	newSnapshot := &IndexSnapshot{
		parent:   s,
		internal: root.internal,
		refs:     1,
	}

	// iterate through current segments
	newSegmentDeleted := roaring.NewBitmap()
	var running uint64
	for i := range root.segment {
		segmentID := root.segment[i].id
		if segSnapAtMerge, ok := nextMerge.old[segmentID]; ok {
			// this segment is going away, see if anything else was deleted since we started the merge
			if segSnapAtMerge != nil && root.segment[i].deleted != nil {
				// assume all these deletes are new
				deletedSince := root.segment[i].deleted
				// if we already knew about some of them, remove
				if segSnapAtMerge.deleted != nil {
					deletedSince = roaring.AndNot(root.segment[i].deleted, segSnapAtMerge.deleted)
				}
				deletedSinceItr := deletedSince.Iterator()
				for deletedSinceItr.HasNext() {
					oldDocNum := deletedSinceItr.Next()
					newDocNum := nextMerge.oldNewDocNums[segmentID][oldDocNum]
					newSegmentDeleted.Add(uint32(newDocNum))
				}
			}
			// clean up the old segment map to figure out the
			// obsolete segments wrt root in meantime, whatever
			// segments left behind in old map after processing
			// the root segments would be the obsolete segment set
			delete(nextMerge.old, segmentID)
		} else if root.segment[i].LiveSize() > 0 {
			// this segment is staying
			newSnapshot.segment = append(newSnapshot.segment, &SegmentSnapshot{
				id:         root.segment[i].id,
				segment:    root.segment[i].segment,
				deleted:    root.segment[i].deleted,
				cachedDocs: root.segment[i].cachedDocs,
			})
			root.segment[i].segment.AddRef()
			newSnapshot.offsets = append(newSnapshot.offsets, running)
			running += root.segment[i].segment.Count()
		}
	}

	// before the newMerge introduction, need to clean the newly
	// merged segment wrt the current root segments, hence
	// applying the obsolete segment contents to newly merged segment
	for segID, ss := range nextMerge.old {
		obsoleted := ss.DocNumbersLive()
		if obsoleted != nil {
			obsoletedIter := obsoleted.Iterator()
			for obsoletedIter.HasNext() {
				oldDocNum := obsoletedIter.Next()
				newDocNum := nextMerge.oldNewDocNums[segID][oldDocNum]
				newSegmentDeleted.Add(uint32(newDocNum))
			}
		}
	}
	// In case where all the docs in the newly merged segment getting
	// deleted by the time we reach here, can skip the introduction.
	if nextMerge.new != nil &&
		nextMerge.new.Count() > newSegmentDeleted.GetCardinality() {
		// put new segment at end
		newSnapshot.segment = append(newSnapshot.segment, &SegmentSnapshot{
			id:         nextMerge.id,
			segment:    nextMerge.new, // take ownership for nextMerge.new's ref-count
			deleted:    newSegmentDeleted,
			cachedDocs: &cachedDocs{cache: nil},
		})
		newSnapshot.offsets = append(newSnapshot.offsets, running)
		atomic.AddUint64(&s.stats.TotIntroducedSegmentsMerge, 1)
	}

	newSnapshot.AddRef() // 1 ref for the nextMerge.notify response

	newSnapshot.updateSize()
	s.rootLock.Lock()
	// swap in new index snapshot
	newSnapshot.epoch = s.nextSnapshotEpoch
	s.nextSnapshotEpoch++
	rootPrev := s.root
	s.root = newSnapshot
	// release lock
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	// notify requester that we incorporated this
	nextMerge.notify <- newSnapshot
	close(nextMerge.notify)
}

func (s *Scorch) revertToSnapshot(revertTo *snapshotReversion) error {
	atomic.AddUint64(&s.stats.TotIntroduceRevertBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroduceRevertEnd, 1)

	if revertTo.snapshot == nil {
		err := fmt.Errorf("Cannot revert to a nil snapshot")
		revertTo.applied <- err
		return err
	}

	// acquire lock
	s.rootLock.Lock()

	// prepare a new index snapshot, based on next snapshot
	newSnapshot := &IndexSnapshot{
		parent:   s,
		segment:  make([]*SegmentSnapshot, len(revertTo.snapshot.segment)),
		offsets:  revertTo.snapshot.offsets,
		internal: revertTo.snapshot.internal,
		epoch:    s.nextSnapshotEpoch,
		refs:     1,
	}
	s.nextSnapshotEpoch++

	// iterate through segments
	for i, segmentSnapshot := range revertTo.snapshot.segment {
		newSnapshot.segment[i] = &SegmentSnapshot{
			id:         segmentSnapshot.id,
			segment:    segmentSnapshot.segment,
			deleted:    segmentSnapshot.deleted,
			cachedDocs: segmentSnapshot.cachedDocs,
		}
		newSnapshot.segment[i].segment.AddRef()

		// remove segment from ineligibleForRemoval map
		filename := zapFileName(segmentSnapshot.id)
		delete(s.ineligibleForRemoval, filename)
	}

	if revertTo.persisted != nil {
		s.rootPersisted = append(s.rootPersisted, revertTo.persisted)
	}

	newSnapshot.updateSize()
	// swap in new snapshot
	rootPrev := s.root
	s.root = newSnapshot
	// release lock
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	close(revertTo.applied)

	return nil
}
<commit_msg>scorch introduce empty deleted bitmaps as nil<commit_after>//  Copyright (c) 2017 Couchbase, Inc.
//
// Licensed 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 scorch

import (
	"fmt"
	"sync/atomic"

	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/index/scorch/segment"
)

type segmentIntroduction struct {
	id        uint64
	data      segment.Segment
	obsoletes map[uint64]*roaring.Bitmap
	ids       []string
	internal  map[string][]byte

	applied   chan error
	persisted chan error
}

type persistIntroduction struct {
	persisted map[uint64]segment.Segment
	applied   notificationChan
}

type epochWatcher struct {
	epoch    uint64
	notifyCh notificationChan
}

type snapshotReversion struct {
	snapshot  *IndexSnapshot
	applied   chan error
	persisted chan error
}

func (s *Scorch) mainLoop() {
	var epochWatchers []*epochWatcher
OUTER:
	for {
		atomic.AddUint64(&s.stats.TotIntroduceLoop, 1)

		select {
		case <-s.closeCh:
			break OUTER

		case epochWatcher := <-s.introducerNotifier:
			epochWatchers = append(epochWatchers, epochWatcher)

		case nextMerge := <-s.merges:
			s.introduceMerge(nextMerge)

		case next := <-s.introductions:
			err := s.introduceSegment(next)
			if err != nil {
				continue OUTER
			}

		case persist := <-s.persists:
			s.introducePersist(persist)

		case revertTo := <-s.revertToSnapshots:
			err := s.revertToSnapshot(revertTo)
			if err != nil {
				continue OUTER
			}
		}

		var epochCurr uint64
		s.rootLock.RLock()
		if s.root != nil {
			epochCurr = s.root.epoch
		}
		s.rootLock.RUnlock()
		var epochWatchersNext []*epochWatcher
		for _, w := range epochWatchers {
			if w.epoch < epochCurr {
				close(w.notifyCh)
			} else {
				epochWatchersNext = append(epochWatchersNext, w)
			}
		}
		epochWatchers = epochWatchersNext
	}

	s.asyncTasks.Done()
}

func (s *Scorch) introduceSegment(next *segmentIntroduction) error {
	atomic.AddUint64(&s.stats.TotIntroduceSegmentBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroduceSegmentEnd, 1)

	s.rootLock.RLock()
	root := s.root
	s.rootLock.RUnlock()

	nsegs := len(root.segment)

	// prepare new index snapshot
	newSnapshot := &IndexSnapshot{
		parent:   s,
		segment:  make([]*SegmentSnapshot, 0, nsegs+1),
		offsets:  make([]uint64, 0, nsegs+1),
		internal: make(map[string][]byte, len(root.internal)),
		refs:     1,
	}

	// iterate through current segments
	var running uint64
	for i := range root.segment {
		// see if optimistic work included this segment
		delta, ok := next.obsoletes[root.segment[i].id]
		if !ok {
			var err error
			delta, err = root.segment[i].segment.DocNumbers(next.ids)
			if err != nil {
				next.applied <- fmt.Errorf("error computing doc numbers: %v", err)
				close(next.applied)
				_ = newSnapshot.DecRef()
				return err
			}
		}

		newss := &SegmentSnapshot{
			id:         root.segment[i].id,
			segment:    root.segment[i].segment,
			cachedDocs: root.segment[i].cachedDocs,
		}

		// apply new obsoletions
		if root.segment[i].deleted == nil {
			newss.deleted = delta
		} else {
			newss.deleted = roaring.Or(root.segment[i].deleted, delta)
		}
		if newss.deleted.IsEmpty() {
			newss.deleted = nil
		}

		// check for live size before copying
		if newss.LiveSize() > 0 {
			newSnapshot.segment = append(newSnapshot.segment, newss)
			root.segment[i].segment.AddRef()
			newSnapshot.offsets = append(newSnapshot.offsets, running)
			running += newss.segment.Count()
		}
	}

	// append new segment, if any, to end of the new index snapshot
	if next.data != nil {
		newSegmentSnapshot := &SegmentSnapshot{
			id:         next.id,
			segment:    next.data, // take ownership of next.data's ref-count
			cachedDocs: &cachedDocs{cache: nil},
		}
		newSnapshot.segment = append(newSnapshot.segment, newSegmentSnapshot)
		newSnapshot.offsets = append(newSnapshot.offsets, running)

		// increment numItemsIntroduced which tracks the number of items
		// queued for persistence.
		atomic.AddUint64(&s.stats.TotIntroducedItems, newSegmentSnapshot.Count())
		atomic.AddUint64(&s.stats.TotIntroducedSegmentsBatch, 1)
	}
	// copy old values
	for key, oldVal := range root.internal {
		newSnapshot.internal[key] = oldVal
	}
	// set new values and apply deletes
	for key, newVal := range next.internal {
		if newVal != nil {
			newSnapshot.internal[key] = newVal
		} else {
			delete(newSnapshot.internal, key)
		}
	}

	newSnapshot.updateSize()
	s.rootLock.Lock()
	if next.persisted != nil {
		s.rootPersisted = append(s.rootPersisted, next.persisted)
	}
	// swap in new index snapshot
	newSnapshot.epoch = s.nextSnapshotEpoch
	s.nextSnapshotEpoch++
	rootPrev := s.root
	s.root = newSnapshot
	// release lock
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	close(next.applied)

	return nil
}

func (s *Scorch) introducePersist(persist *persistIntroduction) {
	atomic.AddUint64(&s.stats.TotIntroducePersistBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroducePersistEnd, 1)

	s.rootLock.Lock()
	root := s.root
	nextSnapshotEpoch := s.nextSnapshotEpoch
	s.nextSnapshotEpoch++
	s.rootLock.Unlock()

	newIndexSnapshot := &IndexSnapshot{
		parent:   s,
		epoch:    nextSnapshotEpoch,
		segment:  make([]*SegmentSnapshot, len(root.segment)),
		offsets:  make([]uint64, len(root.offsets)),
		internal: make(map[string][]byte, len(root.internal)),
		refs:     1,
	}

	for i, segmentSnapshot := range root.segment {
		// see if this segment has been replaced
		if replacement, ok := persist.persisted[segmentSnapshot.id]; ok {
			newSegmentSnapshot := &SegmentSnapshot{
				id:         segmentSnapshot.id,
				segment:    replacement,
				deleted:    segmentSnapshot.deleted,
				cachedDocs: segmentSnapshot.cachedDocs,
			}
			newIndexSnapshot.segment[i] = newSegmentSnapshot
			delete(persist.persisted, segmentSnapshot.id)

			// update items persisted incase of a new segment snapshot
			atomic.AddUint64(&s.stats.TotPersistedItems, newSegmentSnapshot.Count())
			atomic.AddUint64(&s.stats.TotPersistedSegments, 1)
		} else {
			newIndexSnapshot.segment[i] = root.segment[i]
			newIndexSnapshot.segment[i].segment.AddRef()
		}
		newIndexSnapshot.offsets[i] = root.offsets[i]
	}

	for k, v := range root.internal {
		newIndexSnapshot.internal[k] = v
	}

	newIndexSnapshot.updateSize()
	s.rootLock.Lock()
	rootPrev := s.root
	s.root = newIndexSnapshot
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	close(persist.applied)
}

func (s *Scorch) introduceMerge(nextMerge *segmentMerge) {
	atomic.AddUint64(&s.stats.TotIntroduceMergeBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroduceMergeEnd, 1)

	s.rootLock.RLock()
	root := s.root
	s.rootLock.RUnlock()

	newSnapshot := &IndexSnapshot{
		parent:   s,
		internal: root.internal,
		refs:     1,
	}

	// iterate through current segments
	newSegmentDeleted := roaring.NewBitmap()
	var running uint64
	for i := range root.segment {
		segmentID := root.segment[i].id
		if segSnapAtMerge, ok := nextMerge.old[segmentID]; ok {
			// this segment is going away, see if anything else was deleted since we started the merge
			if segSnapAtMerge != nil && root.segment[i].deleted != nil {
				// assume all these deletes are new
				deletedSince := root.segment[i].deleted
				// if we already knew about some of them, remove
				if segSnapAtMerge.deleted != nil {
					deletedSince = roaring.AndNot(root.segment[i].deleted, segSnapAtMerge.deleted)
				}
				deletedSinceItr := deletedSince.Iterator()
				for deletedSinceItr.HasNext() {
					oldDocNum := deletedSinceItr.Next()
					newDocNum := nextMerge.oldNewDocNums[segmentID][oldDocNum]
					newSegmentDeleted.Add(uint32(newDocNum))
				}
			}
			// clean up the old segment map to figure out the
			// obsolete segments wrt root in meantime, whatever
			// segments left behind in old map after processing
			// the root segments would be the obsolete segment set
			delete(nextMerge.old, segmentID)
		} else if root.segment[i].LiveSize() > 0 {
			// this segment is staying
			newSnapshot.segment = append(newSnapshot.segment, &SegmentSnapshot{
				id:         root.segment[i].id,
				segment:    root.segment[i].segment,
				deleted:    root.segment[i].deleted,
				cachedDocs: root.segment[i].cachedDocs,
			})
			root.segment[i].segment.AddRef()
			newSnapshot.offsets = append(newSnapshot.offsets, running)
			running += root.segment[i].segment.Count()
		}
	}

	// before the newMerge introduction, need to clean the newly
	// merged segment wrt the current root segments, hence
	// applying the obsolete segment contents to newly merged segment
	for segID, ss := range nextMerge.old {
		obsoleted := ss.DocNumbersLive()
		if obsoleted != nil {
			obsoletedIter := obsoleted.Iterator()
			for obsoletedIter.HasNext() {
				oldDocNum := obsoletedIter.Next()
				newDocNum := nextMerge.oldNewDocNums[segID][oldDocNum]
				newSegmentDeleted.Add(uint32(newDocNum))
			}
		}
	}
	// In case where all the docs in the newly merged segment getting
	// deleted by the time we reach here, can skip the introduction.
	if nextMerge.new != nil &&
		nextMerge.new.Count() > newSegmentDeleted.GetCardinality() {
		// put new segment at end
		newSnapshot.segment = append(newSnapshot.segment, &SegmentSnapshot{
			id:         nextMerge.id,
			segment:    nextMerge.new, // take ownership for nextMerge.new's ref-count
			deleted:    newSegmentDeleted,
			cachedDocs: &cachedDocs{cache: nil},
		})
		newSnapshot.offsets = append(newSnapshot.offsets, running)
		atomic.AddUint64(&s.stats.TotIntroducedSegmentsMerge, 1)
	}

	newSnapshot.AddRef() // 1 ref for the nextMerge.notify response

	newSnapshot.updateSize()
	s.rootLock.Lock()
	// swap in new index snapshot
	newSnapshot.epoch = s.nextSnapshotEpoch
	s.nextSnapshotEpoch++
	rootPrev := s.root
	s.root = newSnapshot
	// release lock
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	// notify requester that we incorporated this
	nextMerge.notify <- newSnapshot
	close(nextMerge.notify)
}

func (s *Scorch) revertToSnapshot(revertTo *snapshotReversion) error {
	atomic.AddUint64(&s.stats.TotIntroduceRevertBeg, 1)
	defer atomic.AddUint64(&s.stats.TotIntroduceRevertEnd, 1)

	if revertTo.snapshot == nil {
		err := fmt.Errorf("Cannot revert to a nil snapshot")
		revertTo.applied <- err
		return err
	}

	// acquire lock
	s.rootLock.Lock()

	// prepare a new index snapshot, based on next snapshot
	newSnapshot := &IndexSnapshot{
		parent:   s,
		segment:  make([]*SegmentSnapshot, len(revertTo.snapshot.segment)),
		offsets:  revertTo.snapshot.offsets,
		internal: revertTo.snapshot.internal,
		epoch:    s.nextSnapshotEpoch,
		refs:     1,
	}
	s.nextSnapshotEpoch++

	// iterate through segments
	for i, segmentSnapshot := range revertTo.snapshot.segment {
		newSnapshot.segment[i] = &SegmentSnapshot{
			id:         segmentSnapshot.id,
			segment:    segmentSnapshot.segment,
			deleted:    segmentSnapshot.deleted,
			cachedDocs: segmentSnapshot.cachedDocs,
		}
		newSnapshot.segment[i].segment.AddRef()

		// remove segment from ineligibleForRemoval map
		filename := zapFileName(segmentSnapshot.id)
		delete(s.ineligibleForRemoval, filename)
	}

	if revertTo.persisted != nil {
		s.rootPersisted = append(s.rootPersisted, revertTo.persisted)
	}

	newSnapshot.updateSize()
	// swap in new snapshot
	rootPrev := s.root
	s.root = newSnapshot
	// release lock
	s.rootLock.Unlock()

	if rootPrev != nil {
		_ = rootPrev.DecRef()
	}

	close(revertTo.applied)

	return nil
}
<|endoftext|>


<commit_before>package indicators_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/thetruetrade/gotrade/indicators"
)

var _ = Describe("when calculating a minus directional indicator (MinusDi) with DOHLCV source data", func() {
	var (
		period    int = 14
		indicator *MinusDi
		inputs    IndicatorWithFloatBoundsSharedSpecInputs
	)

	BeforeEach(func() {
		indicator, _ = NewMinusDi(period)

		inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
			func() float64 {
				return GetFloatDataMax(indicator.Data)
			},
			func() float64 {
				return GetFloatDataMin(indicator.Data)
			})
	})

	Context("and the indicator has not yet received any ticks", func() {
		ShouldBeAnInitialisedIndicator(&inputs)

		ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
	})

	Context("and the indicator has received less ticks than the lookback period", func() {

		BeforeEach(func() {
			for i := 0; i < indicator.GetLookbackPeriod(); i++ {
				indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
			}
		})

		ShouldBeAnIndicatorThatHasReceivedFewerTicksThanItsLookbackPeriod(&inputs)

		ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
	})

	Context("and the indicator has received ticks equal to the lookback period", func() {

		BeforeEach(func() {
			for i := 0; i <= indicator.GetLookbackPeriod(); i++ {
				indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
			}
		})

		ShouldBeAnIndicatorThatHasReceivedTicksEqualToItsLookbackPeriod(&inputs)

		ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
	})

	Context("and the indicator has received more ticks than the lookback period", func() {

		BeforeEach(func() {
			for i := range sourceDOHLCVData {
				indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
			}
		})

		ShouldBeAnIndicatorThatHasReceivedMoreTicksThanItsLookbackPeriod(&inputs)

		ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
	})

	Context("and the indicator has recieved all of its ticks", func() {
		BeforeEach(func() {
			for i := 0; i < len(sourceDOHLCVData); i++ {
				indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
			}
		})

		ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

		ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
	})
})
<commit_msg>#77 achieve 100% test coverage for indicators - minusdi<commit_after>package indicators_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/thetruetrade/gotrade/indicators"
)

var _ = Describe("when creating an minusdiwithoutstorage", func() {
	var (
		indicator      *indicators.MinusDiWithoutStorage
		indicatorError error
	)

	Context("and the indicator was not given a value available action", func() {
		BeforeEach(func() {
			indicator, indicatorError = indicators.NewMinusDiWithoutStorage(4, nil)
		})

		It("the indicator should not be created and return the appropriate error message", func() {
			Expect(indicator).To(BeNil())
			Expect(indicatorError).To(Equal(indicators.ErrValueAvailableActionIsNil))
		})
	})

	Context("and the indicator was given a timePeriod below the minimum", func() {
		BeforeEach(func() {
			indicator, indicatorError = indicators.NewMinusDiWithoutStorage(0, fakeFloatValAvailable)
		})

		It("the indicator should not be created and return the appropriate error message", func() {
			Expect(indicator).To(BeNil())
		})
	})

	Context("and the indicator was given a timePeriod above the maximum", func() {
		BeforeEach(func() {
			indicator, indicatorError = indicators.NewMinusDiWithoutStorage(indicators.MaximumLookbackPeriod+1, fakeFloatValAvailable)
		})

		It("the indicator should not be created and return the appropriate error message", func() {
			Expect(indicator).To(BeNil())
		})
	})
})

var _ = Describe("when calculating a minus directional indicator (MinusDi) with DOHLCV source data", func() {
	var (
		period    int = 14
		indicator *indicators.MinusDi
		inputs    IndicatorWithFloatBoundsSharedSpecInputs
		stream    *fakeDOHLCVStreamSubscriber
	)

	Context("given the indicator is created via the standard constructor", func() {
		BeforeEach(func() {
			indicator, _ = indicators.NewMinusDi(period)

			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has received less ticks than the lookback period", func() {

			BeforeEach(func() {
				for i := 0; i < indicator.GetLookbackPeriod(); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedFewerTicksThanItsLookbackPeriod(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has received ticks equal to the lookback period", func() {

			BeforeEach(func() {
				for i := 0; i <= indicator.GetLookbackPeriod(); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedTicksEqualToItsLookbackPeriod(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
		})

		Context("and the indicator has received more ticks than the lookback period", func() {

			BeforeEach(func() {
				for i := range sourceDOHLCVData {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedMoreTicksThanItsLookbackPeriod(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
		})
	})

	Context("given the indicator is created via the constructor with defaulted parameters", func() {
		BeforeEach(func() {
			indicator, _ = indicators.NewDefaultMinusDi()
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
		})
	})

	Context("given the indicator is created via the constructor with fixed source length", func() {
		BeforeEach(func() {
			indicator, _ = indicators.NewMinusDiWithSrcLen(uint(len(sourceDOHLCVData)), 4)
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		It("should have pre-allocated storge for the output data", func() {
			Expect(cap(indicator.Data)).To(Equal(len(sourceDOHLCVData) - indicator.GetLookbackPeriod()))
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)

			It("no new storage capcity should have been allocated", func() {
				Expect(len(indicator.Data)).To(Equal(cap(indicator.Data)))
			})
		})
	})

	Context("given the indicator is created via the constructor with defaulted parameters and fixed source length", func() {
		BeforeEach(func() {
			indicator, _ = indicators.NewDefaultMinusDiWithSrcLen(uint(len(sourceDOHLCVData)))
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		It("should have pre-allocated storge for the output data", func() {
			Expect(cap(indicator.Data)).To(Equal(len(sourceDOHLCVData) - indicator.GetLookbackPeriod()))
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)

			It("no new storage capcity should have been allocated", func() {
				Expect(len(indicator.Data)).To(Equal(cap(indicator.Data)))
			})
		})
	})

	Context("given the indicator is created via the constructor for use with a price stream", func() {
		BeforeEach(func() {
			stream = newFakeDOHLCVStreamSubscriber()
			indicator, _ = indicators.NewMinusDiForStream(stream, 4)
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		It("should have requested to be attached to the stream", func() {
			Expect(stream.lastCallToAddTickSubscriptionArg).To(Equal(indicator))
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
		})
	})

	Context("given the indicator is created via the constructor for use with a price stream with defaulted parameters", func() {
		BeforeEach(func() {
			stream = newFakeDOHLCVStreamSubscriber()
			indicator, _ = indicators.NewDefaultMinusDiForStream(stream)
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		It("should have requested to be attached to the stream", func() {
			Expect(stream.lastCallToAddTickSubscriptionArg).To(Equal(indicator))
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)
		})
	})

	Context("given the indicator is created via the constructor for use with a price stream with fixed source length", func() {
		BeforeEach(func() {
			stream = newFakeDOHLCVStreamSubscriber()
			indicator, _ = indicators.NewMinusDiForStreamWithSrcLen(uint(len(sourceDOHLCVData)), stream, 4)
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		It("should have pre-allocated storge for the output data", func() {
			Expect(cap(indicator.Data)).To(Equal(len(sourceDOHLCVData) - indicator.GetLookbackPeriod()))
		})

		It("should have requested to be attached to the stream", func() {
			Expect(stream.lastCallToAddTickSubscriptionArg).To(Equal(indicator))
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)

			It("no new storage capcity should have been allocated", func() {
				Expect(len(indicator.Data)).To(Equal(cap(indicator.Data)))
			})
		})
	})

	Context("given the indicator is created via the constructor for use with a price stream with fixed source length with defaulted parmeters", func() {
		BeforeEach(func() {
			stream = newFakeDOHLCVStreamSubscriber()
			indicator, _ = indicators.NewDefaultMinusDiForStreamWithSrcLen(uint(len(sourceDOHLCVData)), stream)
			inputs = NewIndicatorWithFloatBoundsSharedSpecInputs(indicator, len(sourceDOHLCVData), indicator,
				func() float64 {
					return GetFloatDataMax(indicator.Data)
				},
				func() float64 {
					return GetFloatDataMin(indicator.Data)
				})
		})

		It("should have pre-allocated storge for the output data", func() {
			Expect(cap(indicator.Data)).To(Equal(len(sourceDOHLCVData) - indicator.GetLookbackPeriod()))
		})

		It("should have requested to be attached to the stream", func() {
			Expect(stream.lastCallToAddTickSubscriptionArg).To(Equal(indicator))
		})

		Context("and the indicator has not yet received any ticks", func() {
			ShouldBeAnInitialisedIndicator(&inputs)

			ShouldNotHaveAnyFloatBoundsSetYet(&inputs)
		})

		Context("and the indicator has recieved all of its ticks", func() {
			BeforeEach(func() {
				for i := 0; i < len(sourceDOHLCVData); i++ {
					indicator.ReceiveDOHLCVTick(sourceDOHLCVData[i], i+1)
				}
			})

			ShouldBeAnIndicatorThatHasReceivedAllOfItsTicks(&inputs)

			ShouldHaveFloatBoundsSetToMinMaxOfResults(&inputs)

			It("no new storage capcity should have been allocated", func() {
				Expect(len(indicator.Data)).To(Equal(cap(indicator.Data)))
			})
		})
	})
})
<|endoftext|>


<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 avro

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	avro "github.com/elodina/go-avro"
	"sync"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
	auth         *KafkaAvroAuth
	lock         sync.RWMutex
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return NewCachedSchemaRegistryClientAuth(registryURL, nil)
}

func NewCachedSchemaRegistryClientAuth(registryURL string, auth *KafkaAvroAuth) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
		auth:         auth,
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool

	this.lock.RLock()
	if schemaIdMap, exists = this.schemaCache[subject]; exists {
		var id int32
		if id, exists = schemaIdMap[schema]; exists {
			return id, nil
		}
	}
	this.lock.RUnlock()

	this.lock.Lock()
	defer this.lock.Unlock()
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", SCHEMA_REGISTRY_V1_JSON)
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	if this.auth != nil {
		request.Header.Set("X-Api-User", this.auth.User)
		request.Header.Set("X-Api-Key", this.auth.Key)
	}
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<commit_msg>Fix idCache races<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 avro

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	avro "github.com/elodina/go-avro"
	"sync"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
	auth         *KafkaAvroAuth
	lock         sync.RWMutex
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return NewCachedSchemaRegistryClientAuth(registryURL, nil)
}

func NewCachedSchemaRegistryClientAuth(registryURL string, auth *KafkaAvroAuth) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
		auth:         auth,
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool

	this.lock.RLock()
	if schemaIdMap, exists = this.schemaCache[subject]; exists {
		var id int32
		if id, exists = schemaIdMap[schema]; exists {
			return id, nil
		}
	}
	this.lock.RUnlock()

	this.lock.Lock()
	defer this.lock.Unlock()
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.lock.RLock()
		this.idCache[decodedResponse.Id] = schema
		this.lock.RUnlock()

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	this.lock.RLock()
	if schema, exists = this.idCache[id]; exists {
		this.lock.RUnlock()
		return schema, nil
	}
	this.lock.RUnlock()

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.lock.Lock()
		this.idCache[id] = schema
		this.lock.Unlock()

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", SCHEMA_REGISTRY_V1_JSON)
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	if this.auth != nil {
		request.Header.Set("X-Api-User", this.auth.User)
		request.Header.Set("X-Api-Key", this.auth.Key)
	}
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<|endoftext|>


<commit_before>package testflight_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit"
)

var _ = Describe("Flying", func() {
	var tmpdir string
	var fixture string

	BeforeEach(func() {
		var err error

		writeATCPipeline("noop.yml", nil)

		atcProcess = ifrit.Envoke(atcRunner)
		Consistently(atcProcess.Wait(), 1*time.Second).ShouldNot(Receive())

		tmpdir, err = ioutil.TempDir("", "fly-test")
		Ω(err).ShouldNot(HaveOccurred())

		fixture = filepath.Join(tmpdir, "fixture")

		err = os.MkdirAll(fixture, 0755)
		Ω(err).ShouldNot(HaveOccurred())

		err = ioutil.WriteFile(
			filepath.Join(fixture, "run"),
			[]byte(`#!/bin/bash
echo some output
echo FOO is $FOO
echo ARGS are "$@"
exit 0
`),
			0755,
		)
		Ω(err).ShouldNot(HaveOccurred())

		err = ioutil.WriteFile(
			filepath.Join(fixture, "build.yml"),
			[]byte(fmt.Sprintf(`---
image: %s

params:
  FOO: 1

run:
  path: fixture/run
`, helperRootfs)),
			0644,
		)
		Ω(err).ShouldNot(HaveOccurred())
	})

	AfterEach(func() {
		os.RemoveAll(tmpdir)
	})

	It("works", func() {
		fly := exec.Command(builtComponents["fly"], "--", "SOME", "ARGS")
		fly.Dir = fixture

		session, err := gexec.Start(fly, GinkgoWriter, GinkgoWriter)
		Ω(err).ShouldNot(HaveOccurred())

		Eventually(session, 1*time.Minute).Should(gexec.Exit(0))

		Ω(session).Should(gbytes.Say("some output"))
		Ω(session).Should(gbytes.Say("FOO is 1"))
		Ω(session).Should(gbytes.Say("ARGS are SOME ARGS"))
	})

	Describe("hijacking", func() {
		It("executes an interactive command in a running build's container", func() {
			err := ioutil.WriteFile(
				filepath.Join(fixture, "run"),
				[]byte(`#!/bin/bash
mkfifo /tmp/fifo
echo waiting
cat < /tmp/fifo
`),
				0755,
			)
			Ω(err).ShouldNot(HaveOccurred())

			fly := exec.Command(builtComponents["fly"])
			fly.Dir = fixture

			flyS, err := gexec.Start(fly, GinkgoWriter, GinkgoWriter)
			Ω(err).ShouldNot(HaveOccurred())

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("waiting"))

			// TODO there's a gap between start + attach in turbine
			time.Sleep(5 * time.Second)

			hijack := exec.Command(builtComponents["fly"], "hijack", "--", "bash", "-c", "echo marco > /tmp/fifo")

			hijackS, err := gexec.Start(hijack, GinkgoWriter, GinkgoWriter)
			Ω(err).ShouldNot(HaveOccurred())

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("marco"))

			Eventually(hijackS, 5*time.Second).Should(gexec.Exit())

			Eventually(flyS, 5*time.Second).Should(gexec.Exit(0))
		})
	})

	Describe("aborting", func() {
		It("terminates the running build", func() {
			err := ioutil.WriteFile(
				filepath.Join(fixture, "run"),
				[]byte(`#!/bin/bash
trap "echo build got sigterm; exit 1" SIGTERM
sleep 1000 &
echo waiting
wait
`),
				0755,
			)
			Ω(err).ShouldNot(HaveOccurred())

			fly := exec.Command(builtComponents["fly"])
			fly.Dir = fixture

			flyS, err := gexec.Start(fly, GinkgoWriter, GinkgoWriter)
			Ω(err).ShouldNot(HaveOccurred())

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("waiting"))

			flyS.Signal(syscall.SIGTERM)

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("build got sigterm"))

			// build should have errored
			Eventually(flyS, 5*time.Second).Should(gexec.Exit(2))
		})
	})
})
<commit_msg>fix aborted build exit status check<commit_after>package testflight_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit"
)

var _ = Describe("Flying", func() {
	var tmpdir string
	var fixture string

	BeforeEach(func() {
		var err error

		writeATCPipeline("noop.yml", nil)

		atcProcess = ifrit.Envoke(atcRunner)
		Consistently(atcProcess.Wait(), 1*time.Second).ShouldNot(Receive())

		tmpdir, err = ioutil.TempDir("", "fly-test")
		Ω(err).ShouldNot(HaveOccurred())

		fixture = filepath.Join(tmpdir, "fixture")

		err = os.MkdirAll(fixture, 0755)
		Ω(err).ShouldNot(HaveOccurred())

		err = ioutil.WriteFile(
			filepath.Join(fixture, "run"),
			[]byte(`#!/bin/bash
echo some output
echo FOO is $FOO
echo ARGS are "$@"
exit 0
`),
			0755,
		)
		Ω(err).ShouldNot(HaveOccurred())

		err = ioutil.WriteFile(
			filepath.Join(fixture, "build.yml"),
			[]byte(fmt.Sprintf(`---
image: %s

params:
  FOO: 1

run:
  path: fixture/run
`, helperRootfs)),
			0644,
		)
		Ω(err).ShouldNot(HaveOccurred())
	})

	AfterEach(func() {
		os.RemoveAll(tmpdir)
	})

	It("works", func() {
		fly := exec.Command(builtComponents["fly"], "--", "SOME", "ARGS")
		fly.Dir = fixture

		session, err := gexec.Start(fly, GinkgoWriter, GinkgoWriter)
		Ω(err).ShouldNot(HaveOccurred())

		Eventually(session, 1*time.Minute).Should(gexec.Exit(0))

		Ω(session).Should(gbytes.Say("some output"))
		Ω(session).Should(gbytes.Say("FOO is 1"))
		Ω(session).Should(gbytes.Say("ARGS are SOME ARGS"))
	})

	Describe("hijacking", func() {
		It("executes an interactive command in a running build's container", func() {
			err := ioutil.WriteFile(
				filepath.Join(fixture, "run"),
				[]byte(`#!/bin/bash
mkfifo /tmp/fifo
echo waiting
cat < /tmp/fifo
`),
				0755,
			)
			Ω(err).ShouldNot(HaveOccurred())

			fly := exec.Command(builtComponents["fly"])
			fly.Dir = fixture

			flyS, err := gexec.Start(fly, GinkgoWriter, GinkgoWriter)
			Ω(err).ShouldNot(HaveOccurred())

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("waiting"))

			// TODO there's a gap between start + attach in turbine
			time.Sleep(5 * time.Second)

			hijack := exec.Command(builtComponents["fly"], "hijack", "--", "bash", "-c", "echo marco > /tmp/fifo")

			hijackS, err := gexec.Start(hijack, GinkgoWriter, GinkgoWriter)
			Ω(err).ShouldNot(HaveOccurred())

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("marco"))

			Eventually(hijackS, 5*time.Second).Should(gexec.Exit())

			Eventually(flyS, 5*time.Second).Should(gexec.Exit(0))
		})
	})

	Describe("aborting", func() {
		It("terminates the running build", func() {
			err := ioutil.WriteFile(
				filepath.Join(fixture, "run"),
				[]byte(`#!/bin/bash
trap "echo build got sigterm; exit 1" SIGTERM
sleep 1000 &
echo waiting
wait
`),
				0755,
			)
			Ω(err).ShouldNot(HaveOccurred())

			fly := exec.Command(builtComponents["fly"])
			fly.Dir = fixture

			flyS, err := gexec.Start(fly, GinkgoWriter, GinkgoWriter)
			Ω(err).ShouldNot(HaveOccurred())

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("waiting"))

			flyS.Signal(syscall.SIGTERM)

			Eventually(flyS, 10*time.Second).Should(gbytes.Say("build got sigterm"))

			// build should have been aborted
			Eventually(flyS, 5*time.Second).Should(gexec.Exit(3))
		})
	})
})
<|endoftext|>


<commit_before>package fmttab
import (
	"testing"
	"strings"
	"strconv"
	"bytes"
)

func TestCallAndOut(t *testing.T) {
	const name = "<THIS IS NAME>"
	const val = 123
	arr := []map[string]interface{}{
		map[string]interface{}{
			"ID" : val,
			"NAME" : name,
		},
	}
	l := len(arr)
	cur := 0
	tbl := New("", BORDER_DOUBLE, func() (bool, map[string]interface{}) {
		if cur >= l {
			return false, nil
		}
		r := arr[cur]
		cur++
		return true, r
	})
	tbl.AddColumn("ID", 10, ALIGN_LEFT).
	AddColumn("NAME", 40, ALIGN_RIGHT)

	sout := tbl.String()
	if (cur != 1) {
		t.Error("Excepted call datagetter")
	}

	if strings.IndexAny(sout, name) == -1 {
		t.Errorf("Excepted in output %q", name)
	}
	if strings.IndexAny(sout, strconv.Itoa(val)) == -1 {
		t.Errorf("Excepted in output %v", val)
	}

}

func TestTrim(t *testing.T) {
	test := map[struct {
		val string
		end string
		max int
	}]string{
		{"testing", "...", 6} : "tes...",
		{"testing", "...", 7} : "testing",
		{"testing", "...", 2} : "..",
		{"testing", ">", 5} : "test>",
	}
	for key, pair := range test {
		r := TrimEnds(key.val, key.end, key.max)
		if r != pair {
			t.Errorf("Excepted %q, got %q", pair, r)
		}
	}
}

func TestString(t *testing.T) {
	tab := New("Table", BORDER_THIN, nil)
	tab.AddColumn("Column1", 8, ALIGN_LEFT)
	org := "┌────────┐\n│Column1 │\n├────────┤\n└────────┘\n"
	res := tab.String()
	if org != res {
		t.Errorf("Excepted \n%q, got:\n%q", org, res)
	}

	var buf bytes.Buffer

	if cntwrite,err:=tab.WriteTo(&buf);err!=nil {
		t.Error(err)
	} else {
		if int64(len(res))!=cntwrite {
			t.Errorf("Excepted count write:%d, got: %d",len(res),cntwrite)
		}
	}

	var buf2 bytes.Buffer
	tab2 := New("",BORDER_THIN,nil)
	if cntwrite,err:=tab2.WriteTo(&buf2);err!=nil {
		t.Error(err)
	} else {
		if cntwrite!=0 {
			t.Errorf("Excepted count write:-1, got: %d",cntwrite)
		}
	}

}<commit_msg>Fix test<commit_after>package fmttab

import (
	"bytes"
	"strconv"
	"strings"
	"testing"
)

func TestCallAndOut(t *testing.T) {
	const name = "<THIS IS NAME>"
	const val = 123
	arr := []map[string]interface{}{
		map[string]interface{}{
			"ID":   val,
			"NAME": name,
		},
	}
	l := len(arr)
	cur := 0
	tbl := New("", BORDER_DOUBLE, func() (bool, map[string]interface{}) {
		if cur >= l {
			return false, nil
		}
		r := arr[cur]
		cur++
		return true, r
	})
	tbl.AddColumn("ID", 10, ALIGN_LEFT).
		AddColumn("NAME", 40, ALIGN_RIGHT)

	sout := tbl.String()
	if cur != 1 {
		t.Error("Excepted call datagetter")
	}

	if strings.IndexAny(sout, name) == -1 {
		t.Errorf("Excepted in output %q", name)
	}
	if strings.IndexAny(sout, strconv.Itoa(val)) == -1 {
		t.Errorf("Excepted in output %v", val)
	}

}

func TestTrim(t *testing.T) {
	test := map[struct {
		val string
		end string
		max int
	}]string{
		{"testing", "...", 6,}: "tes...",
		{"testing", "...", 7,}: "testing",
		{"testing", "...", 2,}: "..",
		{"testing", ">", 5}:   "test>",
	}
	for key, pair := range test {
		r := TrimEnds(key.val, key.end, key.max)
		if r != pair {
			t.Errorf("Excepted %q, got %q", pair, r)
		}
	}
}

func TestString(t *testing.T) {
	tab := New("Table", BORDER_THIN, nil)
	tab.AddColumn("Column1", 8, ALIGN_LEFT)
	org := "┌────────┐\n│Column1 │\n├────────┤\n└────────┘\n"
	res := tab.String()
	if org != res {
		t.Errorf("Excepted \n%q, got:\n%q", org, res)
	}

	var buf bytes.Buffer

	if cntwrite, err := tab.WriteTo(&buf); err != nil {
		t.Error(err)
	} else {
		if int64(len(res)) != cntwrite {
			t.Errorf("Excepted count write:%d, got: %d", len(res), cntwrite)
		}
	}

	var buf2 bytes.Buffer
	tab2 := New("", BORDER_THIN, nil)
	if cntwrite, err := tab2.WriteTo(&buf2); err != nil {
		t.Error(err)
	} else {
		if cntwrite != 0 {
			t.Errorf("Excepted count write:-1, got: %d", cntwrite)
		}
	}

}
<|endoftext|>


<commit_before><commit_msg>fix manage request<commit_after><|endoftext|>


<commit_before>package goresponse

type State struct {
	Code    int `json:"code"`    // 执行码
	Message int `json:"message"` // 执行描述码
	Status  int `json:"status"`  // 状态码
}

func NewState() *State {
	return &State{
		Code:    1,
		Message: 0,
		Status:  Status_fail,
	}
}

func NewStateCS(status int) *State {
	return &State{
		Code:    status,
		Message: 0,
		Status:  status,
	}
}

func NewStateSuccess() *State {
	return &State{
		Code:    0,
		Message: 0,
		Status:  Status_success,
	}
}

// 设置执行码
func (c *State) C(code int) *State {
	c.Code = code
	return c
}

// 设置执行描述码
func (c *State) M(message int) *State {
	c.Message = message
	return c
}

// 设置状态码
func (c *State) S(status int) *State {
	c.Status = status
	return c
}

func (c *State) CS(status int) *State {
	return c.C(status).S(status)
}

func (c *State) Success() *State {
	return c.C(0).M(0).S(Status_success)
}

func (c *State) Ignore() *State {
	return c.C(0).M(0).S(Status_ignore)
}
<commit_msg>feat: default state<commit_after>package goresponse

const (
	DefaultStateCode    int = 1
	DefaultStateMessage     = 0
	DefaultStateStatus      = Status_fail
)

type State struct {
	Code    int `json:"code"`    // 执行码
	Message int `json:"message"` // 执行描述码
	Status  int `json:"status"`  // 状态码
}

func NewState() *State {
	return &State{
		Code:    DefaultStateCode,
		Message: DefaultStateMessage,
		Status:  DefaultStateStatus,
	}
}

func NewStateCS(status int) *State {
	return &State{
		Code:    status,
		Message: DefaultStateMessage,
		Status:  status,
	}
}

func NewStateSuccess() *State {
	return &State{
		Code:    0,
		Message: DefaultStateMessage,
		Status:  Status_success,
	}
}

// 设置执行码
func (c *State) C(code int) *State {
	c.Code = code
	return c
}

// 设置执行描述码
func (c *State) M(message int) *State {
	c.Message = message
	return c
}

// 设置状态码
func (c *State) S(status int) *State {
	c.Status = status
	return c
}

func (c *State) CS(status int) *State {
	return c.C(status).S(status)
}

func (c *State) Success() *State {
	return c.C(0).M(0).S(Status_success)
}

func (c *State) Ignore() *State {
	return c.C(0).M(0).S(Status_ignore)
}
<|endoftext|>