<commit_before>// delayd2 is an available setTimeout() service for scheduling message sends
package delayd2

import (
	"log"
	"runtime"
	"sync"
	"time"

	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/pmylund/go-cache"
)

// QueueMessage is a message queued in the database.
type QueueMessage struct {
	QueueID   string
	WorkerID  string
	ReleaseAt time.Time
	RelayTo   string
	Payload   string
}

// WorkerConfig is configuration for Worker.
type WorkerConfig struct {
	ID string

	LeaveMessagesOrphanedAtShutdown bool

	NumConsumerFactor int
	NumRelayFactor    int
}

// Worker is the delayd2 worker.
type Worker struct {
	config *WorkerConfig

	driver   Driver
	consumer *Consumer
	relay    *Relay

	succededIDsCache *cache.Cache

	relayCh chan releaseJob

	// signaled from external
	shutdownCh chan struct{}

	// notified when handleRelease() finishes its task
	stoppedCh chan struct{}
}

// NewWorker creates a new worker.
func NewWorker(c *WorkerConfig, driver Driver, consumer *Consumer, relay *Relay) *Worker {
	return &Worker{
		driver:   driver,
		consumer: consumer,
		relay:    relay,

		// Cache succededIDs to prepare for transient errors in the database (e.g. failover/network problems)
		succededIDsCache: cache.New(15*time.Minute, 1*time.Minute),

		config: c,

		relayCh:    make(chan releaseJob, runtime.NumCPU()*c.NumRelayFactor),
		shutdownCh: make(chan struct{}),
		stoppedCh:  make(chan struct{}),
	}
}

// Run starts the worker process.
func (w *Worker) Run() error {
	log.Print("worker: starting delayd2 worker process")

	log.Print("worker: registering delayd2 worker process")
	if err := w.driver.RegisterSession(); err != nil {
		return err
	}

	defer func() {
		log.Print("worker: resetting aborted messages remaining in active queue")
		if _, err := w.driver.ResetActive(); err != nil {
			log.Print("worker: unable to reset messages in active queue:", err)
		}

		if w.config.LeaveMessagesOrphanedAtShutdown {
			log.Print("worker: leaving messages as orphaned")

			if err := w.driver.MarkOrphaned(); err != nil {
				log.Print("worker: unable to mark messages as orphaned:", err)
			}
		}

		log.Print("worker: deregistering delayd2 worker process")
		if err := w.driver.DeregisterSession(); err != nil {
			log.Print("worker: unable to deregister session:", err)
		}
	}()

	log.Print("worker: starting delayd2 process")

	nCPU := runtime.NumCPU()
	for i := 0; i < nCPU*w.config.NumConsumerFactor; i++ {
		log.Print("worker: launching consumer process")
		go w.handleConsume()
	}

	for i := 0; i < nCPU*w.config.NumRelayFactor; i++ {
		log.Print("worker: launching relay worker process")
		go w.relayWorker()
	}

	go w.handleAdoptOrphans()
	go w.handleKeepAlive()
	go w.handleMarkActive()
	go w.handleRelease()

	select {
	case <-w.shutdownCh:
		log.Print("worker: receiving shutdown signal. waiting for the process finished.")
	}

	select {
	case <-w.stoppedCh:
		log.Print("worker: the process finished.")
	}
	return nil
}

func (w *Worker) Stop() {
	close(w.shutdownCh)
}

func (w *Worker) consume() (int64, error) {
	return w.consumer.ConsumeMessages()
}

func (w *Worker) handleKeepAlive() {
	for range time.Tick(1 * time.Second) {
		if err := w.driver.KeepAliveSession(); err != nil {
			log.Printf("worker: unable to keep alived: %s", err)
		}
	}
}

// handleConsume consumes messages in SQS.
func (w *Worker) handleConsume() {
	for {
		begin := time.Now()
		n, err := w.consume()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to consume messages: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages consumed in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) handleAdoptOrphans() {
	for range time.Tick(10 * time.Second) {
		begin := time.Now()

		n, err := w.driver.AdoptOrphans()

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to adopt orphans: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d orphans adopted in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) markActive(begin time.Time) (int64, error) {
	return w.driver.MarkActive(begin)
}

// handleMarkActive marks coming messages in queue.
func (w *Worker) handleMarkActive() {
	for range time.Tick(10 * time.Millisecond) {
		begin := time.Now().Truncate(time.Second)

		n, err := w.markActive(begin)

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to mark messages active: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages marked as active in %s", n, end.Sub(begin))
		}
		time.Sleep(1 * time.Second)
	}
}

// handleRelease releases messages that is ready to fire.
func (w *Worker) handleRelease() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			close(w.relayCh)
			close(w.stoppedCh)
			return
		default:
		}

		if err := w.release(); err != nil {
			log.Printf("worker: unable to release messages: %s", err)
			continue
		}
	}
}

type releaseJob struct {
	*sync.WaitGroup

	RelayTo  string
	Messages []releaseMessage
}

type releaseMessage struct {
	QueueID string
	Payload string
}

// release releases active messages. It returns the number of released message n.
func (w *Worker) release() error {
	messages, err := w.driver.GetActiveMessages()

	if err != nil {
		return err
	}

	// we need QueueID to delete from queue
	batchMap := BuildBatchMap(messages)

	var wg sync.WaitGroup
	for relayTo, batchMsgs := range batchMap {
		for _, rms := range batchMsgs {
			wg.Add(1)
			w.relayCh <- releaseJob{
				WaitGroup: &wg,

				RelayTo:  relayTo,
				Messages: rms,
			}
		}
	}
	wg.Wait()
	return nil
}

func (w *Worker) relayWorker() {
	for rj := range w.relayCh {
		begin := time.Now()
		n := w.releaseBatch(rj)
		end := time.Now()

		if n > 0 {
			log.Printf("worker: %d messages relayed in %s", n, end.Sub(begin))
		}
		rj.Done()
	}
	log.Print("relayworker: closed")
}

func (w *Worker) releaseBatch(rj releaseJob) int64 {
	payloads := make([]string, 0, len(rj.Messages))
	for _, m := range rj.Messages {
		if _, found := w.succededIDsCache.Get(m.QueueID); found {
			// skip since we already relayed
			log.Printf("worker: seeing queue_id in the cache so continueing...")
			continue
		}
		payloads = append(payloads, m.Payload)
	}

	var n int64
	failedIndex := make(map[int]struct{})
	if len(payloads) > 0 {
		if err := w.relay.Relay(rj.RelayTo, payloads); err != nil {
			// only print log here in batch operation
			// TODO: a dead letter queue support
			berrs, batchOk := queue.IsBatchError(err)
			if !batchOk {
				log.Printf("worker: unable to send message due to non batch error. skipping: %s", err)
				return 0
			}

			for _, berr := range berrs {
				if berr.SenderFault {
					log.Printf("worker: unable to send message due to sender's fault: skipping: %s", berr.Message)
				}
				failedIndex[berr.Index] = struct{}{}
			}
		}
	}

	succededIDs := make([]string, 0, len(rj.Messages))
	for i, m := range rj.Messages {
		if _, failed := failedIndex[i]; failed {
			continue
		}
		succededIDs = append(succededIDs, m.QueueID)

		// remember succededIDs for a while to prevent us from relaying message in transient error
		w.succededIDsCache.Set(m.QueueID, struct{}{}, cache.DefaultExpiration)
	}

	if err := w.driver.RemoveMessages(succededIDs); err != nil {
		log.Printf("worker: unable to remove messages from queue: %s", err)
	} else {
		// reset succededIDs here since we succeeded to remove messages from the database
		// When we succeeded to remove messages, the messages will not appear again so it's safe.
		for _, id := range succededIDs {
			w.succededIDsCache.Delete(id)
		}
		n += int64(len(succededIDs))
	}

	return n
}

func BuildBatchMap(messages []*QueueMessage) map[string][][]releaseMessage {
	batchMap := make(map[string][][]releaseMessage)
	for _, m := range messages {
		rm := releaseMessage{
			QueueID: m.QueueID,
			Payload: m.Payload,
		}

		var i int
		i = len(batchMap[m.RelayTo])
		if i == 0 || len(batchMap[m.RelayTo][i-1]) >= 10 {
			batchMap[m.RelayTo] = append(batchMap[m.RelayTo], []releaseMessage{})
		}

		i = len(batchMap[m.RelayTo])

		batchMap[m.RelayTo][i-1] = append(batchMap[m.RelayTo][i-1], rm)
	}
	return batchMap
}
<commit_msg>worker: Stop handleMarkActive worker when shutdownCh is closed<commit_after>// delayd2 is an available setTimeout() service for scheduling message sends
package delayd2

import (
	"log"
	"runtime"
	"sync"
	"time"

	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/pmylund/go-cache"
)

// QueueMessage is a message queued in the database.
type QueueMessage struct {
	QueueID   string
	WorkerID  string
	ReleaseAt time.Time
	RelayTo   string
	Payload   string
}

// WorkerConfig is configuration for Worker.
type WorkerConfig struct {
	ID string

	LeaveMessagesOrphanedAtShutdown bool

	NumConsumerFactor int
	NumRelayFactor    int
}

// Worker is the delayd2 worker.
type Worker struct {
	config *WorkerConfig

	driver   Driver
	consumer *Consumer
	relay    *Relay

	succededIDsCache *cache.Cache

	relayCh chan releaseJob

	// signaled from external
	shutdownCh chan struct{}

	// notified when handleRelease() finishes its task
	stoppedCh chan struct{}
}

// NewWorker creates a new worker.
func NewWorker(c *WorkerConfig, driver Driver, consumer *Consumer, relay *Relay) *Worker {
	return &Worker{
		driver:   driver,
		consumer: consumer,
		relay:    relay,

		// Cache succededIDs to prepare for transient errors in the database (e.g. failover/network problems)
		succededIDsCache: cache.New(15*time.Minute, 1*time.Minute),

		config: c,

		relayCh:    make(chan releaseJob, runtime.NumCPU()*c.NumRelayFactor),
		shutdownCh: make(chan struct{}),
		stoppedCh:  make(chan struct{}),
	}
}

// Run starts the worker process.
func (w *Worker) Run() error {
	log.Print("worker: starting delayd2 worker process")

	log.Print("worker: registering delayd2 worker process")
	if err := w.driver.RegisterSession(); err != nil {
		return err
	}

	defer func() {
		log.Print("worker: resetting aborted messages remaining in active queue")
		if _, err := w.driver.ResetActive(); err != nil {
			log.Print("worker: unable to reset messages in active queue:", err)
		}

		if w.config.LeaveMessagesOrphanedAtShutdown {
			log.Print("worker: leaving messages as orphaned")

			if err := w.driver.MarkOrphaned(); err != nil {
				log.Print("worker: unable to mark messages as orphaned:", err)
			}
		}

		log.Print("worker: deregistering delayd2 worker process")
		if err := w.driver.DeregisterSession(); err != nil {
			log.Print("worker: unable to deregister session:", err)
		}
	}()

	log.Print("worker: starting delayd2 process")

	nCPU := runtime.NumCPU()
	for i := 0; i < nCPU*w.config.NumConsumerFactor; i++ {
		log.Print("worker: launching consumer process")
		go w.handleConsume()
	}

	for i := 0; i < nCPU*w.config.NumRelayFactor; i++ {
		log.Print("worker: launching relay worker process")
		go w.relayWorker()
	}

	go w.handleAdoptOrphans()
	go w.handleKeepAlive()
	go w.handleMarkActive()
	go w.handleRelease()

	select {
	case <-w.shutdownCh:
		log.Print("worker: receiving shutdown signal. waiting for the process finished.")
	}

	select {
	case <-w.stoppedCh:
		log.Print("worker: the process finished.")
	}
	return nil
}

func (w *Worker) Stop() {
	close(w.shutdownCh)
}

func (w *Worker) consume() (int64, error) {
	return w.consumer.ConsumeMessages()
}

func (w *Worker) handleKeepAlive() {
	for range time.Tick(1 * time.Second) {
		if err := w.driver.KeepAliveSession(); err != nil {
			log.Printf("worker: unable to keep alived: %s", err)
		}
	}
}

// handleConsume consumes messages in SQS.
func (w *Worker) handleConsume() {
	for {
		begin := time.Now()
		n, err := w.consume()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to consume messages: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages consumed in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) handleAdoptOrphans() {
	for range time.Tick(10 * time.Second) {
		begin := time.Now()

		n, err := w.driver.AdoptOrphans()

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to adopt orphans: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d orphans adopted in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) markActive(begin time.Time) (int64, error) {
	return w.driver.MarkActive(begin)
}

// handleMarkActive marks coming messages in queue.
func (w *Worker) handleMarkActive() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			log.Print("worker: shutting down marking handler")
			return
		default:
		}

		begin := time.Now().Truncate(time.Second)

		n, err := w.markActive(begin)

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to mark messages active: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages marked as active in %s", n, end.Sub(begin))
		}
		time.Sleep(1 * time.Second)
	}
}

// handleRelease releases messages that is ready to fire.
func (w *Worker) handleRelease() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			close(w.relayCh)
			close(w.stoppedCh)
			return
		default:
		}

		if err := w.release(); err != nil {
			log.Printf("worker: unable to release messages: %s", err)
			continue
		}
	}
}

type releaseJob struct {
	*sync.WaitGroup

	RelayTo  string
	Messages []releaseMessage
}

type releaseMessage struct {
	QueueID string
	Payload string
}

// release releases active messages. It returns the number of released message n.
func (w *Worker) release() error {
	messages, err := w.driver.GetActiveMessages()

	if err != nil {
		return err
	}

	// we need QueueID to delete from queue
	batchMap := BuildBatchMap(messages)

	var wg sync.WaitGroup
	for relayTo, batchMsgs := range batchMap {
		for _, rms := range batchMsgs {
			wg.Add(1)
			w.relayCh <- releaseJob{
				WaitGroup: &wg,

				RelayTo:  relayTo,
				Messages: rms,
			}
		}
	}
	wg.Wait()
	return nil
}

func (w *Worker) relayWorker() {
	for rj := range w.relayCh {
		begin := time.Now()
		n := w.releaseBatch(rj)
		end := time.Now()

		if n > 0 {
			log.Printf("worker: %d messages relayed in %s", n, end.Sub(begin))
		}
		rj.Done()
	}
	log.Print("relayworker: closed")
}

func (w *Worker) releaseBatch(rj releaseJob) int64 {
	payloads := make([]string, 0, len(rj.Messages))
	for _, m := range rj.Messages {
		if _, found := w.succededIDsCache.Get(m.QueueID); found {
			// skip since we already relayed
			log.Printf("worker: seeing queue_id in the cache so continueing...")
			continue
		}
		payloads = append(payloads, m.Payload)
	}

	var n int64
	failedIndex := make(map[int]struct{})
	if len(payloads) > 0 {
		if err := w.relay.Relay(rj.RelayTo, payloads); err != nil {
			// only print log here in batch operation
			// TODO: a dead letter queue support
			berrs, batchOk := queue.IsBatchError(err)
			if !batchOk {
				log.Printf("worker: unable to send message due to non batch error. skipping: %s", err)
				return 0
			}

			for _, berr := range berrs {
				if berr.SenderFault {
					log.Printf("worker: unable to send message due to sender's fault: skipping: %s", berr.Message)
				}
				failedIndex[berr.Index] = struct{}{}
			}
		}
	}

	succededIDs := make([]string, 0, len(rj.Messages))
	for i, m := range rj.Messages {
		if _, failed := failedIndex[i]; failed {
			continue
		}
		succededIDs = append(succededIDs, m.QueueID)

		// remember succededIDs for a while to prevent us from relaying message in transient error
		w.succededIDsCache.Set(m.QueueID, struct{}{}, cache.DefaultExpiration)
	}

	if err := w.driver.RemoveMessages(succededIDs); err != nil {
		log.Printf("worker: unable to remove messages from queue: %s", err)
	} else {
		// reset succededIDs here since we succeeded to remove messages from the database
		// When we succeeded to remove messages, the messages will not appear again so it's safe.
		for _, id := range succededIDs {
			w.succededIDsCache.Delete(id)
		}
		n += int64(len(succededIDs))
	}

	return n
}

func BuildBatchMap(messages []*QueueMessage) map[string][][]releaseMessage {
	batchMap := make(map[string][][]releaseMessage)
	for _, m := range messages {
		rm := releaseMessage{
			QueueID: m.QueueID,
			Payload: m.Payload,
		}

		var i int
		i = len(batchMap[m.RelayTo])
		if i == 0 || len(batchMap[m.RelayTo][i-1]) >= 10 {
			batchMap[m.RelayTo] = append(batchMap[m.RelayTo], []releaseMessage{})
		}

		i = len(batchMap[m.RelayTo])

		batchMap[m.RelayTo][i-1] = append(batchMap[m.RelayTo][i-1], rm)
	}
	return batchMap
}
<|endoftext|>


<commit_before>package smd

import (
	"fmt"
	"math"
	"time"

	"github.com/soniakeys/meeus/julian"
	"github.com/soniakeys/meeus/planetposition"
)

const (
	// AU is one astronomical unit in kilometers.
	AU = 149597870
)

// CelestialObject defines a celestial object.
// Note: globe and elements may be nil; does not support satellites yet.
type CelestialObject struct {
	Name   string
	Radius float64
	a      float64
	μ      float64
	tilt   float64 // Axial tilt
	incl   float64 // Ecliptic inclination
	SOI    float64 // With respect to the Sun
	J2     float64
	J3     float64
	J4     float64
	PP     *planetposition.V87Planet
}

// GM returns μ (which is unexported because it's a lowercase letter)
func (c CelestialObject) GM() float64 {
	return c.μ
}

// J returns the perturbing J_n factor for the provided n.
// Currently only J2 and J3 are supported.
func (c CelestialObject) J(n uint8) float64 {
	switch n {
	case 2:
		return c.J2
	case 3:
		return c.J3
	case 4:
		return c.J4
	default:
		return 0.0
	}
}

// String implements the Stringer interface.
func (c CelestialObject) String() string {
	return c.Name + " body"
}

// Equals returns whether the provided celestial object is the same.
func (c *CelestialObject) Equals(b CelestialObject) bool {
	return c.Name == b.Name && c.Radius == b.Radius && c.a == b.a && c.μ == b.μ && c.SOI == b.SOI && c.J2 == b.J2
}

// HelioOrbit returns the heliocentric position and velocity of this planet at a given time in equatorial coordinates.
// Note that the whole file is loaded. In fact, if we don't, then whoever is the first to call this function will
// set the Epoch at which the ephemeris are available, and that sucks.
func (c *CelestialObject) HelioOrbit(dt time.Time) Orbit {
	if c.Name == "Sun" {
		return *NewOrbitFromRV([]float64{0, 0, 0}, []float64{0, 0, 0}, *c)
	}
	if smdConfig().VSOP87 {
		if c.PP == nil {
			// Load the planet.
			var vsopPosition int
			switch c.Name {
			case "Venus":
				vsopPosition = 2
			case "Earth":
				vsopPosition = 3
			case "Mars":
				vsopPosition = 4
			case "Jupiter":
				vsopPosition = 5
			default:
				panic(fmt.Errorf("unknown object: %s", c.Name))
			}
			planet, err := planetposition.LoadPlanetPath(vsopPosition-1, "")
			if err != nil {
				panic(fmt.Errorf("could not load planet number %d: %s", vsopPosition, err))
			}
			c.PP = planet
		}
		l, b, r := c.PP.Position2000(julian.TimeToJD(dt))
		r *= AU
		v := math.Sqrt(2*Sun.μ/r - Sun.μ/c.a)
		// Get the Cartesian coordinates from L,B,R.
		R, V := make([]float64, 3), make([]float64, 3)
		sB, cB := math.Sincos(b.Rad())
		sL, cL := math.Sincos(l.Rad())
		R[0] = r * cB * cL
		R[1] = r * cB * sL
		R[2] = r * sB
		// Let's find the direction of the velocity vector.
		vDir := cross(R, []float64{0, 0, -1})
		for i := 0; i < 3; i++ {
			V[i] = v * vDir[i] / norm(vDir)
		}
		return *NewOrbitFromRV(R, V, Sun)
	}
	R, V := config.HelioState(c.Name, dt)
	return *NewOrbitFromRV(R, V, Sun)
}

/* Definitions */

// Sun is our closest star.
var Sun = CelestialObject{"Sun", 695700, -1, 1.32712440018 * 1e11, 0.0, 0.0, -1, 0, 0, 0, nil}

// Venus is poisonous.
var Venus = CelestialObject{"Venus", 6052, 108208601, 3.257 * 1e5, 117.36, 3.39458, 0.616e6, 0.000027, 0, 0, nil}

// Earth is home.
var Earth = CelestialObject{"Earth", 6378.1363, 149598023, 3.986004415 * 1e5, 23.4, 0.00005, 924645.0, 1082.6269e-6, -2.5324e-6, -1.6204e-6, nil}

// Mars is the vacation place.
var Mars = CelestialObject{"Mars", 3397.2, 227939282.5616, 4.305 * 1e4, 25.19, 1.85, 576000, 1964e-6, 36e-6, -18e-6, nil}

// Jupiter is big.
var Jupiter = CelestialObject{"Jupiter", 71492.0, 778298361, 1.268 * 1e8, 3.13, 1.30326966, 48.2e6, 0.01475, 0, -0.00058, nil}
<commit_msg>Fixed loading of VSOP87 files<commit_after>package smd

import (
	"fmt"
	"math"
	"time"

	"github.com/soniakeys/meeus/julian"
	"github.com/soniakeys/meeus/planetposition"
)

const (
	// AU is one astronomical unit in kilometers.
	AU = 149597870
)

// CelestialObject defines a celestial object.
// Note: globe and elements may be nil; does not support satellites yet.
type CelestialObject struct {
	Name   string
	Radius float64
	a      float64
	μ      float64
	tilt   float64 // Axial tilt
	incl   float64 // Ecliptic inclination
	SOI    float64 // With respect to the Sun
	J2     float64
	J3     float64
	J4     float64
	PP     *planetposition.V87Planet
}

// GM returns μ (which is unexported because it's a lowercase letter)
func (c CelestialObject) GM() float64 {
	return c.μ
}

// J returns the perturbing J_n factor for the provided n.
// Currently only J2 and J3 are supported.
func (c CelestialObject) J(n uint8) float64 {
	switch n {
	case 2:
		return c.J2
	case 3:
		return c.J3
	case 4:
		return c.J4
	default:
		return 0.0
	}
}

// String implements the Stringer interface.
func (c CelestialObject) String() string {
	return c.Name + " body"
}

// Equals returns whether the provided celestial object is the same.
func (c *CelestialObject) Equals(b CelestialObject) bool {
	return c.Name == b.Name && c.Radius == b.Radius && c.a == b.a && c.μ == b.μ && c.SOI == b.SOI && c.J2 == b.J2
}

// HelioOrbit returns the heliocentric position and velocity of this planet at a given time in equatorial coordinates.
// Note that the whole file is loaded. In fact, if we don't, then whoever is the first to call this function will
// set the Epoch at which the ephemeris are available, and that sucks.
func (c *CelestialObject) HelioOrbit(dt time.Time) Orbit {
	if c.Name == "Sun" {
		return *NewOrbitFromRV([]float64{0, 0, 0}, []float64{0, 0, 0}, *c)
	}
	if smdConfig().VSOP87 {
		if c.PP == nil {
			// Load the planet.
			var vsopPosition int
			switch c.Name {
			case "Venus":
				vsopPosition = 2
			case "Earth":
				vsopPosition = 3
			case "Mars":
				vsopPosition = 4
			case "Jupiter":
				vsopPosition = 5
			default:
				panic(fmt.Errorf("unknown object: %s", c.Name))
			}
			planet, err := planetposition.LoadPlanetPath(vsopPosition-1, smdConfig().VSOP87Dir)
			if err != nil {
				panic(fmt.Errorf("could not load planet number %d: %s", vsopPosition, err))
			}
			c.PP = planet
		}
		l, b, r := c.PP.Position2000(julian.TimeToJD(dt))
		r *= AU
		v := math.Sqrt(2*Sun.μ/r - Sun.μ/c.a)
		// Get the Cartesian coordinates from L,B,R.
		R, V := make([]float64, 3), make([]float64, 3)
		sB, cB := math.Sincos(b.Rad())
		sL, cL := math.Sincos(l.Rad())
		R[0] = r * cB * cL
		R[1] = r * cB * sL
		R[2] = r * sB
		// Let's find the direction of the velocity vector.
		vDir := cross(R, []float64{0, 0, -1})
		for i := 0; i < 3; i++ {
			V[i] = v * vDir[i] / norm(vDir)
		}
		return *NewOrbitFromRV(R, V, Sun)
	}
	R, V := config.HelioState(c.Name, dt)
	return *NewOrbitFromRV(R, V, Sun)
}

/* Definitions */

// Sun is our closest star.
var Sun = CelestialObject{"Sun", 695700, -1, 1.32712440018 * 1e11, 0.0, 0.0, -1, 0, 0, 0, nil}

// Venus is poisonous.
var Venus = CelestialObject{"Venus", 6052, 108208601, 3.257 * 1e5, 117.36, 3.39458, 0.616e6, 0.000027, 0, 0, nil}

// Earth is home.
var Earth = CelestialObject{"Earth", 6378.1363, 149598023, 3.986004415 * 1e5, 23.4, 0.00005, 924645.0, 1082.6269e-6, -2.5324e-6, -1.6204e-6, nil}

// Mars is the vacation place.
var Mars = CelestialObject{"Mars", 3397.2, 227939282.5616, 4.305 * 1e4, 25.19, 1.85, 576000, 1964e-6, 36e-6, -18e-6, nil}

// Jupiter is big.
var Jupiter = CelestialObject{"Jupiter", 71492.0, 778298361, 1.268 * 1e8, 3.13, 1.30326966, 48.2e6, 0.01475, 0, -0.00058, nil}
<|endoftext|>


<commit_before>package acmewrapper

import (
	"bytes"
	"crypto/tls"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestCert(t *testing.T) {
	os.Remove("cert.crt")
	os.Remove("key.pem")

	w, err := New(Config{
		Server:           TESTAPI,
		TOSCallback:      TOSAgree,
		Domains:          TESTDOMAINS,
		PrivateKeyFile:   "testinguser.key",
		RegistrationFile: "testinguser.reg",
		Address:          TLSADDRESS,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",
	})

	require.NoError(t, err)

	c := w.GetCertificate().Certificate[0]

	// Make sure the files were written
	_, err = os.Stat("cert.crt")
	require.False(t, os.IsNotExist(err))
	_, err = os.Stat("key.pem")
	require.False(t, os.IsNotExist(err))

	// Make sure that the key and cert were generated correctly - make the TOS fail,
	// and the renew callback fail, since ACME shouldn't need to be set at all
	hadfailure := false
	w, err = New(Config{
		Server:           TESTAPI,
		TOSCallback:      TOSDecline,
		Domains:          TESTDOMAINS,
		PrivateKeyFile:   "testinguser.key",
		RegistrationFile: "testinguser.reg",
		Address:          TLSADDRESS,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",

		RenewCallback: func() {
			hadfailure = true
		},
	})

	require.NoError(t, err)
	require.False(t, hadfailure)
	require.True(t, bytes.Equal(w.GetCertificate().Certificate[0], c))

	// Now make sure that we can load without ACME enabled using our currentkeys
	w, err = New(Config{
		Server:       TESTAPI,
		AcmeDisabled: true,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",
	})

	require.NoError(t, err)
	require.False(t, hadfailure)
	require.True(t, bytes.Equal(w.GetCertificate().Certificate[0], c))

	// Lastly: Make sure we can start without ACME enabled, but enable it later.
	// NOTE: This also tests our renewal function
	renewnum := 0
	w, err = New(Config{
		AcmeDisabled: true,
		Server:       TESTAPI,
		TOSCallback:  TOSAgree,
		Domains:      TESTDOMAINS,
		Address:      TLSADDRESS,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",

		RenewCallback: func() {
			renewnum++
		},

		RenewTime:  999999999999999, // A ridiculous value so that renew always happens
		RenewCheck: 3,
		RetryDelay: 3,

		RenewFailedCallback: func(err error) {
			require.NoError(t, err)
		},
	})

	require.NoError(t, err)

	// Now start a server with the config
	listener, err := tls.Listen("tcp", TLSADDRESS, w.TLSConfig())
	require.NoError(t, err)
	go func() {
		http.Serve(listener, nil)
	}()

	require.NoError(t, w.AcmeDisabled(false))

	// Now the certificate should be set
	crt := w.GetCertificate()

	time.Sleep(8 * time.Second)

	// The certificate should be renewed
	require.NotEqual(t, crt, w.GetCertificate())
	require.True(t, renewnum >= 2)

	// Stop it from being annoying in the background anymore
	w.Config.RenewCheck = 9999999999
	w.Config.RetryDelay = 9999999999
	w.Config.RenewTime = 500
}
<commit_msg>Added prints to see issues<commit_after>package acmewrapper

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestCert(t *testing.T) {
	os.Remove("cert.crt")
	os.Remove("key.pem")

	w, err := New(Config{
		Server:           TESTAPI,
		TOSCallback:      TOSAgree,
		Domains:          TESTDOMAINS,
		PrivateKeyFile:   "testinguser.key",
		RegistrationFile: "testinguser.reg",
		Address:          TLSADDRESS,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",
	})

	require.NoError(t, err)

	c := w.GetCertificate().Certificate[0]

	// Make sure the files were written
	_, err = os.Stat("cert.crt")
	require.False(t, os.IsNotExist(err))
	_, err = os.Stat("key.pem")
	require.False(t, os.IsNotExist(err))

	// Make sure that the key and cert were generated correctly - make the TOS fail,
	// and the renew callback fail, since ACME shouldn't need to be set at all
	hadfailure := false
	w, err = New(Config{
		Server:           TESTAPI,
		TOSCallback:      TOSDecline,
		Domains:          TESTDOMAINS,
		PrivateKeyFile:   "testinguser.key",
		RegistrationFile: "testinguser.reg",
		Address:          TLSADDRESS,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",

		RenewCallback: func() {
			hadfailure = true
		},
	})

	require.NoError(t, err)
	require.False(t, hadfailure)
	require.True(t, bytes.Equal(w.GetCertificate().Certificate[0], c))

	// Now make sure that we can load without ACME enabled using our currentkeys
	w, err = New(Config{
		Server:       TESTAPI,
		AcmeDisabled: true,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",
	})

	require.NoError(t, err)
	require.False(t, hadfailure)
	require.True(t, bytes.Equal(w.GetCertificate().Certificate[0], c))

	// Lastly: Make sure we can start without ACME enabled, but enable it later.
	// NOTE: This also tests our renewal function
	renewnum := 0
	w, err = New(Config{
		AcmeDisabled: true,
		Server:       TESTAPI,
		TOSCallback:  TOSAgree,
		Domains:      TESTDOMAINS,
		Address:      TLSADDRESS,

		TLSCertFile: "cert.crt",
		TLSKeyFile:  "key.pem",

		RenewCallback: func() {
			renewnum++
		},

		RenewTime:  999999999999999, // A ridiculous value so that renew always happens
		RenewCheck: 3,
		RetryDelay: 3,

		RenewFailedCallback: func(err error) {
			require.NoError(t, err)
		},
	})

	require.NoError(t, err)

	// Now start a server with the config
	listener, err := tls.Listen("tcp", TLSADDRESS, w.TLSConfig())
	require.NoError(t, err)
	go func() {
		http.Serve(listener, nil)
	}()

	fmt.Printf("acmeenable\n")
	require.NoError(t, w.AcmeDisabled(false))

	// Now the certificate should be set
	fmt.Printf("getcert\n")
	crt := w.GetCertificate()
	fmt.Printf("Sleeping for 8 seconds...\n")
	time.Sleep(8 * time.Second)
	fmt.Printf("Done sleeping\n")

	// The certificate should be renewed
	require.NotEqual(t, crt, w.GetCertificate())
	require.True(t, renewnum >= 2)

	// Stop it from being annoying in the background anymore
	w.Config.RenewCheck = 9999999999
	w.Config.RetryDelay = 9999999999
	w.Config.RenewTime = 500
	listener.Close()
}
<|endoftext|>


<commit_before>package disruptor

type Writer struct {
	written  *Cursor
	upstream Barrier
	capacity int64
	previous int64
	gate     int64
}

func NewWriter(written *Cursor, upstream Barrier, capacity int64) *Writer {
	assertPowerOfTwo(capacity)

	return &Writer{
		upstream: upstream,
		written:  written,
		capacity: capacity,
		previous: InitialSequenceValue,
		gate:     InitialSequenceValue,
	}
}

func assertPowerOfTwo(value int64) {
	if value > 0 && (value&(value-1)) != 0 {
		// Wikipedia entry: http://bit.ly/1krhaSB
		panic("The ring capacity must be a power of two, e.g. 2, 4, 8, 16, 32, 64, etc.")
	}
}

func (this *Writer) Reserve(count int64) int64 {
	this.previous += count
	for this.previous-this.capacity > this.gate {
		this.gate = this.upstream.Read(0)
	}
	return this.previous
}

func (this *Writer) Await(next int64) {
	for next-this.capacity > this.gate {
		this.gate = this.upstream.Read(0)
	}
}
<commit_msg>Potential race condition? Really slow consumers may cause this.gate not to be updated.<commit_after>package disruptor

import "time"

type Writer struct {
	written  *Cursor
	upstream Barrier
	capacity int64
	previous int64
	gate     int64
}

func NewWriter(written *Cursor, upstream Barrier, capacity int64) *Writer {
	assertPowerOfTwo(capacity)

	return &Writer{
		upstream: upstream,
		written:  written,
		capacity: capacity,
		previous: InitialSequenceValue,
		gate:     InitialSequenceValue,
	}
}

func assertPowerOfTwo(value int64) {
	if value > 0 && (value&(value-1)) != 0 {
		// Wikipedia entry: http://bit.ly/1krhaSB
		panic("The ring capacity must be a power of two, e.g. 2, 4, 8, 16, 32, 64, etc.")
	}
}

func (this *Writer) Reserve(count int64) int64 {
	this.previous += count

	for spin := int64(0); this.previous-this.capacity > this.gate; spin++ {
		if spin&SpinMask == 0 {
			time.Sleep(time.Nanosecond)
		}

		this.gate = this.upstream.Read(0)
	}

	return this.previous
}

func (this *Writer) Await(next int64) {
	for next-this.capacity > this.gate {
		this.gate = this.upstream.Read(0)
	}
}

const SpinMask = 1024*16 - 1 // arbitrary; we'll want to experiment with different values
<|endoftext|>


<commit_before>package snappystream

import (
	"bufio"
	"errors"
	"fmt"
	"hash/crc32"
	"io"

	"code.google.com/p/snappy-go/snappy"
)

var errClosed = fmt.Errorf("closed")

// BufferedWriter is an io.WriteCloser with behavior similar to writers
// returned by NewWriter but it buffers written data, maximizing block size (to
// improve the output compression ratio) at the cost of speed. Benefits over
// NewWriter are most noticible when individual writes are small and when
// streams are long.
//
// Failure to call a BufferedWriter's Close or Flush methods after it is done
// being written to will likely result in missing data frames which will be
// undetectable in the decoding process.
//
// NOTE: BufferedWriter cannot be instantiated via struct literal and must
// use NewBufferedWriter (i.e. its zero value is not usable).
type BufferedWriter struct {
	err error
	w   *writer
	bw  *bufio.Writer
}

// NewBufferedWriter allocates and returns a BufferedWriter with an internal
// buffer of MaxBlockSize bytes.  If an error occurs writing a block to w, all
// future writes will fail with the same error.  After all data has been
// written, the client should call the Flush method to guarantee all data has
// been forwarded to the underlying io.Writer.
func NewBufferedWriter(w io.Writer) *BufferedWriter {
	_w := NewWriter(w).(*writer)
	return &BufferedWriter{
		w:  _w,
		bw: bufio.NewWriterSize(_w, MaxBlockSize),
	}
}

// Write buffers p internally, encoding and writing a block to the underlying
// buffer if the buffer grows beyond MaxBlockSize bytes.  The returned int
// will be 0 if there was an error and len(p) otherwise.
func (w *BufferedWriter) Write(p []byte) (int, error) {
	if w.err != nil {
		return 0, w.err
	}

	_, w.err = w.bw.Write(p)
	if w.err != nil {
		return 0, w.err
	}

	return len(p), nil
}

// Flush encodes and writes a block with the contents of w's internal buffer to
// the underlying writer even if the buffer does not contain a full block of
// data (MaxBlockSize bytes).
func (w *BufferedWriter) Flush() error {
	if w.err == nil {
		w.err = w.bw.Flush()
	}

	return w.err
}

// Close flushes w's internal buffer and tears down internal data structures.
// After a successful call to Close method calls on w return an error.  Close
// makes no attempt to close the underlying writer.
func (w *BufferedWriter) Close() error {
	if w.err != nil {
		return w.err
	}

	w.err = w.bw.Flush()
	w.w = nil
	w.bw = nil

	if w.err != nil {
		return w.err
	}

	w.err = errClosed
	return nil
}

type writer struct {
	writer io.Writer

	hdr []byte
	dst []byte

	sentStreamID bool
}

// NewWriter returns an io.Writer that writes its input to an underlying
// io.Writer encoded as a snappy framed stream.  A stream identifier block is
// written to w preceding the first data block.  The returned writer will never
// emit a block with length in bytes greater than MaxBlockSize+4 nor one
// containing more than MaxBlockSize bytes of (uncompressed) data.
//
// For each Write, the returned length will only ever be len(p) or 0,
// regardless of the length of *compressed* bytes written to the wrapped
// io.Writer.  If the returned length is 0 then error will be non-nil.  If
// len(p) exceeds 65536, the slice will be automatically chunked into smaller
// blocks which are all emitted before the call returns.
func NewWriter(w io.Writer) io.Writer {
	return &writer{
		writer: w,

		hdr: make([]byte, 8),
		dst: make([]byte, 4096),
	}
}

func (w *writer) Write(p []byte) (int, error) {
	total := 0
	sz := MaxBlockSize
	var n int
	for i := 0; i < len(p); i += n {
		if i+sz > len(p) {
			sz = len(p) - i
		}

		var err error
		n, err = w.write(p[i : i+sz])
		if err != nil {
			return 0, err
		}
		total += n
	}
	return total, nil
}

// write attempts to encode p as a block and write it to the underlying writer.
// The returned int may not equal p's length if compression below
// MaxBlockSize-4 could not be achieved.
func (w *writer) write(p []byte) (int, error) {
	var err error

	if len(p) > MaxBlockSize {
		return 0, errors.New(fmt.Sprintf("block too large %d > %d", len(p), MaxBlockSize))
	}

	w.dst, err = snappy.Encode(w.dst, p)
	if err != nil {
		return 0, err
	}
	block := w.dst
	n := len(p)
	compressed := true

	// check for data which is better left uncompressed.  this is determined if
	// the encoded content is longer than the source.
	if len(w.dst) >= len(p) {
		compressed = false
		block = p[:n]
	}

	if !w.sentStreamID {
		_, err := w.writer.Write(streamID)
		if err != nil {
			return 0, err
		}
		w.sentStreamID = true
	}

	// set the block type
	if compressed {
		w.hdr[0] = blockCompressed
	} else {
		w.hdr[0] = blockUncompressed
	}

	// 3 byte little endian length of encoded content
	length := uint32(len(block)) + 4 // +4 for checksum
	w.hdr[1] = byte(length)
	w.hdr[2] = byte(length >> 8)
	w.hdr[3] = byte(length >> 16)

	// 4 byte little endian CRC32 checksum of decoded content
	checksum := maskChecksum(crc32.Checksum(p[:n], crcTable))
	w.hdr[4] = byte(checksum)
	w.hdr[5] = byte(checksum >> 8)
	w.hdr[6] = byte(checksum >> 16)
	w.hdr[7] = byte(checksum >> 24)

	_, err = w.writer.Write(w.hdr)
	if err != nil {
		return 0, err
	}

	_, err = w.writer.Write(block)
	if err != nil {
		return 0, err
	}

	return n, nil
}
<commit_msg>avoid unnecessary allocations in snappy.Encode<commit_after>package snappystream

import (
	"bufio"
	"errors"
	"fmt"
	"hash/crc32"
	"io"

	"code.google.com/p/snappy-go/snappy"
)

var errClosed = fmt.Errorf("closed")

// BufferedWriter is an io.WriteCloser with behavior similar to writers
// returned by NewWriter but it buffers written data, maximizing block size (to
// improve the output compression ratio) at the cost of speed. Benefits over
// NewWriter are most noticible when individual writes are small and when
// streams are long.
//
// Failure to call a BufferedWriter's Close or Flush methods after it is done
// being written to will likely result in missing data frames which will be
// undetectable in the decoding process.
//
// NOTE: BufferedWriter cannot be instantiated via struct literal and must
// use NewBufferedWriter (i.e. its zero value is not usable).
type BufferedWriter struct {
	err error
	w   *writer
	bw  *bufio.Writer
}

// NewBufferedWriter allocates and returns a BufferedWriter with an internal
// buffer of MaxBlockSize bytes.  If an error occurs writing a block to w, all
// future writes will fail with the same error.  After all data has been
// written, the client should call the Flush method to guarantee all data has
// been forwarded to the underlying io.Writer.
func NewBufferedWriter(w io.Writer) *BufferedWriter {
	_w := NewWriter(w).(*writer)
	return &BufferedWriter{
		w:  _w,
		bw: bufio.NewWriterSize(_w, MaxBlockSize),
	}
}

// Write buffers p internally, encoding and writing a block to the underlying
// buffer if the buffer grows beyond MaxBlockSize bytes.  The returned int
// will be 0 if there was an error and len(p) otherwise.
func (w *BufferedWriter) Write(p []byte) (int, error) {
	if w.err != nil {
		return 0, w.err
	}

	_, w.err = w.bw.Write(p)
	if w.err != nil {
		return 0, w.err
	}

	return len(p), nil
}

// Flush encodes and writes a block with the contents of w's internal buffer to
// the underlying writer even if the buffer does not contain a full block of
// data (MaxBlockSize bytes).
func (w *BufferedWriter) Flush() error {
	if w.err == nil {
		w.err = w.bw.Flush()
	}

	return w.err
}

// Close flushes w's internal buffer and tears down internal data structures.
// After a successful call to Close method calls on w return an error.  Close
// makes no attempt to close the underlying writer.
func (w *BufferedWriter) Close() error {
	if w.err != nil {
		return w.err
	}

	w.err = w.bw.Flush()
	w.w = nil
	w.bw = nil

	if w.err != nil {
		return w.err
	}

	w.err = errClosed
	return nil
}

type writer struct {
	writer io.Writer

	hdr []byte
	dst []byte

	sentStreamID bool
}

// NewWriter returns an io.Writer that writes its input to an underlying
// io.Writer encoded as a snappy framed stream.  A stream identifier block is
// written to w preceding the first data block.  The returned writer will never
// emit a block with length in bytes greater than MaxBlockSize+4 nor one
// containing more than MaxBlockSize bytes of (uncompressed) data.
//
// For each Write, the returned length will only ever be len(p) or 0,
// regardless of the length of *compressed* bytes written to the wrapped
// io.Writer.  If the returned length is 0 then error will be non-nil.  If
// len(p) exceeds 65536, the slice will be automatically chunked into smaller
// blocks which are all emitted before the call returns.
func NewWriter(w io.Writer) io.Writer {
	return &writer{
		writer: w,

		hdr: make([]byte, 8),
		dst: make([]byte, 4096),
	}
}

func (w *writer) Write(p []byte) (int, error) {
	total := 0
	sz := MaxBlockSize
	var n int
	for i := 0; i < len(p); i += n {
		if i+sz > len(p) {
			sz = len(p) - i
		}

		var err error
		n, err = w.write(p[i : i+sz])
		if err != nil {
			return 0, err
		}
		total += n
	}
	return total, nil
}

// write attempts to encode p as a block and write it to the underlying writer.
// The returned int may not equal p's length if compression below
// MaxBlockSize-4 could not be achieved.
func (w *writer) write(p []byte) (int, error) {
	var err error

	if len(p) > MaxBlockSize {
		return 0, errors.New(fmt.Sprintf("block too large %d > %d", len(p), MaxBlockSize))
	}

	w.dst = w.dst[:cap(w.dst)] // Encode does dumb resize w/o context. reslice avoids alloc.
	w.dst, err = snappy.Encode(w.dst, p)
	if err != nil {
		return 0, err
	}
	block := w.dst
	n := len(p)
	compressed := true

	// check for data which is better left uncompressed.  this is determined if
	// the encoded content is longer than the source.
	if len(w.dst) >= len(p) {
		compressed = false
		block = p[:n]
	}

	if !w.sentStreamID {
		_, err := w.writer.Write(streamID)
		if err != nil {
			return 0, err
		}
		w.sentStreamID = true
	}

	// set the block type
	if compressed {
		w.hdr[0] = blockCompressed
	} else {
		w.hdr[0] = blockUncompressed
	}

	// 3 byte little endian length of encoded content
	length := uint32(len(block)) + 4 // +4 for checksum
	w.hdr[1] = byte(length)
	w.hdr[2] = byte(length >> 8)
	w.hdr[3] = byte(length >> 16)

	// 4 byte little endian CRC32 checksum of decoded content
	checksum := maskChecksum(crc32.Checksum(p[:n], crcTable))
	w.hdr[4] = byte(checksum)
	w.hdr[5] = byte(checksum >> 8)
	w.hdr[6] = byte(checksum >> 16)
	w.hdr[7] = byte(checksum >> 24)

	_, err = w.writer.Write(w.hdr)
	if err != nil {
		return 0, err
	}

	_, err = w.writer.Write(block)
	if err != nil {
		return 0, err
	}

	return n, nil
}
<|endoftext|>


<commit_before>package saml

import (
	"errors"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

const (
	xmlResponseID = "urn:oasis:names:tc:SAML:2.0:protocol:Response"
	xmlRequestID  = "urn:oasis:names:tc:SAML:2.0:protocol:AuthnRequest"
)

// SignRequest sign a SAML 2.0 AuthnRequest
// `privateKeyPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func SignRequest(xml string, privateKeyPath string) (string, error) {
	return sign(xml, privateKeyPath, xmlRequestID)
}

// SignResponse sign a SAML 2.0 Response
// `privateKeyPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func SignResponse(xml string, privateKeyPath string) (string, error) {
	return sign(xml, privateKeyPath, xmlResponseID)
}

func sign(xml string, privateKeyPath string, id string) (string, error) {

	samlXmlsecInput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecInput.Name())
	samlXmlsecInput.WriteString("<?xml version='1.0' encoding='UTF-8'?>\n")
	samlXmlsecInput.WriteString(xml)
	samlXmlsecInput.Close()

	samlXmlsecOutput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecOutput.Name())
	samlXmlsecOutput.Close()

	// fmt.Println("xmlsec1", "--sign", "--privkey-pem", privateKeyPath,
	// 	"--id-attr:ID", id,
	// 	"--output", samlXmlsecOutput.Name(), samlXmlsecInput.Name())
	output, err := exec.Command("xmlsec1", "--sign", "--privkey-pem", privateKeyPath,
		"--id-attr:ID", id,
		"--output", samlXmlsecOutput.Name(), samlXmlsecInput.Name()).CombinedOutput()
	if err != nil {
		return "", errors.New(err.Error() + " : " + string(output))
	}

	samlSignedRequest, err := ioutil.ReadFile(samlXmlsecOutput.Name())
	if err != nil {
		return "", err
	}
	samlSignedRequestXML := strings.Trim(string(samlSignedRequest), "\n")
	return samlSignedRequestXML, nil
}

// DecryptResponse decrypt a SAML 2.0 xml using his private key
func DecryptResponse(xml string, privateKeyPath string) (string, error) {
	return decrypt(xml, privateKeyPath)
}

func decrypt(xml string, privateKeyPath string) (string, error) {
	samlXmlsecInput, err := ioutil.TempFile(os.TempDir(), "saml-resp")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecInput.Name())
	samlXmlsecInput.WriteString(xml)
	samlXmlsecInput.Close()

	samlXmlsecOutput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecOutput.Name())
	samlXmlsecOutput.Close()

	output, err := exec.Command("xmlsec1", "--decrypt", "--privkey-pem", privateKeyPath,
		"--output", samlXmlsecOutput.Name(), samlXmlsecInput.Name()).CombinedOutput()
	defer deleteTempFile(samlXmlsecInput.Name())
	if err != nil {
		return "", errors.New(err.Error() + " : " + string(output))
	}
	decryptedSAMLResponse, err := ioutil.ReadFile(samlXmlsecOutput.Name())
	if err != nil {
		return "", err
	}
	decryptedSAMLResponseXML := strings.Trim(string(decryptedSAMLResponse), "\n")
	return decryptedSAMLResponseXML, nil
}

// VerifyResponseSignature verify signature of a SAML 2.0 Response document
// `publicCertPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func VerifyResponseSignature(xml string, publicCertPath string) error {
	return verify(xml, publicCertPath, xmlResponseID)
}

// VerifyRequestSignature verify signature of a SAML 2.0 AuthnRequest document
// `publicCertPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func VerifyRequestSignature(xml string, publicCertPath string) error {
	return verify(xml, publicCertPath, xmlRequestID)
}

func verify(xml string, publicCertPath string, id string) error {
	//Write saml to
	samlXmlsecInput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return err
	}

	samlXmlsecInput.WriteString(xml)
	samlXmlsecInput.Close()
	defer deleteTempFile(samlXmlsecInput.Name())

	//fmt.Println("xmlsec1", "--verify", "--pubkey-cert-pem", publicCertPath, "--id-attr:ID", id, samlXmlsecInput.Name())
	_, err = exec.Command("xmlsec1", "--verify", "--pubkey-cert-pem", publicCertPath, "--id-attr:ID", id, samlXmlsecInput.Name()).CombinedOutput()
	if err != nil {
		return errors.New("error verifing signature: " + err.Error())
	}
	return nil
}

// deleteTempFile remove a file and ignore error
// Intended to be called in a defer after the creation of a temp file to ensure cleanup
func deleteTempFile(filename string) {
	_ = os.Remove(filename)
}
<commit_msg>add id for assertion. also add method to verify assertion using said id<commit_after>package saml

import (
	"errors"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

const (
	xmlResponseID  = "urn:oasis:names:tc:SAML:2.0:protocol:Response"
	xmlAssertionID = "Assertion"
	//xmlResponseID = "urn:oasis:names:tc:SAML:2.0:protocol:Assertion"
	xmlRequestID = "urn:oasis:names:tc:SAML:2.0:protocol:AuthnRequest"
)

// SignRequest sign a SAML 2.0 AuthnRequest
// `privateKeyPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func SignRequest(xml string, privateKeyPath string) (string, error) {
	return sign(xml, privateKeyPath, xmlRequestID)
}

// SignResponse sign a SAML 2.0 Response
// `privateKeyPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func SignResponse(xml string, privateKeyPath string) (string, error) {
	return sign(xml, privateKeyPath, xmlResponseID)
}

func sign(xml string, privateKeyPath string, id string) (string, error) {

	samlXmlsecInput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecInput.Name())
	samlXmlsecInput.WriteString("<?xml version='1.0' encoding='UTF-8'?>\n")
	samlXmlsecInput.WriteString(xml)
	samlXmlsecInput.Close()

	samlXmlsecOutput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecOutput.Name())
	samlXmlsecOutput.Close()

	// fmt.Println("xmlsec1", "--sign", "--privkey-pem", privateKeyPath,
	// 	"--id-attr:ID", id,
	// 	"--output", samlXmlsecOutput.Name(), samlXmlsecInput.Name())
	output, err := exec.Command("xmlsec1", "--sign", "--privkey-pem", privateKeyPath,
		"--id-attr:ID", id,
		"--output", samlXmlsecOutput.Name(), samlXmlsecInput.Name()).CombinedOutput()
	if err != nil {
		return "", errors.New(err.Error() + " : " + string(output))
	}

	samlSignedRequest, err := ioutil.ReadFile(samlXmlsecOutput.Name())
	if err != nil {
		return "", err
	}
	samlSignedRequestXML := strings.Trim(string(samlSignedRequest), "\n")
	return samlSignedRequestXML, nil
}

// DecryptResponse decrypt a SAML 2.0 xml using his private key
func DecryptResponse(xml string, privateKeyPath string) (string, error) {
	return decrypt(xml, privateKeyPath)
}

func decrypt(xml string, privateKeyPath string) (string, error) {
	samlXmlsecInput, err := ioutil.TempFile(os.TempDir(), "saml-resp")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecInput.Name())
	samlXmlsecInput.WriteString(xml)
	samlXmlsecInput.Close()

	samlXmlsecOutput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return "", err
	}
	defer deleteTempFile(samlXmlsecOutput.Name())
	samlXmlsecOutput.Close()

	output, err := exec.Command("xmlsec1", "--decrypt", "--privkey-pem", privateKeyPath,
		"--output", samlXmlsecOutput.Name(), samlXmlsecInput.Name()).CombinedOutput()
	defer deleteTempFile(samlXmlsecInput.Name())
	if err != nil {
		return "", errors.New(err.Error() + " : " + string(output))
	}
	decryptedSAMLResponse, err := ioutil.ReadFile(samlXmlsecOutput.Name())
	if err != nil {
		return "", err
	}
	decryptedSAMLResponseXML := strings.Trim(string(decryptedSAMLResponse), "\n")
	return decryptedSAMLResponseXML, nil
}

// VerifyResponseSignature verify signature of a SAML 2.0 Response document
// `publicCertPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func VerifyResponseSignature(xml string, publicCertPath string) error {
	return verify(xml, publicCertPath, xmlResponseID)
}

// VerifyAssertionSignature verifies the signature of the assertion
// in a SAML 2.0 Response document. If encrypted, decryption must occur prior.
func VerifyAssertionSignature(xml string, publicCertPath string) error {
	return verify(xml, publicCertPath, xmlAssertionID)
}

// VerifyRequestSignature verify signature of a SAML 2.0 AuthnRequest document
// `publicCertPath` must be a path on the filesystem, xmlsec1 is run out of process
// through `exec`
func VerifyRequestSignature(xml string, publicCertPath string) error {
	return verify(xml, publicCertPath, xmlRequestID)
}

func verify(xml string, publicCertPath string, id string) error {
	//Write saml to
	samlXmlsecInput, err := ioutil.TempFile(os.TempDir(), "tmpgs")
	if err != nil {
		return err
	}

	samlXmlsecInput.WriteString(xml)
	samlXmlsecInput.Close()
	defer deleteTempFile(samlXmlsecInput.Name())

	//fmt.Println("xmlsec1", "--verify", "--pubkey-cert-pem", publicCertPath, "--id-attr:ID", id, samlXmlsecInput.Name())
	_, err = exec.Command("xmlsec1", "--verify", "--pubkey-cert-pem", publicCertPath, "--id-attr:ID", id, samlXmlsecInput.Name()).CombinedOutput()
	if err != nil {
		return errors.New("error verifing signature: " + err.Error())
	}
	return nil
}

// deleteTempFile remove a file and ignore error
// Intended to be called in a defer after the creation of a temp file to ensure cleanup
func deleteTempFile(filename string) {
	_ = os.Remove(filename)
}
<|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 cjk

import (
	"github.com/blevesearch/bleve/analysis"
	"github.com/blevesearch/bleve/registry"

	"github.com/blevesearch/bleve/analysis/token/lowercase"
	"github.com/blevesearch/bleve/analysis/tokenizer/unicode"
)

const AnalyzerName = "cjk"

func AnalyzerConstructor(config map[string]interface{}, cache *registry.Cache) (*analysis.Analyzer, error) {
	whitespaceTokenizer, err := cache.TokenizerNamed(unicode.Name)
	if err != nil {
		return nil, err
	}
	widthFilter, err := cache.TokenFilterNamed(WidthName)
	if err != nil {
		return nil, err
	}
	toLowerFilter, err := cache.TokenFilterNamed(lowercase.Name)
	if err != nil {
		return nil, err
	}
	bigramFilter, err := cache.TokenFilterNamed(BigramName)
	if err != nil {
		return nil, err
	}
	rv := analysis.Analyzer{
		Tokenizer: whitespaceTokenizer,
		TokenFilters: []analysis.TokenFilter{
			widthFilter,
			toLowerFilter,
			bigramFilter,
		},
	}
	return &rv, nil
}

func init() {
	registry.RegisterAnalyzer(AnalyzerName, AnalyzerConstructor)
}
<commit_msg>fix the confusing variable name<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 cjk

import (
	"github.com/blevesearch/bleve/analysis"
	"github.com/blevesearch/bleve/registry"

	"github.com/blevesearch/bleve/analysis/token/lowercase"
	"github.com/blevesearch/bleve/analysis/tokenizer/unicode"
)

const AnalyzerName = "cjk"

func AnalyzerConstructor(config map[string]interface{}, cache *registry.Cache) (*analysis.Analyzer, error) {
	tokenizer, err := cache.TokenizerNamed(unicode.Name)
	if err != nil {
		return nil, err
	}
	widthFilter, err := cache.TokenFilterNamed(WidthName)
	if err != nil {
		return nil, err
	}
	toLowerFilter, err := cache.TokenFilterNamed(lowercase.Name)
	if err != nil {
		return nil, err
	}
	bigramFilter, err := cache.TokenFilterNamed(BigramName)
	if err != nil {
		return nil, err
	}
	rv := analysis.Analyzer{
		Tokenizer: tokenizer,
		TokenFilters: []analysis.TokenFilter{
			widthFilter,
			toLowerFilter,
			bigramFilter,
		},
	}
	return &rv, nil
}

func init() {
	registry.RegisterAnalyzer(AnalyzerName, AnalyzerConstructor)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 persist provides persistent implementation of the interfaces
// defined in core.name_mapping.go (please see the documentation in there)
package persist

import (
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"sync"
	"time"

	"flag"
	"github.com/ligato/cn-infra/core"
	log "github.com/ligato/cn-infra/logging/logrus"
	"github.com/ligato/vpp-agent/idxvpp"
	"github.com/ligato/vpp-agent/idxvpp/nametoidx"
)

var (
	// configuration file path
	idxMapConfigFile string

	// configuration common to all mappings
	gConfig *nametoidx.Config
)

// init is here only for parsing program arguments
func init() {
	flag.StringVar(&idxMapConfigFile, "idxmap-config", "",
		"Location of the configuration file for index-to-name maps; also set via 'IDXMAP_CONFIG' env variable.")

	if gConfig == nil {
		var err error
		gConfig, err = nametoidx.ConfigFromFile(idxMapConfigFile)
		if err != nil {
			log.DefaultLogger().WithFields(log.Fields{"filepath": idxMapConfigFile, "err": err}).Warn(
				"Failed to load idxmap configuration file")
		} else {
			log.DefaultLogger().WithFields(log.Fields{"filepath": idxMapConfigFile}).Debug(
				"Loaded idxmap configuration file")
		}
	}
}

// Marshalling loads the config and starts watching for changes.
func Marshalling(agentLabel string, idxMap idxvpp.NameToIdx, loadedFromFile idxvpp.NameToIdxRW) error {
	log.DefaultLogger().Debug("Persistence")

	changes := make(chan idxvpp.NameToIdxDto, 1000)
	fileName := idxMap.GetRegistryTitle() + ".json"
	persist := NewNameToIdxPersist(fileName, gConfig, agentLabel, changes)
	err := persist.Init()
	if err != nil {
		return err
	}

	idxMap.Watch(core.PluginName("idxpersist"), nametoidx.ToChan(changes))

	err = persist.loadIdxMapFile(loadedFromFile)
	if err != nil {
		return err
	}

	return nil
}

// NameToIdxPersist is a decorator for NameToIdxRW implementing persistent storage.
type NameToIdxPersist struct {
	// to now about registration
	registrations chan idxvpp.NameToIdxDto

	// Configuration associated with this mapping.
	// Unless this is a unit test, it is the same as the global configuration.
	config *nametoidx.Config

	// (de)serialization of mapping data
	nameToIdx map[string]uint32

	// persistent storage location
	fileDir  string
	filePath string

	// synchronization between the underlying registry and the persistent storage
	syncLock  sync.Mutex
	syncCh    chan bool
	syncAckCh chan error
}

// NewNameToIdxPersist initializes decorator for persistent storage of index to name mapping.
func NewNameToIdxPersist(fileName string, config *nametoidx.Config, namespace string,
	registrations chan idxvpp.NameToIdxDto) *NameToIdxPersist {

	persist := NameToIdxPersist{}
	persist.config = config
	persist.nameToIdx = map[string]uint32{}

	persist.fileDir = path.Join(persist.config.PersistentStorage.Location, namespace)
	persist.filePath = path.Join(persist.fileDir, fileName)

	persist.registrations = registrations

	return &persist
}

// Init starts go routine that watches chan idxvpp.NameToIdxDto
func (persist *NameToIdxPersist) Init() error {
	persist.syncCh = make(chan bool)
	persist.syncAckCh = make(chan error)

	offset := rand.Int63n(int64(persist.config.PersistentStorage.MaxSyncStartDelay))
	go persist.periodicIdxMapSync(time.Duration(offset))

	return nil
}

// loadIdxMapFile loads persistently stored entries of the associated registry.
func (persist *NameToIdxPersist) loadIdxMapFile(loadedFromFile idxvpp.NameToIdxRW) error {
	if _, err := os.Stat(persist.filePath); os.IsNotExist(err) {
		log.DefaultLogger().WithFields(log.Fields{"Filepath": persist.filePath}).Debug(
			"Persistent storage for name to index mapping doesn't exist yet")
		return nil
	}
	idxMapData, err := ioutil.ReadFile(persist.filePath)
	if err != nil {
		return err
	}

	err = json.Unmarshal(idxMapData, &persist.nameToIdx)
	if err != nil {
		return err
	}

	for name, idx := range persist.nameToIdx {
		log.DefaultLogger().WithFields(log.Fields{"name": name, "idx": idx}).Debug(
			"Loaded mapping from the persistent storage")
		loadedFromFile.RegisterName(name, idx, nil)
	}
	return nil
}

// periodicIdxMapSync periodically synchronizes the underlying registry with the persistent storage.
func (persist *NameToIdxPersist) periodicIdxMapSync(offset time.Duration) error {
	for {
		select {
		case reg := <-persist.registrations:
			if reg.Del {
				persist.unregisterName(reg.Name)
			} else {
				persist.registerName(reg.Name, reg.Idx)
			}
		case <-persist.syncCh:
			persist.syncAckCh <- persist.syncMapping()
		case <-time.After(persist.config.PersistentStorage.SyncInterval + offset):
			offset = 0
			err := persist.syncMapping()
			if err != nil {
				log.DefaultLogger().WithFields(log.Fields{"Error": err, "Filepath": persist.filePath}).Error(
					"Failed to sync idxMap with the persistent storage")
			}
		}
	}
}

// syncMapping updates the persistent  storage with the new mappings.
// Current implementation simply re-builds the file content from the scratch.
// TODO: NICE-TO-HAVE incremental update
func (persist *NameToIdxPersist) syncMapping() error {
	persist.syncLock.Lock()
	defer persist.syncLock.Unlock()

	idxMapData, err := json.Marshal(persist.nameToIdx)
	if err != nil {
		return err
	}

	err = os.MkdirAll(persist.fileDir, 0777)
	if err != nil {
		return err
	}

	//log.Debug("Persist len=", len(persist.nameToIdx)," ", persist.filePath)

	return ioutil.WriteFile(persist.filePath, idxMapData, 0644)
}

// RegisterName from NameToIdxPersist allows to add a name-to-index mapping into both the underlying registry and
// the persistent storage (with some delay in synchronization).
func (persist *NameToIdxPersist) registerName(name string, idx uint32) {
	persist.nameToIdx[name] = idx
}

// UnregisterName from NameToIdxPersist allows to remove mapping from both the underlying registry and the persistent
// storage.
func (persist *NameToIdxPersist) unregisterName(name string) {
	delete(persist.nameToIdx, name)
}

// Close triggers explicit synchronization and closes the underlying mapping.
func (persist *NameToIdxPersist) Close() error {
	err := persist.Sync()
	if err != nil {
		return err
	}
	return nil
}

// Sync triggers immediate synchronization between the underlying registry and the persistent storage.
// The function doesn't return until the operation has fully finished.
func (persist *NameToIdxPersist) Sync() error {
	persist.syncCh <- true
	return <-persist.syncAckCh
}
<commit_msg>Update persistent_name_mapping.go comments<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 persist asynchronously writes changes in the map (name->idx) to
// file.
package persist

import (
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"sync"
	"time"

	"flag"
	"github.com/ligato/cn-infra/core"
	log "github.com/ligato/cn-infra/logging/logrus"
	"github.com/ligato/vpp-agent/idxvpp"
	"github.com/ligato/vpp-agent/idxvpp/nametoidx"
)

var (
	// configuration file path
	idxMapConfigFile string

	// configuration common to all mappings
	gConfig *nametoidx.Config
)

// init is here only for parsing program arguments
func init() {
	flag.StringVar(&idxMapConfigFile, "idxmap-config", "",
		"Location of the configuration file for index-to-name maps; also set via 'IDXMAP_CONFIG' env variable.")

	if gConfig == nil {
		var err error
		gConfig, err = nametoidx.ConfigFromFile(idxMapConfigFile)
		if err != nil {
			log.DefaultLogger().WithFields(log.Fields{"filepath": idxMapConfigFile, "err": err}).Warn(
				"Failed to load idxmap configuration file")
		} else {
			log.DefaultLogger().WithFields(log.Fields{"filepath": idxMapConfigFile}).Debug(
				"Loaded idxmap configuration file")
		}
	}
}

// Marshalling loads the config and starts watching for changes.
func Marshalling(agentLabel string, idxMap idxvpp.NameToIdx, loadedFromFile idxvpp.NameToIdxRW) error {
	log.DefaultLogger().Debug("Persistence")

	changes := make(chan idxvpp.NameToIdxDto, 1000)
	fileName := idxMap.GetRegistryTitle() + ".json"
	persist := NewNameToIdxPersist(fileName, gConfig, agentLabel, changes)
	err := persist.Init()
	if err != nil {
		return err
	}

	idxMap.Watch(core.PluginName("idxpersist"), nametoidx.ToChan(changes))

	err = persist.loadIdxMapFile(loadedFromFile)
	if err != nil {
		return err
	}

	return nil
}

// NameToIdxPersist is a watcher for changes in index to name mapping
// that persists changes in name to idx mapping.
type NameToIdxPersist struct {
	// to now about registration
	registrations chan idxvpp.NameToIdxDto

	// Configuration associated with this mapping.
	// Unless this is a unit test, it is the same as the global configuration.
	config *nametoidx.Config

	// (de)serialization of mapping data
	nameToIdx map[string]uint32

	// persistent storage location
	fileDir  string
	filePath string

	// synchronization between the underlying registry and the persistent storage
	syncLock  sync.Mutex
	syncCh    chan bool
	syncAckCh chan error
}

// NewNameToIdxPersist creates new instance of watcher of index to name mapping
// that persists changes in name to idx mapping.
func NewNameToIdxPersist(fileName string, config *nametoidx.Config, namespace string,
	registrations chan idxvpp.NameToIdxDto) *NameToIdxPersist {

	persist := NameToIdxPersist{}
	persist.config = config
	persist.nameToIdx = map[string]uint32{}

	persist.fileDir = path.Join(persist.config.PersistentStorage.Location, namespace)
	persist.filePath = path.Join(persist.fileDir, fileName)

	persist.registrations = registrations

	return &persist
}

// Init starts go routine that watches chan idxvpp.NameToIdxDto
func (persist *NameToIdxPersist) Init() error {
	persist.syncCh = make(chan bool)
	persist.syncAckCh = make(chan error)

	offset := rand.Int63n(int64(persist.config.PersistentStorage.MaxSyncStartDelay))
	go persist.periodicIdxMapSync(time.Duration(offset))

	return nil
}

// loadIdxMapFile loads persistently stored entries of the associated registry.
func (persist *NameToIdxPersist) loadIdxMapFile(loadedFromFile idxvpp.NameToIdxRW) error {
	if _, err := os.Stat(persist.filePath); os.IsNotExist(err) {
		log.DefaultLogger().WithFields(log.Fields{"Filepath": persist.filePath}).Debug(
			"Persistent storage for name to index mapping doesn't exist yet")
		return nil
	}
	idxMapData, err := ioutil.ReadFile(persist.filePath)
	if err != nil {
		return err
	}

	err = json.Unmarshal(idxMapData, &persist.nameToIdx)
	if err != nil {
		return err
	}

	for name, idx := range persist.nameToIdx {
		log.DefaultLogger().WithFields(log.Fields{"name": name, "idx": idx}).Debug(
			"Loaded mapping from the persistent storage")
		loadedFromFile.RegisterName(name, idx, nil)
	}
	return nil
}

// periodicIdxMapSync periodically synchronizes the underlying registry with the persistent storage.
func (persist *NameToIdxPersist) periodicIdxMapSync(offset time.Duration) error {
	for {
		select {
		case reg := <-persist.registrations:
			if reg.Del {
				persist.unregisterName(reg.Name)
			} else {
				persist.registerName(reg.Name, reg.Idx)
			}
		case <-persist.syncCh:
			persist.syncAckCh <- persist.syncMapping()
		case <-time.After(persist.config.PersistentStorage.SyncInterval + offset):
			offset = 0
			err := persist.syncMapping()
			if err != nil {
				log.DefaultLogger().WithFields(log.Fields{"Error": err, "Filepath": persist.filePath}).Error(
					"Failed to sync idxMap with the persistent storage")
			}
		}
	}
}

// syncMapping updates the persistent  storage with the new mappings.
// Current implementation simply re-builds the file content from the scratch.
// TODO: NICE-TO-HAVE incremental update
func (persist *NameToIdxPersist) syncMapping() error {
	persist.syncLock.Lock()
	defer persist.syncLock.Unlock()

	idxMapData, err := json.Marshal(persist.nameToIdx)
	if err != nil {
		return err
	}

	err = os.MkdirAll(persist.fileDir, 0777)
	if err != nil {
		return err
	}

	//log.Debug("Persist len=", len(persist.nameToIdx)," ", persist.filePath)

	return ioutil.WriteFile(persist.filePath, idxMapData, 0644)
}

// RegisterName from NameToIdxPersist allows to add a name-to-index mapping into both the underlying registry and
// the persistent storage (with some delay in synchronization).
func (persist *NameToIdxPersist) registerName(name string, idx uint32) {
	persist.nameToIdx[name] = idx
}

// UnregisterName from NameToIdxPersist allows to remove mapping from both the underlying registry and the persistent
// storage.
func (persist *NameToIdxPersist) unregisterName(name string) {
	delete(persist.nameToIdx, name)
}

// Close triggers explicit synchronization and closes the underlying mapping.
func (persist *NameToIdxPersist) Close() error {
	err := persist.Sync()
	if err != nil {
		return err
	}
	return nil
}

// Sync triggers immediate synchronization between the underlying registry and the persistent storage.
// The function doesn't return until the operation has fully finished.
func (persist *NameToIdxPersist) Sync() error {
	persist.syncCh <- true
	return <-persist.syncAckCh
}
<|endoftext|>


<commit_before>// +build linux

package store

import (
	"runtime"
	"testing"
)

func StoreLinux__cadir(t *testing.T) {
	// just grab the linuxStore and make sure it has a cadir member
	s := platform()
	if s.ca == nil || s.ca.empty() {
		t.Errorf("no cadir found on platform: %s", runtime.GOOS)
	}
}
<commit_msg>store/linux: cast to linuxStore to grab linux specific fields<commit_after>// +build linux

package store

import (
	"runtime"
	"testing"
)

func StoreLinux__cadir(t *testing.T) {
	// just grab the linuxStore and make sure it has a cadir member
	s, ok := platform().(linuxStore)
	if !ok {
		t.Error("error casting to linuxStore")
	}
	if s.ca == nil || s.ca.empty() {
		t.Errorf("no cadir found on platform: %s", runtime.GOOS)
	}
}
<|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 mergeplan

import (
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"sort"
	"testing"
)

// Implements the Segment interface for testing,
type segment struct {
	MyId       uint64
	MyFullSize int64
	MyLiveSize int64
}

func (s *segment) Id() uint64      { return s.MyId }
func (s *segment) FullSize() int64 { return s.MyFullSize }
func (s *segment) LiveSize() int64 { return s.MyLiveSize }

func makeLinearSegments(n int) (rv []Segment) {
	for i := 0; i < n; i++ {
		rv = append(rv, &segment{
			MyId:       uint64(i),
			MyFullSize: int64(i),
			MyLiveSize: int64(i),
		})
	}
	return rv
}

// ----------------------------------------

func TestSimplePlan(t *testing.T) {
	segs := makeLinearSegments(10)

	tests := []struct {
		Desc       string
		Segments   []Segment
		Options    *MergePlanOptions
		ExpectPlan *MergePlan
		ExpectErr  error
	}{
		{"nil segments",
			nil, nil, nil, nil},
		{"empty segments",
			[]Segment{}, nil, nil, nil},
		{"1 segment",
			[]Segment{segs[1]},
			nil,
			nil,
			nil,
		},
		{"2 segments",
			[]Segment{
				segs[1],
				segs[2],
			},
			nil,
			&MergePlan{},
			nil,
		},
		{"3 segments",
			[]Segment{
				segs[1],
				segs[2],
				segs[9],
			},
			nil,
			&MergePlan{},
			nil,
		},
		{"many segments",
			[]Segment{
				segs[1],
				segs[2],
				segs[3],
				segs[4],
				segs[5],
				segs[6],
			},
			&MergePlanOptions{
				MaxSegmentsPerTier:   1,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			&MergePlan{
				Tasks: []*MergeTask{
					&MergeTask{
						Segments: []Segment{
							segs[6],
							segs[5],
						},
					},
				},
			},
			nil,
		},
	}

	for testi, test := range tests {
		plan, err := Plan(test.Segments, test.Options)
		if err != test.ExpectErr {
			testj, _ := json.Marshal(&test)
			t.Errorf("testi: %d, test: %s, got err: %v",
				testi, testj, err)
		}
		if !reflect.DeepEqual(plan, test.ExpectPlan) {
			testj, _ := json.Marshal(&test)
			planj, _ := json.Marshal(&plan)
			t.Errorf("testi: %d, test: %s, got plan: %s",
				testi, testj, planj)
		}
	}
}

// ----------------------------------------

func TestSort(t *testing.T) {
	segs := makeLinearSegments(10)

	sort.Sort(byLiveSizeDescending(segs))

	for i := 1; i < len(segs); i++ {
		if segs[i].LiveSize() >= segs[i-1].LiveSize() {
			t.Errorf("not descending")
		}
	}
}

// ----------------------------------------

func TestCalcBudget(t *testing.T) {
	tests := []struct {
		totalSize     int64
		firstTierSize int64
		o             MergePlanOptions
		expect        int
	}{
		{0, 0, MergePlanOptions{}, 0},
		{1, 0, MergePlanOptions{}, 1},
		{9, 0, MergePlanOptions{}, 9},
		{1, 1,
			MergePlanOptions{
				MaxSegmentsPerTier:   1,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			1,
		},
		{21, 1,
			MergePlanOptions{
				MaxSegmentsPerTier:   1,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			5,
		},
		{21, 1,
			MergePlanOptions{
				MaxSegmentsPerTier:   2,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			7,
		},
	}

	for testi, test := range tests {
		res := CalcBudget(test.totalSize, test.firstTierSize, &test.o)
		if res != test.expect {
			t.Errorf("testi: %d, test: %#v, res: %v",
				testi, test, res)
		}
	}
}

// ----------------------------------------

func TestInsert1SameSizedSegmentBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "i1sssbm",
		verbose: os.Getenv("VERBOSE") == "i1sssbm" || os.Getenv("VERBOSE") == "y",
		n:       200,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			spec.segments = append(spec.segments, &segment{
				MyId:       spec.nextSegmentId,
				MyFullSize: 1,
				MyLiveSize: 1,
			})
			spec.nextSegmentId++
		},
	}

	spec.runCycles(t)
}

func TestInsertManySameSizedSegmentsBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "imsssbm",
		verbose: os.Getenv("VERBOSE") == "imsssbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			for i := 0; i < 10; i++ {
				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: 1,
					MyLiveSize: 1,
				})
				spec.nextSegmentId++
			}
		},
	}

	spec.runCycles(t)
}

func TestInsertManySameSizedSegmentsWithDeletionsBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "imssswdbm",
		verbose: os.Getenv("VERBOSE") == "imssswdbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			for i := 0; i < 10; i++ {
				// Deletions are a shrinking of the live size.
				for i, seg := range spec.segments {
					if (spec.cycle+i)%5 == 0 {
						s := seg.(*segment)
						if s.MyLiveSize > 0 {
							s.MyLiveSize -= 1
						}
					}
				}

				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: 1,
					MyLiveSize: 1,
				})
				spec.nextSegmentId++
			}
		},
	}

	spec.runCycles(t)
}

func TestInsertManyDifferentSizedSegmentsBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "imdssbm",
		verbose: os.Getenv("VERBOSE") == "imdssbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			for i := 0; i < 10; i++ {
				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: int64(1 + (i % 5)),
					MyLiveSize: int64(1 + (i % 5)),
				})
				spec.nextSegmentId++
			}
		},
	}

	spec.runCycles(t)
}

func TestManySameSizedSegmentsWithDeletesBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	var numPlansWithTasks int

	spec := testCyclesSpec{
		descrip: "mssswdbm",
		verbose: os.Getenv("VERBOSE") == "mssswdbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			// Deletions are a shrinking of the live size.
			for i, seg := range spec.segments {
				if (spec.cycle+i)%5 == 0 {
					s := seg.(*segment)
					if s.MyLiveSize > 0 {
						s.MyLiveSize -= 1
					}
				}
			}

			for i := 0; i < 10; i++ {
				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: 1,
					MyLiveSize: 1,
				})
				spec.nextSegmentId++
			}
		},
		afterPlan: func(spec *testCyclesSpec, plan *MergePlan) {
			if plan != nil && len(plan.Tasks) > 0 {
				numPlansWithTasks++
			}
		},
	}

	spec.runCycles(t)

	if numPlansWithTasks <= 0 {
		t.Errorf("expected some plans with tasks")
	}
}

// ----------------------------------------

type testCyclesSpec struct {
	descrip string
	verbose bool

	n int // Number of cycles to run.
	o *MergePlanOptions

	beforePlan func(*testCyclesSpec)
	afterPlan  func(*testCyclesSpec, *MergePlan)

	cycle         int
	segments      []Segment
	nextSegmentId uint64
}

func (spec *testCyclesSpec) runCycles(t *testing.T) {
	numPlansWithTasks := 0

	for spec.cycle < spec.n {
		if spec.verbose {
			emit(spec.descrip, spec.cycle, 0, spec.segments, nil)
		}

		if spec.beforePlan != nil {
			spec.beforePlan(spec)
		}

		if spec.verbose {
			emit(spec.descrip, spec.cycle, 1, spec.segments, nil)
		}

		plan, err := Plan(spec.segments, spec.o)
		if err != nil {
			t.Fatalf("expected no err, got: %v", err)
		}

		if spec.afterPlan != nil {
			spec.afterPlan(spec, plan)
		}

		if spec.verbose {
			emit(spec.descrip, spec.cycle, 2, spec.segments, plan)
		}

		if plan != nil {
			if len(plan.Tasks) > 0 {
				numPlansWithTasks++
			}

			for _, task := range plan.Tasks {
				spec.segments = removeSegments(spec.segments, task.Segments)

				var totLiveSize int64
				for _, segment := range task.Segments {
					totLiveSize += segment.LiveSize()
				}

				if totLiveSize > 0 {
					spec.segments = append(spec.segments, &segment{
						MyId:       spec.nextSegmentId,
						MyFullSize: totLiveSize,
						MyLiveSize: totLiveSize,
					})
					spec.nextSegmentId++
				}
			}
		}

		spec.cycle++
	}

	if numPlansWithTasks <= 0 {
		t.Errorf("expected some plans with tasks")
	}
}

func emit(descrip string, cycle int, step int, segments []Segment, plan *MergePlan) {
	if os.Getenv("VERBOSE") == "" {
		return
	}

	suffix := ""
	if plan != nil && len(plan.Tasks) > 0 {
		suffix = "hasPlan"
	}

	fmt.Printf("%s %d.%d ---------- %s\n", descrip, cycle, step, suffix)
	fmt.Printf("%s\n", ToBarChart(descrip, 100, segments, plan))
}
<commit_msg>more zap merge-planner CalcBudget tests at larger sizes<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 mergeplan

import (
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"sort"
	"testing"
)

// Implements the Segment interface for testing,
type segment struct {
	MyId       uint64
	MyFullSize int64
	MyLiveSize int64
}

func (s *segment) Id() uint64      { return s.MyId }
func (s *segment) FullSize() int64 { return s.MyFullSize }
func (s *segment) LiveSize() int64 { return s.MyLiveSize }

func makeLinearSegments(n int) (rv []Segment) {
	for i := 0; i < n; i++ {
		rv = append(rv, &segment{
			MyId:       uint64(i),
			MyFullSize: int64(i),
			MyLiveSize: int64(i),
		})
	}
	return rv
}

// ----------------------------------------

func TestSimplePlan(t *testing.T) {
	segs := makeLinearSegments(10)

	tests := []struct {
		Desc       string
		Segments   []Segment
		Options    *MergePlanOptions
		ExpectPlan *MergePlan
		ExpectErr  error
	}{
		{"nil segments",
			nil, nil, nil, nil},
		{"empty segments",
			[]Segment{}, nil, nil, nil},
		{"1 segment",
			[]Segment{segs[1]},
			nil,
			nil,
			nil,
		},
		{"2 segments",
			[]Segment{
				segs[1],
				segs[2],
			},
			nil,
			&MergePlan{},
			nil,
		},
		{"3 segments",
			[]Segment{
				segs[1],
				segs[2],
				segs[9],
			},
			nil,
			&MergePlan{},
			nil,
		},
		{"many segments",
			[]Segment{
				segs[1],
				segs[2],
				segs[3],
				segs[4],
				segs[5],
				segs[6],
			},
			&MergePlanOptions{
				MaxSegmentsPerTier:   1,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			&MergePlan{
				Tasks: []*MergeTask{
					&MergeTask{
						Segments: []Segment{
							segs[6],
							segs[5],
						},
					},
				},
			},
			nil,
		},
	}

	for testi, test := range tests {
		plan, err := Plan(test.Segments, test.Options)
		if err != test.ExpectErr {
			testj, _ := json.Marshal(&test)
			t.Errorf("testi: %d, test: %s, got err: %v",
				testi, testj, err)
		}
		if !reflect.DeepEqual(plan, test.ExpectPlan) {
			testj, _ := json.Marshal(&test)
			planj, _ := json.Marshal(&plan)
			t.Errorf("testi: %d, test: %s, got plan: %s",
				testi, testj, planj)
		}
	}
}

// ----------------------------------------

func TestSort(t *testing.T) {
	segs := makeLinearSegments(10)

	sort.Sort(byLiveSizeDescending(segs))

	for i := 1; i < len(segs); i++ {
		if segs[i].LiveSize() >= segs[i-1].LiveSize() {
			t.Errorf("not descending")
		}
	}
}

// ----------------------------------------

func TestCalcBudget(t *testing.T) {
	tests := []struct {
		totalSize     int64
		firstTierSize int64
		o             MergePlanOptions
		expect        int
	}{
		{0, 0, MergePlanOptions{}, 0},
		{1, 0, MergePlanOptions{}, 1},
		{9, 0, MergePlanOptions{}, 9},
		{1, 1,
			MergePlanOptions{
				MaxSegmentsPerTier:   1,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			1,
		},
		{21, 1,
			MergePlanOptions{
				MaxSegmentsPerTier:   1,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			5,
		},
		{21, 1,
			MergePlanOptions{
				MaxSegmentsPerTier:   2,
				MaxSegmentSize:       1000,
				TierGrowth:           2.0,
				SegmentsPerMergeTask: 2,
				FloorSegmentSize:     1,
			},
			7,
		},
		{1000, 2000, DefaultMergePlanOptions,
			1},
		{5000, 2000, DefaultMergePlanOptions,
			3},
		{10000, 2000, DefaultMergePlanOptions,
			5},
		{30000, 2000, DefaultMergePlanOptions,
			11},
		{1000000, 2000, DefaultMergePlanOptions,
			24},
		{1000000000, 2000, DefaultMergePlanOptions,
			54},
	}

	for testi, test := range tests {
		res := CalcBudget(test.totalSize, test.firstTierSize, &test.o)
		if res != test.expect {
			t.Errorf("testi: %d, test: %#v, res: %v",
				testi, test, res)
		}
	}
}

// ----------------------------------------

func TestInsert1SameSizedSegmentBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "i1sssbm",
		verbose: os.Getenv("VERBOSE") == "i1sssbm" || os.Getenv("VERBOSE") == "y",
		n:       200,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			spec.segments = append(spec.segments, &segment{
				MyId:       spec.nextSegmentId,
				MyFullSize: 1,
				MyLiveSize: 1,
			})
			spec.nextSegmentId++
		},
	}

	spec.runCycles(t)
}

func TestInsertManySameSizedSegmentsBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "imsssbm",
		verbose: os.Getenv("VERBOSE") == "imsssbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			for i := 0; i < 10; i++ {
				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: 1,
					MyLiveSize: 1,
				})
				spec.nextSegmentId++
			}
		},
	}

	spec.runCycles(t)
}

func TestInsertManySameSizedSegmentsWithDeletionsBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "imssswdbm",
		verbose: os.Getenv("VERBOSE") == "imssswdbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			for i := 0; i < 10; i++ {
				// Deletions are a shrinking of the live size.
				for i, seg := range spec.segments {
					if (spec.cycle+i)%5 == 0 {
						s := seg.(*segment)
						if s.MyLiveSize > 0 {
							s.MyLiveSize -= 1
						}
					}
				}

				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: 1,
					MyLiveSize: 1,
				})
				spec.nextSegmentId++
			}
		},
	}

	spec.runCycles(t)
}

func TestInsertManyDifferentSizedSegmentsBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	spec := testCyclesSpec{
		descrip: "imdssbm",
		verbose: os.Getenv("VERBOSE") == "imdssbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			for i := 0; i < 10; i++ {
				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: int64(1 + (i % 5)),
					MyLiveSize: int64(1 + (i % 5)),
				})
				spec.nextSegmentId++
			}
		},
	}

	spec.runCycles(t)
}

func TestManySameSizedSegmentsWithDeletesBetweenMerges(t *testing.T) {
	o := &MergePlanOptions{
		MaxSegmentSize:       1000,
		MaxSegmentsPerTier:   3,
		TierGrowth:           3.0,
		SegmentsPerMergeTask: 3,
	}

	var numPlansWithTasks int

	spec := testCyclesSpec{
		descrip: "mssswdbm",
		verbose: os.Getenv("VERBOSE") == "mssswdbm" || os.Getenv("VERBOSE") == "y",
		n:       20,
		o:       o,
		beforePlan: func(spec *testCyclesSpec) {
			// Deletions are a shrinking of the live size.
			for i, seg := range spec.segments {
				if (spec.cycle+i)%5 == 0 {
					s := seg.(*segment)
					if s.MyLiveSize > 0 {
						s.MyLiveSize -= 1
					}
				}
			}

			for i := 0; i < 10; i++ {
				spec.segments = append(spec.segments, &segment{
					MyId:       spec.nextSegmentId,
					MyFullSize: 1,
					MyLiveSize: 1,
				})
				spec.nextSegmentId++
			}
		},
		afterPlan: func(spec *testCyclesSpec, plan *MergePlan) {
			if plan != nil && len(plan.Tasks) > 0 {
				numPlansWithTasks++
			}
		},
	}

	spec.runCycles(t)

	if numPlansWithTasks <= 0 {
		t.Errorf("expected some plans with tasks")
	}
}

// ----------------------------------------

type testCyclesSpec struct {
	descrip string
	verbose bool

	n int // Number of cycles to run.
	o *MergePlanOptions

	beforePlan func(*testCyclesSpec)
	afterPlan  func(*testCyclesSpec, *MergePlan)

	cycle         int
	segments      []Segment
	nextSegmentId uint64
}

func (spec *testCyclesSpec) runCycles(t *testing.T) {
	numPlansWithTasks := 0

	for spec.cycle < spec.n {
		if spec.verbose {
			emit(spec.descrip, spec.cycle, 0, spec.segments, nil)
		}

		if spec.beforePlan != nil {
			spec.beforePlan(spec)
		}

		if spec.verbose {
			emit(spec.descrip, spec.cycle, 1, spec.segments, nil)
		}

		plan, err := Plan(spec.segments, spec.o)
		if err != nil {
			t.Fatalf("expected no err, got: %v", err)
		}

		if spec.afterPlan != nil {
			spec.afterPlan(spec, plan)
		}

		if spec.verbose {
			emit(spec.descrip, spec.cycle, 2, spec.segments, plan)
		}

		if plan != nil {
			if len(plan.Tasks) > 0 {
				numPlansWithTasks++
			}

			for _, task := range plan.Tasks {
				spec.segments = removeSegments(spec.segments, task.Segments)

				var totLiveSize int64
				for _, segment := range task.Segments {
					totLiveSize += segment.LiveSize()
				}

				if totLiveSize > 0 {
					spec.segments = append(spec.segments, &segment{
						MyId:       spec.nextSegmentId,
						MyFullSize: totLiveSize,
						MyLiveSize: totLiveSize,
					})
					spec.nextSegmentId++
				}
			}
		}

		spec.cycle++
	}

	if numPlansWithTasks <= 0 {
		t.Errorf("expected some plans with tasks")
	}
}

func emit(descrip string, cycle int, step int, segments []Segment, plan *MergePlan) {
	if os.Getenv("VERBOSE") == "" {
		return
	}

	suffix := ""
	if plan != nil && len(plan.Tasks) > 0 {
		suffix = "hasPlan"
	}

	fmt.Printf("%s %d.%d ---------- %s\n", descrip, cycle, step, suffix)
	fmt.Printf("%s\n", ToBarChart(descrip, 100, segments, plan))
}
<|endoftext|>