<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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 cache

import (
	"sync"
	"time"

	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/util/clock"
)

// Config contains all the settings for a Controller.
type Config struct {
	// The queue for your objects; either a FIFO or
	// a DeltaFIFO. Your Process() function should accept
	// the output of this Queue's Pop() method.
	Queue

	// Something that can list and watch your objects.
	ListerWatcher

	// Something that can process your objects.
	Process ProcessFunc

	// The type of your objects.
	ObjectType runtime.Object

	// Reprocess everything at least this often.
	// Note that if it takes longer for you to clear the queue than this
	// period, you will end up processing items in the order determined
	// by FIFO.Replace(). Currently, this is random. If this is a
	// problem, we can change that replacement policy to append new
	// things to the end of the queue instead of replacing the entire
	// queue.
	FullResyncPeriod time.Duration

	// ShouldResync, if specified, is invoked when the controller's reflector determines the next
	// periodic sync should occur. If this returns true, it means the reflector should proceed with
	// the resync.
	ShouldResync ShouldResyncFunc

	// If true, when Process() returns an error, re-enqueue the object.
	// TODO: add interface to let you inject a delay/backoff or drop
	//       the object completely if desired. Pass the object in
	//       question to this interface as a parameter.
	RetryOnError bool
}

// ShouldResyncFunc is a type of function that indicates if a reflector should perform a
// resync or not. It can be used by a shared informer to support multiple event handlers with custom
// resync periods.
type ShouldResyncFunc func() bool

// ProcessFunc processes a single object.
type ProcessFunc func(obj interface{}) error

// Controller is a generic controller framework.
type controller struct {
	config         Config
	reflector      *Reflector
	reflectorMutex sync.RWMutex
	clock          clock.Clock
}

type Controller interface {
	Run(stopCh <-chan struct{})
	HasSynced() bool
	LastSyncResourceVersion() string
}

// New makes a new Controller from the given Config.
func New(c *Config) Controller {
	ctlr := &controller{
		config: *c,
		clock:  &clock.RealClock{},
	}
	return ctlr
}

// Run begins processing items, and will continue until a value is sent down stopCh.
// It's an error to call Run more than once.
// Run blocks; call via go.
func (c *controller) Run(stopCh <-chan struct{}) {
	defer utilruntime.HandleCrash()
	go func() {
		<-stopCh
		c.config.Queue.Close()
	}()
	r := NewReflector(
		c.config.ListerWatcher,
		c.config.ObjectType,
		c.config.Queue,
		c.config.FullResyncPeriod,
	)
	r.ShouldResync = c.config.ShouldResync
	r.clock = c.clock

	c.reflectorMutex.Lock()
	c.reflector = r
	c.reflectorMutex.Unlock()

	r.RunUntil(stopCh)

	wait.Until(c.processLoop, time.Second, stopCh)
}

// Returns true once this controller has completed an initial resource listing
func (c *controller) HasSynced() bool {
	return c.config.Queue.HasSynced()
}

func (c *controller) LastSyncResourceVersion() string {
	if c.reflector == nil {
		return ""
	}
	return c.reflector.LastSyncResourceVersion()
}

// processLoop drains the work queue.
// TODO: Consider doing the processing in parallel. This will require a little thought
// to make sure that we don't end up processing the same object multiple times
// concurrently.
//
// TODO: Plumb through the stopCh here (and down to the queue) so that this can
// actually exit when the controller is stopped. Or just give up on this stuff
// ever being stoppable. Converting this whole package to use Context would
// also be helpful.
func (c *controller) processLoop() {
	for {
		obj, err := c.config.Queue.Pop(PopProcessFunc(c.config.Process))
		if err != nil {
			if err == FIFOClosedError {
				return
			}
			if c.config.RetryOnError {
				// This is the safe way to re-enqueue.
				c.config.Queue.AddIfNotPresent(obj)
			}
		}
	}
}

// ResourceEventHandler can handle notifications for events that happen to a
// resource. The events are informational only, so you can't return an
// error.
//  * OnAdd is called when an object is added.
//  * OnUpdate is called when an object is modified. Note that oldObj is the
//      last known state of the object-- it is possible that several changes
//      were combined together, so you can't use this to see every single
//      change. OnUpdate is also called when a re-list happens, and it will
//      get called even if nothing changed. This is useful for periodically
//      evaluating or syncing something.
//  * OnDelete will get the final state of the item if it is known, otherwise
//      it will get an object of type DeletedFinalStateUnknown. This can
//      happen if the watch is closed and misses the delete event and we don't
//      notice the deletion until the subsequent re-list.
type ResourceEventHandler interface {
	OnAdd(obj interface{})
	OnUpdate(oldObj, newObj interface{})
	OnDelete(obj interface{})
}

// ResourceEventHandlerFuncs is an adaptor to let you easily specify as many or
// as few of the notification functions as you want while still implementing
// ResourceEventHandler.
type ResourceEventHandlerFuncs struct {
	AddFunc    func(obj interface{})
	UpdateFunc func(oldObj, newObj interface{})
	DeleteFunc func(obj interface{})
}

// OnAdd calls AddFunc if it's not nil.
func (r ResourceEventHandlerFuncs) OnAdd(obj interface{}) {
	if r.AddFunc != nil {
		r.AddFunc(obj)
	}
}

// OnUpdate calls UpdateFunc if it's not nil.
func (r ResourceEventHandlerFuncs) OnUpdate(oldObj, newObj interface{}) {
	if r.UpdateFunc != nil {
		r.UpdateFunc(oldObj, newObj)
	}
}

// OnDelete calls DeleteFunc if it's not nil.
func (r ResourceEventHandlerFuncs) OnDelete(obj interface{}) {
	if r.DeleteFunc != nil {
		r.DeleteFunc(obj)
	}
}

// DeletionHandlingMetaNamespaceKeyFunc checks for
// DeletedFinalStateUnknown objects before calling
// MetaNamespaceKeyFunc.
func DeletionHandlingMetaNamespaceKeyFunc(obj interface{}) (string, error) {
	if d, ok := obj.(DeletedFinalStateUnknown); ok {
		return d.Key, nil
	}
	return MetaNamespaceKeyFunc(obj)
}

// NewInformer returns a Store and a controller for populating the store
// while also providing event notifications. You should only used the returned
// Store for Get/List operations; Add/Modify/Deletes will cause the event
// notifications to be faulty.
//
// Parameters:
//  * lw is list and watch functions for the source of the resource you want to
//    be informed of.
//  * objType is an object of the type that you expect to receive.
//  * resyncPeriod: if non-zero, will re-list this often (you will get OnUpdate
//    calls, even if nothing changed). Otherwise, re-list will be delayed as
//    long as possible (until the upstream source closes the watch or times out,
//    or you stop the controller).
//  * h is the object you want notifications sent to.
//
func NewInformer(
	lw ListerWatcher,
	objType runtime.Object,
	resyncPeriod time.Duration,
	h ResourceEventHandler,
) (Store, Controller) {
	// This will hold the client state, as we know it.
	clientState := NewStore(DeletionHandlingMetaNamespaceKeyFunc)

	// This will hold incoming changes. Note how we pass clientState in as a
	// KeyLister, that way resync operations will result in the correct set
	// of update/delete deltas.
	fifo := NewDeltaFIFO(MetaNamespaceKeyFunc, nil, clientState)

	cfg := &Config{
		Queue:            fifo,
		ListerWatcher:    lw,
		ObjectType:       objType,
		FullResyncPeriod: resyncPeriod,
		RetryOnError:     false,

		Process: func(obj interface{}) error {
			// from oldest to newest
			for _, d := range obj.(Deltas) {
				switch d.Type {
				case Sync, Added, Updated:
					if old, exists, err := clientState.Get(d.Object); err == nil && exists {
						if err := clientState.Update(d.Object); err != nil {
							return err
						}
						h.OnUpdate(old, d.Object)
					} else {
						if err := clientState.Add(d.Object); err != nil {
							return err
						}
						h.OnAdd(d.Object)
					}
				case Deleted:
					if err := clientState.Delete(d.Object); err != nil {
						return err
					}
					h.OnDelete(d.Object)
				}
			}
			return nil
		},
	}
	return clientState, New(cfg)
}

// NewIndexerInformer returns a Indexer and a controller for populating the index
// while also providing event notifications. You should only used the returned
// Index for Get/List operations; Add/Modify/Deletes will cause the event
// notifications to be faulty.
//
// Parameters:
//  * lw is list and watch functions for the source of the resource you want to
//    be informed of.
//  * objType is an object of the type that you expect to receive.
//  * resyncPeriod: if non-zero, will re-list this often (you will get OnUpdate
//    calls, even if nothing changed). Otherwise, re-list will be delayed as
//    long as possible (until the upstream source closes the watch or times out,
//    or you stop the controller).
//  * h is the object you want notifications sent to.
//
func NewIndexerInformer(
	lw ListerWatcher,
	objType runtime.Object,
	resyncPeriod time.Duration,
	h ResourceEventHandler,
	indexers Indexers,
) (Indexer, Controller) {
	// This will hold the client state, as we know it.
	clientState := NewIndexer(DeletionHandlingMetaNamespaceKeyFunc, indexers)

	// This will hold incoming changes. Note how we pass clientState in as a
	// KeyLister, that way resync operations will result in the correct set
	// of update/delete deltas.
	fifo := NewDeltaFIFO(MetaNamespaceKeyFunc, nil, clientState)

	cfg := &Config{
		Queue:            fifo,
		ListerWatcher:    lw,
		ObjectType:       objType,
		FullResyncPeriod: resyncPeriod,
		RetryOnError:     false,

		Process: func(obj interface{}) error {
			// from oldest to newest
			for _, d := range obj.(Deltas) {
				switch d.Type {
				case Sync, Added, Updated:
					if old, exists, err := clientState.Get(d.Object); err == nil && exists {
						if err := clientState.Update(d.Object); err != nil {
							return err
						}
						h.OnUpdate(old, d.Object)
					} else {
						if err := clientState.Add(d.Object); err != nil {
							return err
						}
						h.OnAdd(d.Object)
					}
				case Deleted:
					if err := clientState.Delete(d.Object); err != nil {
						return err
					}
					h.OnDelete(d.Object)
				}
			}
			return nil
		},
	}
	return clientState, New(cfg)
}
<commit_msg>Added indexer description for the API of NewIndexerInformer.<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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 cache

import (
	"sync"
	"time"

	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/util/clock"
)

// Config contains all the settings for a Controller.
type Config struct {
	// The queue for your objects; either a FIFO or
	// a DeltaFIFO. Your Process() function should accept
	// the output of this Queue's Pop() method.
	Queue

	// Something that can list and watch your objects.
	ListerWatcher

	// Something that can process your objects.
	Process ProcessFunc

	// The type of your objects.
	ObjectType runtime.Object

	// Reprocess everything at least this often.
	// Note that if it takes longer for you to clear the queue than this
	// period, you will end up processing items in the order determined
	// by FIFO.Replace(). Currently, this is random. If this is a
	// problem, we can change that replacement policy to append new
	// things to the end of the queue instead of replacing the entire
	// queue.
	FullResyncPeriod time.Duration

	// ShouldResync, if specified, is invoked when the controller's reflector determines the next
	// periodic sync should occur. If this returns true, it means the reflector should proceed with
	// the resync.
	ShouldResync ShouldResyncFunc

	// If true, when Process() returns an error, re-enqueue the object.
	// TODO: add interface to let you inject a delay/backoff or drop
	//       the object completely if desired. Pass the object in
	//       question to this interface as a parameter.
	RetryOnError bool
}

// ShouldResyncFunc is a type of function that indicates if a reflector should perform a
// resync or not. It can be used by a shared informer to support multiple event handlers with custom
// resync periods.
type ShouldResyncFunc func() bool

// ProcessFunc processes a single object.
type ProcessFunc func(obj interface{}) error

// Controller is a generic controller framework.
type controller struct {
	config         Config
	reflector      *Reflector
	reflectorMutex sync.RWMutex
	clock          clock.Clock
}

type Controller interface {
	Run(stopCh <-chan struct{})
	HasSynced() bool
	LastSyncResourceVersion() string
}

// New makes a new Controller from the given Config.
func New(c *Config) Controller {
	ctlr := &controller{
		config: *c,
		clock:  &clock.RealClock{},
	}
	return ctlr
}

// Run begins processing items, and will continue until a value is sent down stopCh.
// It's an error to call Run more than once.
// Run blocks; call via go.
func (c *controller) Run(stopCh <-chan struct{}) {
	defer utilruntime.HandleCrash()
	go func() {
		<-stopCh
		c.config.Queue.Close()
	}()
	r := NewReflector(
		c.config.ListerWatcher,
		c.config.ObjectType,
		c.config.Queue,
		c.config.FullResyncPeriod,
	)
	r.ShouldResync = c.config.ShouldResync
	r.clock = c.clock

	c.reflectorMutex.Lock()
	c.reflector = r
	c.reflectorMutex.Unlock()

	r.RunUntil(stopCh)

	wait.Until(c.processLoop, time.Second, stopCh)
}

// Returns true once this controller has completed an initial resource listing
func (c *controller) HasSynced() bool {
	return c.config.Queue.HasSynced()
}

func (c *controller) LastSyncResourceVersion() string {
	if c.reflector == nil {
		return ""
	}
	return c.reflector.LastSyncResourceVersion()
}

// processLoop drains the work queue.
// TODO: Consider doing the processing in parallel. This will require a little thought
// to make sure that we don't end up processing the same object multiple times
// concurrently.
//
// TODO: Plumb through the stopCh here (and down to the queue) so that this can
// actually exit when the controller is stopped. Or just give up on this stuff
// ever being stoppable. Converting this whole package to use Context would
// also be helpful.
func (c *controller) processLoop() {
	for {
		obj, err := c.config.Queue.Pop(PopProcessFunc(c.config.Process))
		if err != nil {
			if err == FIFOClosedError {
				return
			}
			if c.config.RetryOnError {
				// This is the safe way to re-enqueue.
				c.config.Queue.AddIfNotPresent(obj)
			}
		}
	}
}

// ResourceEventHandler can handle notifications for events that happen to a
// resource. The events are informational only, so you can't return an
// error.
//  * OnAdd is called when an object is added.
//  * OnUpdate is called when an object is modified. Note that oldObj is the
//      last known state of the object-- it is possible that several changes
//      were combined together, so you can't use this to see every single
//      change. OnUpdate is also called when a re-list happens, and it will
//      get called even if nothing changed. This is useful for periodically
//      evaluating or syncing something.
//  * OnDelete will get the final state of the item if it is known, otherwise
//      it will get an object of type DeletedFinalStateUnknown. This can
//      happen if the watch is closed and misses the delete event and we don't
//      notice the deletion until the subsequent re-list.
type ResourceEventHandler interface {
	OnAdd(obj interface{})
	OnUpdate(oldObj, newObj interface{})
	OnDelete(obj interface{})
}

// ResourceEventHandlerFuncs is an adaptor to let you easily specify as many or
// as few of the notification functions as you want while still implementing
// ResourceEventHandler.
type ResourceEventHandlerFuncs struct {
	AddFunc    func(obj interface{})
	UpdateFunc func(oldObj, newObj interface{})
	DeleteFunc func(obj interface{})
}

// OnAdd calls AddFunc if it's not nil.
func (r ResourceEventHandlerFuncs) OnAdd(obj interface{}) {
	if r.AddFunc != nil {
		r.AddFunc(obj)
	}
}

// OnUpdate calls UpdateFunc if it's not nil.
func (r ResourceEventHandlerFuncs) OnUpdate(oldObj, newObj interface{}) {
	if r.UpdateFunc != nil {
		r.UpdateFunc(oldObj, newObj)
	}
}

// OnDelete calls DeleteFunc if it's not nil.
func (r ResourceEventHandlerFuncs) OnDelete(obj interface{}) {
	if r.DeleteFunc != nil {
		r.DeleteFunc(obj)
	}
}

// DeletionHandlingMetaNamespaceKeyFunc checks for
// DeletedFinalStateUnknown objects before calling
// MetaNamespaceKeyFunc.
func DeletionHandlingMetaNamespaceKeyFunc(obj interface{}) (string, error) {
	if d, ok := obj.(DeletedFinalStateUnknown); ok {
		return d.Key, nil
	}
	return MetaNamespaceKeyFunc(obj)
}

// NewInformer returns a Store and a controller for populating the store
// while also providing event notifications. You should only used the returned
// Store for Get/List operations; Add/Modify/Deletes will cause the event
// notifications to be faulty.
//
// Parameters:
//  * lw is list and watch functions for the source of the resource you want to
//    be informed of.
//  * objType is an object of the type that you expect to receive.
//  * resyncPeriod: if non-zero, will re-list this often (you will get OnUpdate
//    calls, even if nothing changed). Otherwise, re-list will be delayed as
//    long as possible (until the upstream source closes the watch or times out,
//    or you stop the controller).
//  * h is the object you want notifications sent to.
//
func NewInformer(
	lw ListerWatcher,
	objType runtime.Object,
	resyncPeriod time.Duration,
	h ResourceEventHandler,
) (Store, Controller) {
	// This will hold the client state, as we know it.
	clientState := NewStore(DeletionHandlingMetaNamespaceKeyFunc)

	// This will hold incoming changes. Note how we pass clientState in as a
	// KeyLister, that way resync operations will result in the correct set
	// of update/delete deltas.
	fifo := NewDeltaFIFO(MetaNamespaceKeyFunc, nil, clientState)

	cfg := &Config{
		Queue:            fifo,
		ListerWatcher:    lw,
		ObjectType:       objType,
		FullResyncPeriod: resyncPeriod,
		RetryOnError:     false,

		Process: func(obj interface{}) error {
			// from oldest to newest
			for _, d := range obj.(Deltas) {
				switch d.Type {
				case Sync, Added, Updated:
					if old, exists, err := clientState.Get(d.Object); err == nil && exists {
						if err := clientState.Update(d.Object); err != nil {
							return err
						}
						h.OnUpdate(old, d.Object)
					} else {
						if err := clientState.Add(d.Object); err != nil {
							return err
						}
						h.OnAdd(d.Object)
					}
				case Deleted:
					if err := clientState.Delete(d.Object); err != nil {
						return err
					}
					h.OnDelete(d.Object)
				}
			}
			return nil
		},
	}
	return clientState, New(cfg)
}

// NewIndexerInformer returns a Indexer and a controller for populating the index
// while also providing event notifications. You should only used the returned
// Index for Get/List operations; Add/Modify/Deletes will cause the event
// notifications to be faulty.
//
// Parameters:
//  * lw is list and watch functions for the source of the resource you want to
//    be informed of.
//  * objType is an object of the type that you expect to receive.
//  * resyncPeriod: if non-zero, will re-list this often (you will get OnUpdate
//    calls, even if nothing changed). Otherwise, re-list will be delayed as
//    long as possible (until the upstream source closes the watch or times out,
//    or you stop the controller).
//  * h is the object you want notifications sent to.
//  * indexers is the indexer for the received object type.
//
func NewIndexerInformer(
	lw ListerWatcher,
	objType runtime.Object,
	resyncPeriod time.Duration,
	h ResourceEventHandler,
	indexers Indexers,
) (Indexer, Controller) {
	// This will hold the client state, as we know it.
	clientState := NewIndexer(DeletionHandlingMetaNamespaceKeyFunc, indexers)

	// This will hold incoming changes. Note how we pass clientState in as a
	// KeyLister, that way resync operations will result in the correct set
	// of update/delete deltas.
	fifo := NewDeltaFIFO(MetaNamespaceKeyFunc, nil, clientState)

	cfg := &Config{
		Queue:            fifo,
		ListerWatcher:    lw,
		ObjectType:       objType,
		FullResyncPeriod: resyncPeriod,
		RetryOnError:     false,

		Process: func(obj interface{}) error {
			// from oldest to newest
			for _, d := range obj.(Deltas) {
				switch d.Type {
				case Sync, Added, Updated:
					if old, exists, err := clientState.Get(d.Object); err == nil && exists {
						if err := clientState.Update(d.Object); err != nil {
							return err
						}
						h.OnUpdate(old, d.Object)
					} else {
						if err := clientState.Add(d.Object); err != nil {
							return err
						}
						h.OnAdd(d.Object)
					}
				case Deleted:
					if err := clientState.Delete(d.Object); err != nil {
						return err
					}
					h.OnDelete(d.Object)
				}
			}
			return nil
		},
	}
	return clientState, New(cfg)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Workiva, LLC

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 queue

import "errors"

var (
	// ErrDisposed is returned when an operation is performed on a disposed
	// queue.
	ErrDisposed = errors.New(`Queue has been disposed.`)

	// ErrTimeout is returned when an applicable queue operation times out.
	ErrTimeout = errors.New(`Poll timed out.`)
)
<commit_msg>Make error messages more idiomatic<commit_after>/*
Copyright 2014 Workiva, LLC

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 queue

import "errors"

var (
	// ErrDisposed is returned when an operation is performed on a disposed
	// queue.
	ErrDisposed = errors.New(`queue: disposed`)

	// ErrTimeout is returned when an applicable queue operation times out.
	ErrTimeout = errors.New(`queue: poll timed out`)
)
<|endoftext|>


<commit_before>package queue

import (
	"github.com/nathan-osman/go-cannon/email"
	"github.com/nathan-osman/go-cannon/util"

	"io/ioutil"
	"log"
	"path"
	"time"
)

// Mail queue managing the sending of emails to hosts.
type Queue struct {
	newEmail *util.NonBlockingChan
	stop     chan bool
}

// Attempt to load all emails from the specified directory and send them on the
// specified channel..
func loadEmails(directory string, newEmail *util.NonBlockingChan) error {

	// Enumerate the files in the directory
	files, err := ioutil.ReadDir(directory)
	if err != nil {
		return err
	}

	// Attempt to load each file and ignore ones that fail
	for _, f := range files {
		if e, err := email.LoadEmail(path.Join(directory, f.Name())); err == nil {
			newEmail.Send <- e
		}
	}

	return nil
}

// Create a new mail queue.
func NewQueue(directory string) (*Queue, error) {

	// Create the two channels the queue will need
	q := &Queue{
		newEmail: util.NewNonBlockingChan(),
		stop:     make(chan bool),
	}

	// Load any emails in the storage directory
	if err := loadEmails(directory, q.newEmail); err != nil {
		return nil, err
	}

	// Start a goroutine to manage the lifecycle of the queue
	go func() {

		// Close the stop channel when the goroutine exits
		defer close(q.stop)

		// Create a map of hosts and a ticker for freeing up unused hosts
		var (
			hosts  = make(map[string]*Host)
			ticker = time.NewTicker(5 * time.Minute)
		)

		// Stop the ticker when the goroutine exits
		defer ticker.Stop()

		// Main "loop" of the queue
	loop:
		for {
			select {
			case i := <-q.newEmail.Recv:

				log.Println("received new email in mail queue")

				// Convert to an Email pointer
				e := i.(*email.Email)

				// Create the specified host if it doesn't exist
				if _, ok := hosts[e.Host]; !ok {
					hosts[e.Host] = NewHost(e.Host)
				}

				// Deliver the message to the host
				hosts[e.Host].Deliver(e)

			case <-ticker.C:

				// Loop through all of the hosts and remove ones that have been
				// idle for longer than 5 minutes and stops them
				for h := range hosts {
					if hosts[h].Idle() > 5*time.Minute {
						hosts[h].Stop()
						delete(hosts, h)
					}
				}

			case <-q.stop:
				break loop
			}
		}

		log.Println("shutting down mail queue")

		// Stop all host queues
		for h := range hosts {
			hosts[h].Stop()
		}
	}()

	return q, nil
}

// Deliver the provided email.
func (q *Queue) Deliver(e *email.Email) {
	q.newEmail.Send <- e
}

// Stop all active host queues.
func (q *Queue) Stop() {

	// Send on the channel to stop it and wait for it to be closed
	q.stop <- true
	<-q.stop
}
<commit_msg>Added check for existence of directory when loading emails.<commit_after>package queue

import (
	"github.com/nathan-osman/go-cannon/email"
	"github.com/nathan-osman/go-cannon/util"

	"io/ioutil"
	"log"
	"os"
	"path"
	"time"
)

// Mail queue managing the sending of emails to hosts.
type Queue struct {
	newEmail *util.NonBlockingChan
	stop     chan bool
}

// Attempt to load all emails from the specified directory and send them on the
// specified channel.
func loadEmails(directory string, newEmail *util.NonBlockingChan) error {

	// If the directory does not exist, quit
	if _, err := os.Stat(directory); os.IsNotExist(err) {
		return nil
	}

	// Enumerate the files in the directory
	files, err := ioutil.ReadDir(directory)
	if err != nil {
		return err
	}

	// Attempt to load each file and ignore ones that fail
	for _, f := range files {
		if e, err := email.LoadEmail(path.Join(directory, f.Name())); err == nil {
			newEmail.Send <- e
		}
	}

	return nil
}

// Create a new mail queue.
func NewQueue(directory string) (*Queue, error) {

	// Create the two channels the queue will need
	q := &Queue{
		newEmail: util.NewNonBlockingChan(),
		stop:     make(chan bool),
	}

	// Load any emails in the storage directory
	if err := loadEmails(directory, q.newEmail); err != nil {
		return nil, err
	}

	// Start a goroutine to manage the lifecycle of the queue
	go func() {

		// Close the stop channel when the goroutine exits
		defer close(q.stop)

		// Create a map of hosts and a ticker for freeing up unused hosts
		var (
			hosts  = make(map[string]*Host)
			ticker = time.NewTicker(5 * time.Minute)
		)

		// Stop the ticker when the goroutine exits
		defer ticker.Stop()

		// Main "loop" of the queue
	loop:
		for {
			select {
			case i := <-q.newEmail.Recv:

				log.Println("received new email in mail queue")

				// Convert to an Email pointer
				e := i.(*email.Email)

				// Create the specified host if it doesn't exist
				if _, ok := hosts[e.Host]; !ok {
					hosts[e.Host] = NewHost(e.Host)
				}

				// Deliver the message to the host
				hosts[e.Host].Deliver(e)

			case <-ticker.C:

				// Loop through all of the hosts and remove ones that have been
				// idle for longer than 5 minutes and stops them
				for h := range hosts {
					if hosts[h].Idle() > 5*time.Minute {
						hosts[h].Stop()
						delete(hosts, h)
					}
				}

			case <-q.stop:
				break loop
			}
		}

		log.Println("shutting down mail queue")

		// Stop all host queues
		for h := range hosts {
			hosts[h].Stop()
		}
	}()

	return q, nil
}

// Deliver the provided email.
func (q *Queue) Deliver(e *email.Email) {
	q.newEmail.Send <- e
}

// Stop all active host queues.
func (q *Queue) Stop() {

	// Send on the channel to stop it and wait for it to be closed
	q.stop <- true
	<-q.stop
}
<|endoftext|>


<commit_before>/* Copyright 2017 The Bazel Authors. All rights reserved.

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.
*/

// reproducible_binary_test checks that a given binary DOES NOT contain
// strings that match the GOROOT, the current user's name, or the
// current user's home directory.
package main

import (
	"bytes"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"regexp"
	"testing"

	"github.com/bazelbuild/rules_go/go/tools/bazel"
)

var allStrings [][]byte
var currentUser *user.User

func TestMain(m *testing.M) {
	log.SetFlags(0)
	if len(os.Args) != 2 {
		log.Fatalf("usage: %s <binary>\n", os.Args[0])
	}

	binary, err := bazel.Runfile(os.Args[1])
	if err != nil {
		log.Fatalf("Could not find runfile %s: %q", os.Args[1], err)
	}
	binaryData, err := ioutil.ReadFile(binary)
	if err != nil {
		log.Fatal(err)
	}
	stringRex := regexp.MustCompile(`[[:graph:]]{3,}`)
	allStrings = stringRex.FindAll(binaryData, -1)

	currentUser, err = user.Current()
	if err != nil {
		currentUser = nil
	}

	os.Exit(m.Run())
}

// TestStandardPath checks that source paths from the standard library
// are trimmed. We just check one known source file.
func TestStandardPath(t *testing.T) {
	want := []byte("GOROOT/src/fmt/format.go")
	for _, s := range allStrings {
		if bytes.HasSuffix(s, []byte("fmt/format.go")) && !bytes.Equal(s, want) {
			t.Fatalf("got %q; want %q", s, want)
		}
	}
}

// TestSandbox checks that the bazel-sandbox path does not appear in strings
// from the binary.
func TestSandboxPath(t *testing.T) {
	for _, s := range allStrings {
		if bytes.Contains(s, []byte("bazel-sandbox")) {
			t.Errorf("binary contains bazel sandbox path: %s", s)
		}
	}
}

// TestUserNameAndHome checks the user name and home directory do not
// appear in strings from the binary.
func TestUserNameAndHome(t *testing.T) {
	if currentUser == nil || len(currentUser.Username) < 4 {
		t.Skip()
	}
	for _, s := range allStrings {
		if currentUser.Username != "" && bytes.Contains(s, []byte(currentUser.Username)) {
			t.Errorf("binary contains username %q in string %q", currentUser.Username, s)
			continue
		}
		if currentUser.HomeDir != "" && bytes.Contains(s, []byte(currentUser.HomeDir)) {
			t.Errorf("binary contains home dir %q in string %q", currentUser.HomeDir, s)
		}
	}
}
<commit_msg>Skip part of //tests/reproducibility:go_default_test in BuildKite (#1902)<commit_after>/* Copyright 2017 The Bazel Authors. All rights reserved.

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.
*/

// reproducible_binary_test checks that a given binary DOES NOT contain
// strings that match the GOROOT, the current user's name, or the
// current user's home directory.
package main

import (
	"bytes"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"regexp"
	"strings"
	"testing"

	"github.com/bazelbuild/rules_go/go/tools/bazel"
)

var allStrings [][]byte
var currentUser *user.User

func TestMain(m *testing.M) {
	log.SetFlags(0)
	if len(os.Args) != 2 {
		log.Fatalf("usage: %s <binary>\n", os.Args[0])
	}

	binary, err := bazel.Runfile(os.Args[1])
	if err != nil {
		log.Fatalf("Could not find runfile %s: %q", os.Args[1], err)
	}
	binaryData, err := ioutil.ReadFile(binary)
	if err != nil {
		log.Fatal(err)
	}
	stringRex := regexp.MustCompile(`[[:graph:]]{3,}`)
	allStrings = stringRex.FindAll(binaryData, -1)

	currentUser, err = user.Current()
	if err != nil {
		currentUser = nil
	}

	os.Exit(m.Run())
}

// TestStandardPath checks that source paths from the standard library
// are trimmed. We just check one known source file.
func TestStandardPath(t *testing.T) {
	want := []byte("GOROOT/src/fmt/format.go")
	for _, s := range allStrings {
		if bytes.HasSuffix(s, []byte("fmt/format.go")) && !bytes.Equal(s, want) {
			t.Fatalf("got %q; want %q", s, want)
		}
	}
}

// TestSandbox checks that the bazel-sandbox path does not appear in strings
// from the binary.
func TestSandboxPath(t *testing.T) {
	for _, s := range allStrings {
		if bytes.Contains(s, []byte("bazel-sandbox")) {
			t.Errorf("binary contains bazel sandbox path: %s", s)
		}
	}
}

// TestUserNameAndHome checks the user name and home directory do not
// appear in strings from the binary.
func TestUserNameAndHome(t *testing.T) {
	if currentUser == nil || len(currentUser.Username) < 4 || strings.Contains(currentUser.Username, "bazel") {
		t.Skip()
	}
	for _, s := range allStrings {
		if currentUser.Username != "" && bytes.Contains(s, []byte(currentUser.Username)) {
			t.Errorf("binary contains username %q in string %q", currentUser.Username, s)
			continue
		}
		if currentUser.HomeDir != "" && bytes.Contains(s, []byte(currentUser.HomeDir)) {
			t.Errorf("binary contains home dir %q in string %q", currentUser.HomeDir, s)
		}
	}
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/state/multiwatcher"
	"launchpad.net/juju-core/state/watcher"
	"reflect"
)

// allWatcherStateBacking implements allWatcherBacking by
// fetching entities from the State.
type allWatcherStateBacking struct {
	st *State
	// collections
	collectionByName map[string]allWatcherStateCollection
	collectionByKind map[string]allWatcherStateCollection
}

// allWatcherStateCollection holds information about a
// collection watched by an allWatcher and the
// type of value we use to store entity information
// for that collection.
type allWatcherStateCollection struct {
	*mgo.Collection
	// infoSliceType stores the type of a slice of the info type
	// that we use for this collection.  In Go 1.1 we can change
	// this to use the type itself, as we'll have reflect.SliceOf.
	infoSliceType reflect.Type
	// idOf returns the id of the given info.
	idOf func(info params.EntityInfo) interface{}
}

func newAllWatcherStateBacking(st *State) multiwatcher.Backing {
	b := &allWatcherStateBacking{
		st:               st,
		collectionByName: make(map[string]allWatcherStateCollection),
		collectionByKind: make(map[string]allWatcherStateCollection),
	}
	collections := []allWatcherStateCollection{{
		Collection:    st.machines,
		infoSliceType: reflect.TypeOf([]params.MachineInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.MachineInfo).Id
		},
	}, {
		Collection:    st.units,
		infoSliceType: reflect.TypeOf([]params.UnitInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.UnitInfo).Name
		},
	}, {
		Collection:    st.services,
		infoSliceType: reflect.TypeOf([]params.ServiceInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.ServiceInfo).Name
		},
	}, {
		Collection:    st.relations,
		infoSliceType: reflect.TypeOf([]params.RelationInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.RelationInfo).Key
		},
	}, {
		Collection:    st.annotations,
		infoSliceType: reflect.TypeOf([]params.AnnotationInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.AnnotationInfo).GlobalKey
		},
	}}
	// Populate the collection maps from the above set of collections.
	for _, c := range collections {
		// Create a new instance of the info type so we can
		// find out its kind.
		info := reflect.New(c.infoSliceType.Elem()).Interface().(params.EntityInfo)
		kind := info.EntityKind()
		if _, ok := b.collectionByKind[kind]; ok {
			panic(fmt.Errorf("duplicate collection kind %q", kind))
		}
		b.collectionByKind[kind] = c
		if _, ok := b.collectionByName[c.Name]; ok {
			panic(fmt.Errorf("duplicate collection name %q", kind))
		}
		b.collectionByName[c.Name] = c
	}
	return b
}

// Watch watches all the collections.
func (b *allWatcherStateBacking) Watch(in chan<- watcher.Change) {
	for _, c := range b.collectionByName {
		b.st.watcher.WatchCollection(c.Name, in)
	}
}

// Unwatch unwatches all the collections.
func (b *allWatcherStateBacking) Unwatch(in chan<- watcher.Change) {
	for _, c := range b.collectionByName {
		b.st.watcher.UnwatchCollection(c.Name, in)
	}
}

// GetAll fetches all items that we want to watch from the state.
func (b *allWatcherStateBacking) GetAll(all *multiwatcher.Store) error {
	// TODO(rog) fetch collections concurrently?
	for _, c := range b.collectionByName {
		infoSlicePtr := reflect.New(c.infoSliceType).Interface()
		if err := c.Find(nil).All(infoSlicePtr); err != nil {
			return fmt.Errorf("cannot get all %s: %v", c.Name, err)
		}
		infos := reflect.ValueOf(infoSlicePtr).Elem()
		for i := 0; i < infos.Len(); i++ {
			info := infos.Index(i).Addr().Interface().(params.EntityInfo)
			all.Update(b.IdForInfo(info), info)
		}
	}
	return nil
}

// entityId holds the mongo identifier of an entity.
type entityId struct {
	collection string
	id         interface{}
}

// changed updates the allWatcher's idea of the current state
// in response to the given change.
func (b *allWatcherStateBacking) Changed(all *multiwatcher.Store, change watcher.Change) error {
	id := entityId{
		collection: change.C,
		id:         change.Id,
	}
	// TODO(rog) investigate ways that this can be made more efficient
	// than simply fetching each entity in turn.
	info, err := b.fetch(id)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}
	log.Infof("updated %+v to %#v", id, info)
	all.Update(id, info)
	return nil
}

func (b *allWatcherStateBacking) fetch(id entityId) (params.EntityInfo, error) {
	c, ok := b.collectionByName[id.collection]
	if !ok {
		panic(fmt.Errorf("unknown collection %q in fetch request", id.collection))
	}
	info := reflect.New(c.infoSliceType.Elem()).Interface().(params.EntityInfo)
	if err := c.FindId(id.id).One(info); err != nil {
		return nil, err
	}
	return info, nil
}

// idForInfo returns the info id of the given entity document.
func (b *allWatcherStateBacking) IdForInfo(info params.EntityInfo) multiwatcher.InfoId {
	c, ok := b.collectionByKind[info.EntityKind()]
	if !ok {
		panic(fmt.Errorf("entity with unknown kind %q", info.EntityKind()))
	}
	return entityId{
		collection: c.Name,
		id:         c.idOf(info),
	}
}
<commit_msg>state: remove log<commit_after>package state

import (
	"fmt"
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/state/multiwatcher"
	"launchpad.net/juju-core/state/watcher"
	"reflect"
)

// allWatcherStateBacking implements allWatcherBacking by
// fetching entities from the State.
type allWatcherStateBacking struct {
	st *State
	// collections
	collectionByName map[string]allWatcherStateCollection
	collectionByKind map[string]allWatcherStateCollection
}

// allWatcherStateCollection holds information about a
// collection watched by an allWatcher and the
// type of value we use to store entity information
// for that collection.
type allWatcherStateCollection struct {
	*mgo.Collection
	// infoSliceType stores the type of a slice of the info type
	// that we use for this collection.  In Go 1.1 we can change
	// this to use the type itself, as we'll have reflect.SliceOf.
	infoSliceType reflect.Type
	// idOf returns the id of the given info.
	idOf func(info params.EntityInfo) interface{}
}

func newAllWatcherStateBacking(st *State) multiwatcher.Backing {
	b := &allWatcherStateBacking{
		st:               st,
		collectionByName: make(map[string]allWatcherStateCollection),
		collectionByKind: make(map[string]allWatcherStateCollection),
	}
	collections := []allWatcherStateCollection{{
		Collection:    st.machines,
		infoSliceType: reflect.TypeOf([]params.MachineInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.MachineInfo).Id
		},
	}, {
		Collection:    st.units,
		infoSliceType: reflect.TypeOf([]params.UnitInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.UnitInfo).Name
		},
	}, {
		Collection:    st.services,
		infoSliceType: reflect.TypeOf([]params.ServiceInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.ServiceInfo).Name
		},
	}, {
		Collection:    st.relations,
		infoSliceType: reflect.TypeOf([]params.RelationInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.RelationInfo).Key
		},
	}, {
		Collection:    st.annotations,
		infoSliceType: reflect.TypeOf([]params.AnnotationInfo(nil)),
		idOf: func(info params.EntityInfo) interface{} {
			return info.(*params.AnnotationInfo).GlobalKey
		},
	}}
	// Populate the collection maps from the above set of collections.
	for _, c := range collections {
		// Create a new instance of the info type so we can
		// find out its kind.
		info := reflect.New(c.infoSliceType.Elem()).Interface().(params.EntityInfo)
		kind := info.EntityKind()
		if _, ok := b.collectionByKind[kind]; ok {
			panic(fmt.Errorf("duplicate collection kind %q", kind))
		}
		b.collectionByKind[kind] = c
		if _, ok := b.collectionByName[c.Name]; ok {
			panic(fmt.Errorf("duplicate collection name %q", kind))
		}
		b.collectionByName[c.Name] = c
	}
	return b
}

// Watch watches all the collections.
func (b *allWatcherStateBacking) Watch(in chan<- watcher.Change) {
	for _, c := range b.collectionByName {
		b.st.watcher.WatchCollection(c.Name, in)
	}
}

// Unwatch unwatches all the collections.
func (b *allWatcherStateBacking) Unwatch(in chan<- watcher.Change) {
	for _, c := range b.collectionByName {
		b.st.watcher.UnwatchCollection(c.Name, in)
	}
}

// GetAll fetches all items that we want to watch from the state.
func (b *allWatcherStateBacking) GetAll(all *multiwatcher.Store) error {
	// TODO(rog) fetch collections concurrently?
	for _, c := range b.collectionByName {
		infoSlicePtr := reflect.New(c.infoSliceType).Interface()
		if err := c.Find(nil).All(infoSlicePtr); err != nil {
			return fmt.Errorf("cannot get all %s: %v", c.Name, err)
		}
		infos := reflect.ValueOf(infoSlicePtr).Elem()
		for i := 0; i < infos.Len(); i++ {
			info := infos.Index(i).Addr().Interface().(params.EntityInfo)
			all.Update(b.IdForInfo(info), info)
		}
	}
	return nil
}

// entityId holds the mongo identifier of an entity.
type entityId struct {
	collection string
	id         interface{}
}

// changed updates the allWatcher's idea of the current state
// in response to the given change.
func (b *allWatcherStateBacking) Changed(all *multiwatcher.Store, change watcher.Change) error {
	id := entityId{
		collection: change.C,
		id:         change.Id,
	}
	// TODO(rog) investigate ways that this can be made more efficient
	// than simply fetching each entity in turn.
	info, err := b.fetch(id)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}
	all.Update(id, info)
	return nil
}

func (b *allWatcherStateBacking) fetch(id entityId) (params.EntityInfo, error) {
	c, ok := b.collectionByName[id.collection]
	if !ok {
		panic(fmt.Errorf("unknown collection %q in fetch request", id.collection))
	}
	info := reflect.New(c.infoSliceType.Elem()).Interface().(params.EntityInfo)
	if err := c.FindId(id.id).One(info); err != nil {
		return nil, err
	}
	return info, nil
}

// idForInfo returns the info id of the given entity document.
func (b *allWatcherStateBacking) IdForInfo(info params.EntityInfo) multiwatcher.InfoId {
	c, ok := b.collectionByKind[info.EntityKind()]
	if !ok {
		panic(fmt.Errorf("entity with unknown kind %q", info.EntityKind()))
	}
	return entityId{
		collection: c.Name,
		id:         c.idOf(info),
	}
}
<|endoftext|>


<commit_before>package minhashlsh

import (
	"math/rand"
	"testing"
)

func randomSignature(size int, seed int64) Signature {
	r := rand.New(rand.NewSource(seed))
	sig := make(Signature, size)
	for i := range sig {
		sig[i] = uint64(r.Int63())
	}
	return sig
}

func Test_HashKeyFunc16(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(2)
	hashKey := f(sig)
	if len(hashKey) != 2*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_HashKeyFunc64(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(8)
	hashKey := f(sig)
	if len(hashKey) != 8*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_MinhashLSH(t *testing.T) {
	f := NewMinhashLSH16(256, 0.6)
	// sig1 is different from sig2 and sig3
	// sig2 and sig3 are identical
	sig1 := randomSignature(256, 1)
	sig2 := randomSignature(256, 2)
	sig3 := randomSignature(256, 2)

	f.Add("sig1", sig1)
	f.Add("sig2", sig2)
	f.Add("sig3", sig3)
	f.Index()
	// sig1 should be in its own bucket
	// sig2 and sig3 are in another bucket
	for i := range f.hashTables {
		if len(f.hashTables[i]) != 2 {
			t.Fatal(f.hashTables[i])
		}
	}

	found := 0
	for _, key := range f.Query(sig3) {
		if key == "sig3" || key == "sig2" {
			found++
		}
	}
	if found != 2 {
		t.Fatal("unable to retrieve inserted keys")
	}
}

func Test_MinhashLSH2(t *testing.T) {
	minhashLsh := NewMinhashLSH16(256, 0.5)
	seed := 1
	numHash := 256
	mh := NewMinhash(seed, numHash)
	words := []string{"hello", "world", "minhash", "one", "two", "three", "four",
		"five", "six", "seven", "eight", "nine", "ten"}
	for _, word := range words {
		mh.Push([]byte(word))
	}
	sig1 := mh.Signature()
	minhashLsh.Add("s1", sig1)
	minhashLsh.Index()
	k, l := minhashLsh.Params()
	t.Logf("Minhash LSH params: k = %d, l = %d", k, l)

	mh = NewMinhash(seed, numHash)
	words = []string{"one", "two", "three", "four",
		"five", "six", "seven", "eight", "nine", "ten"}
	for _, word := range words {
		mh.Push([]byte(word))
	}
	sig2 := mh.Signature()

	results := minhashLsh.Query(sig2)
	t.Log(results)
	if len(results) < 1 {
		t.Fail()
	}
}
<commit_msg>fix test<commit_after>package minhashlsh

import (
	"math/rand"
	"testing"
)

func randomSignature(size int, seed int64) Signature {
	r := rand.New(rand.NewSource(seed))
	sig := make(Signature, size)
	for i := range sig {
		sig[i] = uint64(r.Int63())
	}
	return sig
}

func Test_HashKeyFunc16(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(2)
	hashKey := f(sig)
	if len(hashKey) != 2*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_HashKeyFunc64(t *testing.T) {
	sig := randomSignature(2, 1)
	f := hashKeyFuncGen(8)
	hashKey := f(sig)
	if len(hashKey) != 8*2 {
		t.Fatal(len(hashKey))
	}
}

func Test_MinhashLSH(t *testing.T) {
	f := NewMinhashLSH16(256, 0.6)
	// sig1 is different from sig2 and sig3
	// sig2 and sig3 are identical
	sig1 := randomSignature(256, 1)
	sig2 := randomSignature(256, 2)
	sig3 := randomSignature(256, 2)

	f.Add("sig1", sig1)
	f.Add("sig2", sig2)
	f.Add("sig3", sig3)
	f.Index()
	// sig1 should be in its own bucket
	// sig2 and sig3 are in another bucket
	for i := range f.hashTables {
		if len(f.hashTables[i]) != 2 {
			t.Fatal(f.hashTables[i])
		}
	}

	found := 0
	for _, key := range f.Query(sig3) {
		if key == "sig3" || key == "sig2" {
			found++
		}
	}
	if found != 2 {
		t.Fatal("unable to retrieve inserted keys")
	}
}

func Test_MinhashLSH2(t *testing.T) {
	minhashLsh := NewMinhashLSH16(256, 0.5)
	seed := int64(1)
	numHash := 256
	mh := NewMinhash(seed, numHash)
	words := []string{"hello", "world", "minhash", "one", "two", "three", "four",
		"five", "six", "seven", "eight", "nine", "ten"}
	for _, word := range words {
		mh.Push([]byte(word))
	}
	sig1 := mh.Signature()
	minhashLsh.Add("s1", sig1)
	minhashLsh.Index()
	k, l := minhashLsh.Params()
	t.Logf("Minhash LSH params: k = %d, l = %d", k, l)

	mh = NewMinhash(seed, numHash)
	words = []string{"one", "two", "three", "four",
		"five", "six", "seven", "eight", "nine", "ten"}
	for _, word := range words {
		mh.Push([]byte(word))
	}
	sig2 := mh.Signature()

	results := minhashLsh.Query(sig2)
	t.Log(results)
	if len(results) < 1 {
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"

	markdown "github.com/shurcooL/github_flavored_markdown"
)

type PageData struct {
	Title   string
	Content template.HTML
}

type markdownFile struct {
	*bytes.Reader

	file http.File
}

var _ http.File = &markdownFile{}

func openMarkdown(name string, f http.File) (*markdownFile, error) {
	bs, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	body_bs := markdown.Markdown(bs)
	d := PageData{
		Title:   strings.TrimPrefix(strings.Replace(name, "-", " ", -1), "/"),
		Content: template.HTML(string(body_bs)),
	}
	buffer := new(bytes.Buffer)
	t, err := template.ParseFiles("public-wiki.html")
	if err != nil {
		return nil, err
	}
	err = t.Execute(buffer, d)
	if err != nil {
		return nil, err
	}
	return &markdownFile{bytes.NewReader(buffer.Bytes()), f}, nil
}

func (f *markdownFile) Readdir(count int) ([]os.FileInfo, error) {
	log.Panicf("Readdir() called on file")
	return nil, nil
}

func (f *markdownFile) Stat() (os.FileInfo, error) {
	fi, err := f.file.Stat()
	if err != nil {
		return nil, err
	}
	return &markdownFileInfo{thinMarkdownFileInfo{fi}, f}, nil
}

func (f *markdownFile) Close() error {
	return f.file.Close()
}

// Doesn't need to read the file; Size() panics
type thinMarkdownFileInfo struct {
	os.FileInfo
}

func (fi *thinMarkdownFileInfo) Name() string {
	return strings.TrimSuffix(fi.FileInfo.Name(), ".md")
}

func (fi *thinMarkdownFileInfo) Size() int64 {
	log.Panicf("Size() called on thinMarkdownFileInfo")
	return 0
}

type markdownFileInfo struct {
	thinMarkdownFileInfo

	*markdownFile
}

func (fi *markdownFileInfo) Size() int64 {
	return fi.Reader.Size()
}

type markdownDir struct {
	http.File
}

func (d *markdownDir) Readdir(count int) ([]os.FileInfo, error) {
	fis, err := d.File.Readdir(count)
	if err != nil {
		return nil, err
	}
	fakefis := []os.FileInfo{}
	for _, fi := range fis {
		fakefis = append(fakefis, &thinMarkdownFileInfo{fi})
	}
	return fakefis, nil
}

type markdownFs struct {
	fs http.FileSystem
}

var _ http.FileSystem = &markdownFs{}

func (fs *markdownFs) Open(name string) (http.File, error) {
	realf, err := fs.fs.Open(name)
	if os.IsNotExist(err) {
		realf, err = fs.fs.Open(name + ".md")
	}
	if err != nil {
		return nil, err
	}
	finfo, err := realf.Stat()
	if err != nil {
		return nil, err
	}
	if finfo.IsDir() {
		return realf, nil
	} else {
		return openMarkdown(name, realf)
	}
}

func MarkdownFileSystem(fs http.FileSystem) http.FileSystem {
	return &markdownFs{fs}
}
<commit_msg>only markdownify .md files<commit_after>package main

import (
	"bytes"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"

	markdown "github.com/shurcooL/github_flavored_markdown"
)

type PageData struct {
	Title   string
	Content template.HTML
}

type markdownFile struct {
	*bytes.Reader

	file http.File
}

var _ http.File = &markdownFile{}

func openMarkdown(name string, f http.File) (*markdownFile, error) {
	bs, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	body_bs := markdown.Markdown(bs)
	d := PageData{
		Title:   strings.TrimPrefix(strings.Replace(name, "-", " ", -1), "/"),
		Content: template.HTML(string(body_bs)),
	}
	buffer := new(bytes.Buffer)
	t, err := template.ParseFiles("public-wiki.html")
	if err != nil {
		return nil, err
	}
	err = t.Execute(buffer, d)
	if err != nil {
		return nil, err
	}
	return &markdownFile{bytes.NewReader(buffer.Bytes()), f}, nil
}

func (f *markdownFile) Readdir(count int) ([]os.FileInfo, error) {
	log.Panicf("Readdir() called on file")
	return nil, nil
}

func (f *markdownFile) Stat() (os.FileInfo, error) {
	fi, err := f.file.Stat()
	if err != nil {
		return nil, err
	}
	return &markdownFileInfo{thinMarkdownFileInfo{fi}, f}, nil
}

func (f *markdownFile) Close() error {
	return f.file.Close()
}

// Doesn't need to read the file; Size() panics
type thinMarkdownFileInfo struct {
	os.FileInfo
}

func (fi *thinMarkdownFileInfo) Name() string {
	return strings.TrimSuffix(fi.FileInfo.Name(), ".md")
}

func (fi *thinMarkdownFileInfo) Size() int64 {
	log.Panicf("Size() called on thinMarkdownFileInfo")
	return 0
}

type markdownFileInfo struct {
	thinMarkdownFileInfo

	*markdownFile
}

func (fi *markdownFileInfo) Size() int64 {
	return fi.Reader.Size()
}

type markdownDir struct {
	http.File
}

func (d *markdownDir) Readdir(count int) ([]os.FileInfo, error) {
	fis, err := d.File.Readdir(count)
	if err != nil {
		return nil, err
	}
	fakefis := []os.FileInfo{}
	for _, fi := range fis {
		fakefis = append(fakefis, &thinMarkdownFileInfo{fi})
	}
	return fakefis, nil
}

type markdownFs struct {
	fs http.FileSystem
}

var _ http.FileSystem = &markdownFs{}

func (fs *markdownFs) Open(name string) (http.File, error) {
	realf, err := fs.fs.Open(name)
	if os.IsNotExist(err) {
		realf, err = fs.fs.Open(name + ".md")
	}
	if err != nil {
		return nil, err
	}
	finfo, err := realf.Stat()
	if err != nil {
		return nil, err
	}
	if finfo.IsDir() || !strings.HasSuffix(finfo.Name(), ".md") {
		return realf, nil
	} else {
		return openMarkdown(name, realf)
	}
}

func MarkdownFileSystem(fs http.FileSystem) http.FileSystem {
	return &markdownFs{fs}
}
<|endoftext|>


<commit_before>package lz4_test

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"testing"

	"github.com/pierrec/lz4/v4"
)

func _o(s ...lz4.Option) []lz4.Option {
	return s
}

func TestReader(t *testing.T) {
	goldenFiles := []string{
		"testdata/e.txt.lz4",
		"testdata/gettysburg.txt.lz4",
		"testdata/Mark.Twain-Tom.Sawyer.txt.lz4",
		"testdata/Mark.Twain-Tom.Sawyer_long.txt.lz4",
		"testdata/pg1661.txt.lz4",
		"testdata/pi.txt.lz4",
		"testdata/random.data.lz4",
		"testdata/repeat.txt.lz4",
		"testdata/pg_control.tar.lz4",
	}

	for _, fname := range goldenFiles {
		for _, opts := range [][]lz4.Option{
			nil,
			_o(lz4.ConcurrencyOption(-1)),
		} {
			label := fmt.Sprintf("%s %v", fname, opts)
			t.Run(label, func(t *testing.T) {
				fname := fname
				t.Parallel()

				f, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f.Close()

				rawfile := strings.TrimSuffix(fname, ".lz4")
				raw, err := ioutil.ReadFile(rawfile)
				if err != nil {
					t.Fatal(err)
				}

				out := new(bytes.Buffer)
				zr := lz4.NewReader(f)
				if err := zr.Apply(opts...); err != nil {
					t.Fatal(err)
				}
				n, err := io.Copy(out, zr)
				if err != nil {
					t.Error(err)
				}

				if got, want := int(n), len(raw); got != want {
					t.Errorf("invalid size: got %d; want %d", got, want)
				}

				if got, want := out.Bytes(), raw; !reflect.DeepEqual(got, want) {
					t.Fatal("uncompressed data does not match original")
				}

				if len(raw) < 20 {
					return
				}

				f2, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f2.Close()

				out.Reset()
				zr = lz4.NewReader(f2)
				_, err = io.CopyN(out, zr, 10)
				if err != nil {
					t.Fatal(err)
				}
				if !reflect.DeepEqual(out.Bytes(), raw[:10]) {
					t.Fatal("partial read does not match original")
				}
			})
		}
	}
}

func TestReader_Reset(t *testing.T) {
	data := pg1661LZ4
	buf := new(bytes.Buffer)
	src := bytes.NewReader(data)
	zr := lz4.NewReader(src)

	// Partial read.
	_, _ = io.CopyN(buf, zr, int64(len(data))/2)

	buf.Reset()
	src.Reset(data)
	// Another time to maybe trigger some edge case.
	src.Reset(data)
	zr.Reset(src)
	if _, err := io.Copy(buf, zr); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(buf.Bytes(), pg1661) {
		t.Fatal("result does not match original")
	}
}

type brokenWriter int

func (w *brokenWriter) Write(p []byte) (n int, err error) {
	n = len(p)
	if n > int(*w) {
		n = int(*w)
		err = errors.New("broken")
	}
	*w -= brokenWriter(n)
	return
}

// WriteTo should report the number of bytes successfully written,
// not the number successfully decompressed.
func TestWriteToBrokenWriter(t *testing.T) {
	const capacity = 10
	w := brokenWriter(capacity)
	r := lz4.NewReader(bytes.NewReader(pg1661LZ4))

	n, err := r.WriteTo(&w)
	switch {
	case n > capacity:
		t.Errorf("reported number of bytes written %d too big", n)
	case err == nil:
		t.Error("no error from broken Writer")
	case err.Error() != "broken":
		t.Errorf("unexpected error %q", err.Error())
	}
}

func TestReaderLegacy(t *testing.T) {
	goldenFiles := []string{
		"testdata/vmlinux_LZ4_19377.lz4",
		"testdata/bzImage_lz4_isolated.lz4",
	}

	for _, fname := range goldenFiles {
		for _, opts := range [][]lz4.Option{
			nil,
			_o(lz4.ConcurrencyOption(-1)),
		} {
			label := fmt.Sprintf("%s %v", fname, opts)
			t.Run(label, func(t *testing.T) {
				fname := fname
				t.Parallel()

				var out bytes.Buffer
				rawfile := strings.TrimSuffix(fname, ".lz4")
				raw, err := ioutil.ReadFile(rawfile)
				if err != nil {
					t.Fatal(err)
				}

				f, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f.Close()

				zr := lz4.NewReader(f)
				if err := zr.Apply(opts...); err != nil {
					t.Fatal(err)
				}
				n, err := io.Copy(&out, zr)
				if err != nil {
					t.Fatal(err, n)
				}

				if got, want := int(n), len(raw); got != want {
					t.Errorf("invalid sizes: got %d; want %d", got, want)
				}

				if got, want := out.Bytes(), raw; !bytes.Equal(got, want) {
					t.Fatal("uncompressed data does not match original")
				}

				if len(raw) < 20 {
					return
				}

				f2, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f2.Close()

				out.Reset()
				zr = lz4.NewReader(f2)
				_, err = io.CopyN(&out, zr, 10)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(out.Bytes(), raw[:10]) {
					t.Fatal("partial read does not match original")
				}

				out.Reset()
				_, err = io.CopyN(&out, zr, 10)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(out.Bytes(), raw[10:20]) {
					t.Fatal("after seek, partial read does not match original")
				}
			})
		}
	}
}
<commit_msg>Fix test on Windows<commit_after>package lz4_test

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"runtime"
	"strings"
	"testing"

	"github.com/pierrec/lz4/v4"
)

func _o(s ...lz4.Option) []lz4.Option {
	return s
}

func TestReader(t *testing.T) {
	goldenFiles := []struct {
		name   string
		isText bool
	}{
		{
			name:   "testdata/e.txt.lz4",
			isText: true,
		},
		{
			name:   "testdata/gettysburg.txt.lz4",
			isText: true,
		},
		{
			name:   "testdata/Mark.Twain-Tom.Sawyer.txt.lz4",
			isText: true,
		},
		{
			name:   "testdata/Mark.Twain-Tom.Sawyer_long.txt.lz4",
			isText: true,
		},
		{
			name:   "testdata/pg1661.txt.lz4",
			isText: false,
		},
		{
			name:   "testdata/pi.txt.lz4",
			isText: true,
		},
		{
			name:   "testdata/random.data.lz4",
			isText: false,
		},
		{
			name:   "testdata/repeat.txt.lz4",
			isText: true,
		},
		{
			name:   "testdata/pg_control.tar.lz4",
			isText: false,
		},
	}

	for _, golden := range goldenFiles {
		for _, opts := range [][]lz4.Option{
			nil,
			_o(lz4.ConcurrencyOption(-1)),
		} {
			fname := golden.name
			isText := golden.isText
			label := fmt.Sprintf("%s %v", fname, opts)
			t.Run(label, func(t *testing.T) {
				t.Parallel()

				f, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f.Close()

				rawfile := strings.TrimSuffix(fname, ".lz4")
				_raw, err := ioutil.ReadFile(rawfile)
				if err != nil {
					t.Fatal(err)
				}
				var raw []byte
				if isText && runtime.GOOS == "windows" {
					raw = []byte(strings.ReplaceAll(string(_raw), "\r\n", "\n"))
				} else {
					raw = _raw
				}

				out := new(bytes.Buffer)
				zr := lz4.NewReader(f)
				if err := zr.Apply(opts...); err != nil {
					t.Fatal(err)
				}
				n, err := io.Copy(out, zr)
				if err != nil {
					t.Error(err)
				}

				if got, want := int(n), len(raw); got != want {
					t.Errorf("invalid size: got %d; want %d", got, want)
				}

				if got, want := out.Bytes(), raw; !reflect.DeepEqual(got, want) {
					t.Fatal("uncompressed data does not match original")
				}

				if len(raw) < 20 {
					return
				}

				f2, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f2.Close()

				out.Reset()
				zr = lz4.NewReader(f2)
				_, err = io.CopyN(out, zr, 10)
				if err != nil {
					t.Fatal(err)
				}
				if !reflect.DeepEqual(out.Bytes(), raw[:10]) {
					t.Fatal("partial read does not match original")
				}
			})
		}
	}
}

func TestReader_Reset(t *testing.T) {
	data := pg1661LZ4
	buf := new(bytes.Buffer)
	src := bytes.NewReader(data)
	zr := lz4.NewReader(src)

	// Partial read.
	_, _ = io.CopyN(buf, zr, int64(len(data))/2)

	buf.Reset()
	src.Reset(data)
	// Another time to maybe trigger some edge case.
	src.Reset(data)
	zr.Reset(src)
	if _, err := io.Copy(buf, zr); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(buf.Bytes(), pg1661) {
		t.Fatal("result does not match original")
	}
}

type brokenWriter int

func (w *brokenWriter) Write(p []byte) (n int, err error) {
	n = len(p)
	if n > int(*w) {
		n = int(*w)
		err = errors.New("broken")
	}
	*w -= brokenWriter(n)
	return
}

// WriteTo should report the number of bytes successfully written,
// not the number successfully decompressed.
func TestWriteToBrokenWriter(t *testing.T) {
	const capacity = 10
	w := brokenWriter(capacity)
	r := lz4.NewReader(bytes.NewReader(pg1661LZ4))

	n, err := r.WriteTo(&w)
	switch {
	case n > capacity:
		t.Errorf("reported number of bytes written %d too big", n)
	case err == nil:
		t.Error("no error from broken Writer")
	case err.Error() != "broken":
		t.Errorf("unexpected error %q", err.Error())
	}
}

func TestReaderLegacy(t *testing.T) {
	goldenFiles := []string{
		"testdata/vmlinux_LZ4_19377.lz4",
		"testdata/bzImage_lz4_isolated.lz4",
	}

	for _, fname := range goldenFiles {
		for _, opts := range [][]lz4.Option{
			nil,
			_o(lz4.ConcurrencyOption(-1)),
		} {
			label := fmt.Sprintf("%s %v", fname, opts)
			t.Run(label, func(t *testing.T) {
				fname := fname
				t.Parallel()

				var out bytes.Buffer
				rawfile := strings.TrimSuffix(fname, ".lz4")
				raw, err := ioutil.ReadFile(rawfile)
				if err != nil {
					t.Fatal(err)
				}

				f, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f.Close()

				zr := lz4.NewReader(f)
				if err := zr.Apply(opts...); err != nil {
					t.Fatal(err)
				}
				n, err := io.Copy(&out, zr)
				if err != nil {
					t.Fatal(err, n)
				}

				if got, want := int(n), len(raw); got != want {
					t.Errorf("invalid sizes: got %d; want %d", got, want)
				}

				if got, want := out.Bytes(), raw; !bytes.Equal(got, want) {
					t.Fatal("uncompressed data does not match original")
				}

				if len(raw) < 20 {
					return
				}

				f2, err := os.Open(fname)
				if err != nil {
					t.Fatal(err)
				}
				defer f2.Close()

				out.Reset()
				zr = lz4.NewReader(f2)
				_, err = io.CopyN(&out, zr, 10)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(out.Bytes(), raw[:10]) {
					t.Fatal("partial read does not match original")
				}

				out.Reset()
				_, err = io.CopyN(&out, zr, 10)
				if err != nil {
					t.Fatal(err)
				}

				if !bytes.Equal(out.Bytes(), raw[10:20]) {
					t.Fatal("after seek, partial read does not match original")
				}
			})
		}
	}
}
<|endoftext|>


<commit_before>package unilog

import (
	"io"
	"strings"
	"testing"
	"time"
)

type op struct {
	read   int
	expect string
}

func TestUnilogReader(t *testing.T) {
	var twoLines = "hello world\nsecond line\n"
	tests := []struct {
		in  string
		ops []op
	}{
		{twoLines, []op{
			{5, "hello"},
			{-1, ""},
			{10, " "},
			{10, "w"},
			{10, "o"},
			{10, "r"},
			{1, "l"},
			{1, "d"},
			{5, "\n"},
			{5, ""},
			{1, ""},
		}},
		{twoLines, []op{
			{12, "hello world\n"},
			{-1, ""},
			{10, ""},
			{1, ""},
		}},
	}
	for i, tc := range tests {
		shutdown := make(chan struct{})
		rd := NewUnilogReader(strings.NewReader(tc.in), shutdown)
		buf := make([]byte, 128)
		for _, o := range tc.ops {
			if o.read < 0 {
				close(shutdown)
				// HACK: Wait for the shutdown to propagate.
				for !rd.(*UnilogReader).shuttingDown {
					time.Sleep(1 * time.Millisecond)
				}
				continue
			}
			if len(buf) < o.read {
				buf = make([]byte, o.read)
			}
			n, e := rd.Read(buf[:o.read])
			if n != len(o.expect) {
				t.Errorf("test %d, read(%d) expected %v, but got %d bytes",
					i, o.read, o.expect, n)
				break
			}
			got := string(buf[:n])
			if got != o.expect {
				t.Errorf("test %d, read(%d) expected %v, but got %v",
					i, o.read, o.expect, got)
				break
			}
			if e != nil && len(o.expect) > 0 {
				t.Errorf("test %d, read(%d) expected %v, but got err: %s",
					i, o.read, o.expect, e.Error())
				break
			}
			if len(o.expect) == 0 && e != io.EOF {
				t.Errorf("expected EOF, got %v", e)
				break
			}
		}
	}
}
<commit_msg>Update reader tests and deprecate old test<commit_after>package unilog

import (
	"bufio"
	"io"
	"strings"
	"testing"
	"time"
)

type op struct {
	read   int
	expect string
}

func testUnilogReader(t *testing.T) {
	var twoLines = "hello world\nsecond line\n"
	tests := []struct {
		in  string
		ops []op
	}{
		{twoLines, []op{
			{5, "hello"},
			{-1, ""},
			{10, " "},
			{10, "w"},
			{10, "o"},
			{10, "r"},
			{1, "l"},
			{1, "d"},
			{5, "\n"},
			{5, ""},
			{1, ""},
		}},
		{twoLines, []op{
			{12, "hello world\n"},
			{-1, ""},
			{10, ""},
			{1, ""},
		}},
	}
	for i, tc := range tests {
		shutdown := make(chan struct{})
		rd := NewUnilogReader(strings.NewReader(tc.in), shutdown)
		buf := make([]byte, 128)
		for _, o := range tc.ops {
			if o.read < 0 {
				close(shutdown)
				// HACK: Wait for the shutdown to propagate.
				for !rd.(*UnilogReader).shuttingDown {
					time.Sleep(1 * time.Millisecond)
				}
				continue
			}
			if len(buf) < o.read {
				buf = make([]byte, o.read)
			}
			n, e := rd.Read(buf[:o.read])
			if n != len(o.expect) {
				t.Errorf("test %d, read(%d) expected %v, but got %d bytes",
					i, o.read, o.expect, n)
				break
			}
			got := string(buf[:n])
			if got != o.expect {
				t.Errorf("test %d, read(%d) expected %v, but got %v",
					i, o.read, o.expect, got)
				break
			}
			if e != nil && len(o.expect) > 0 {
				t.Errorf("test %d, read(%d) expected %v, but got err: %s",
					i, o.read, o.expect, e.Error())
				break
			}
			if len(o.expect) == 0 && e != io.EOF {
				t.Errorf("expected EOF, got %v", e)
				break
			}
		}
	}
}

type OurReader struct {
	input chan []byte
}

func (or *OurReader) Read(p []byte) (n int, err error) {
	bs, ok := <-or.input
	if !ok {
		return 0, io.EOF
	}
	return copy(p, bs), nil
}

func TestReaderAlone(t *testing.T) {
	t.Run("basic", func(t *testing.T) {
		input := make(chan []byte, 1)
		inRdr := OurReader{input}
		shutdown := make(chan struct{}, 1)
		close(shutdown)
		rdr := NewUnilogReader(&inRdr, shutdown)
		bufRdr := bufio.NewReader(rdr)
		input <- []byte("hello I am a line\n")
		line, err := bufRdr.ReadString('\n')
		if err != nil {
			t.Error(err)
		}
		if line != "hello I am a line\n" {
			t.Error(line)
		}
	})

	// Less happy case: shutdown with a missing newline.
	t.Run("complex", func(t *testing.T) {
		input := make(chan []byte)
		inRdr := OurReader{input}
		shutdown := make(chan struct{})
		close(shutdown)
		rdr := NewUnilogReader(&inRdr, shutdown)
		bufRdr := bufio.NewReader(rdr)

		go func() {
			input <- []byte("unterminated")
			input <- []byte("done\n")
		}()
		line, err := bufRdr.ReadString('\n')
		if err != nil {
			t.Error(err)
		}
		if line != "unterminateddone\n" {
			t.Error(line)
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the LGPLv3, see LICENCE file for details.

package debugstatus_test

import (
	"errors"
	"time"

	jujutesting "github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
	"gopkg.in/mgo.v2"

	"github.com/juju/utils/debugstatus"
)

type statusSuite struct {
	jujutesting.IsolationSuite
}

var _ = gc.Suite(&statusSuite{})

func makeCheckerFunc(key, name, value string, passed bool) debugstatus.CheckerFunc {
	return func() (string, debugstatus.CheckResult) {
		return key, debugstatus.CheckResult{
			Name:   name,
			Value:  value,
			Passed: passed,
		}
	}
}

func (s *statusSuite) TestCheck(c *gc.C) {
	results := debugstatus.Check(
		makeCheckerFunc("check1", "check1 name", "value1", true),
		makeCheckerFunc("check2", "check2 name", "value2", false),
		makeCheckerFunc("check3", "check3 name", "value3", true),
	)

	c.Assert(results, jc.DeepEquals, map[string]debugstatus.CheckResult{
		"check1": {
			Name:   "check1 name",
			Value:  "value1",
			Passed: true,
		},
		"check2": {
			Name:   "check2 name",
			Value:  "value2",
			Passed: false,
		},
		"check3": {
			Name:   "check3 name",
			Value:  "value3",
			Passed: true,
		},
	})
}

func (s *statusSuite) TestServerStartTime(c *gc.C) {
	startTime := time.Now()
	original := debugstatus.StartTime
	debugstatus.StartTime = startTime
	s.AddCleanup(func(*gc.C) {
		debugstatus.StartTime = original
	})
	key, result := debugstatus.ServerStartTime()
	c.Assert(key, gc.Equals, "server_started")
	c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
		Name:   "Server started",
		Value:  startTime.String(),
		Passed: true,
	})
}

func (s *statusSuite) TestConnection(c *gc.C) {
	// Ensure a connection established is properly reported.
	check := debugstatus.Connection(pinger{nil})
	key, result := check()
	c.Assert(key, gc.Equals, "mongo_connected")
	c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
		Name:   "MongoDB is connected",
		Value:  "Connected",
		Passed: true,
	})

	// An error is reported if ping fails.
	check = debugstatus.Connection(pinger{errors.New("bad wolf")})
	key, result = check()
	c.Assert(key, gc.Equals, "mongo_connected")
	c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
		Name:   "MongoDB is connected",
		Value:  "Ping error: bad wolf",
		Passed: false,
	})
}

// pinger implements a debugstatus.Pinger used for tests.
type pinger struct {
	err error
}

func (p pinger) Ping() error {
	return p.err
}

var mongoCollectionsTests = []struct {
	about        string
	collector    collector
	expectValue  string
	expectPassed bool
}{{
	about: "all collection exist",
	collector: collector{
		expected: []string{"coll1", "coll2"},
		obtained: []string{"coll1", "coll2"},
	},
	expectValue:  "All required collections exist",
	expectPassed: true,
}, {
	about:        "no collections",
	expectValue:  "All required collections exist",
	expectPassed: true,
}, {
	about: "missing collections",
	collector: collector{
		expected: []string{"coll1", "coll2", "coll3"},
		obtained: []string{"coll2"},
	},
	expectValue: "Missing collections: [coll1 coll3]",
}, {
	about: "error retrieving collections",
	collector: collector{
		err: errors.New("bad wolf"),
	},
	expectValue: "Cannot get collections: bad wolf",
}}

func (s *statusSuite) TestMongoCollections(c *gc.C) {
	for i, test := range mongoCollectionsTests {
		c.Logf("test %d: %s", i, test.about)

		// Ensure a connection established is properly reported.
		check := debugstatus.MongoCollections(test.collector)
		key, result := check()
		c.Assert(key, gc.Equals, "mongo_collections")
		c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
			Name:   "MongoDB collections",
			Value:  test.expectValue,
			Passed: test.expectPassed,
		})
	}
}

// collector implements a debugstatus.Collector used for tests.
type collector struct {
	expected []string
	obtained []string
	err      error
}

func (c collector) CollectionNames() ([]string, error) {
	return c.obtained, c.err
}

func (c collector) Collections() []*mgo.Collection {
	collections := make([]*mgo.Collection, len(c.expected))
	for i, name := range c.expected {
		collections[i] = &mgo.Collection{Name: name}
	}
	return collections
}

var renameTests = []struct {
	about string
	key   string
	name  string
}{{
	about: "rename key",
	key:   "new-key",
}, {
	about: "rename name",
	name:  "new name",
}, {
	about: "rename both",
	key:   "another-key",
	name:  "another name",
}, {
	about: "do not rename",
}}

func (s *statusSuite) TestRename(c *gc.C) {
	check := makeCheckerFunc("old-key", "old name", "value", true)
	for i, test := range renameTests {
		c.Logf("test %d: %s", i, test.about)

		// Rename and run the check.
		key, result := debugstatus.Rename(test.key, test.name, check)()

		// Ensure the results are successfully renamed.
		expectKey := test.key
		if expectKey == "" {
			expectKey = "old-key"
		}
		expectName := test.name
		if expectName == "" {
			expectName = "old name"
		}
		c.Assert(key, gc.Equals, expectKey)
		c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
			Name:   expectName,
			Value:  "value",
			Passed: true,
		})
	}
}
<commit_msg>Use PatchValue to patch debugstatus.StartTime.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the LGPLv3, see LICENCE file for details.

package debugstatus_test

import (
	"errors"
	"time"

	jujutesting "github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
	"gopkg.in/mgo.v2"

	"github.com/juju/utils/debugstatus"
)

type statusSuite struct {
	jujutesting.IsolationSuite
}

var _ = gc.Suite(&statusSuite{})

func makeCheckerFunc(key, name, value string, passed bool) debugstatus.CheckerFunc {
	return func() (string, debugstatus.CheckResult) {
		return key, debugstatus.CheckResult{
			Name:   name,
			Value:  value,
			Passed: passed,
		}
	}
}

func (s *statusSuite) TestCheck(c *gc.C) {
	results := debugstatus.Check(
		makeCheckerFunc("check1", "check1 name", "value1", true),
		makeCheckerFunc("check2", "check2 name", "value2", false),
		makeCheckerFunc("check3", "check3 name", "value3", true),
	)

	c.Assert(results, jc.DeepEquals, map[string]debugstatus.CheckResult{
		"check1": {
			Name:   "check1 name",
			Value:  "value1",
			Passed: true,
		},
		"check2": {
			Name:   "check2 name",
			Value:  "value2",
			Passed: false,
		},
		"check3": {
			Name:   "check3 name",
			Value:  "value3",
			Passed: true,
		},
	})
}

func (s *statusSuite) TestServerStartTime(c *gc.C) {
	startTime := time.Now()
	s.PatchValue(&debugstatus.StartTime, startTime)
	key, result := debugstatus.ServerStartTime()
	c.Assert(key, gc.Equals, "server_started")
	c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
		Name:   "Server started",
		Value:  startTime.String(),
		Passed: true,
	})
}

func (s *statusSuite) TestConnection(c *gc.C) {
	// Ensure a connection established is properly reported.
	check := debugstatus.Connection(pinger{nil})
	key, result := check()
	c.Assert(key, gc.Equals, "mongo_connected")
	c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
		Name:   "MongoDB is connected",
		Value:  "Connected",
		Passed: true,
	})

	// An error is reported if ping fails.
	check = debugstatus.Connection(pinger{errors.New("bad wolf")})
	key, result = check()
	c.Assert(key, gc.Equals, "mongo_connected")
	c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
		Name:   "MongoDB is connected",
		Value:  "Ping error: bad wolf",
		Passed: false,
	})
}

// pinger implements a debugstatus.Pinger used for tests.
type pinger struct {
	err error
}

func (p pinger) Ping() error {
	return p.err
}

var mongoCollectionsTests = []struct {
	about        string
	collector    collector
	expectValue  string
	expectPassed bool
}{{
	about: "all collection exist",
	collector: collector{
		expected: []string{"coll1", "coll2"},
		obtained: []string{"coll1", "coll2"},
	},
	expectValue:  "All required collections exist",
	expectPassed: true,
}, {
	about:        "no collections",
	expectValue:  "All required collections exist",
	expectPassed: true,
}, {
	about: "missing collections",
	collector: collector{
		expected: []string{"coll1", "coll2", "coll3"},
		obtained: []string{"coll2"},
	},
	expectValue: "Missing collections: [coll1 coll3]",
}, {
	about: "error retrieving collections",
	collector: collector{
		err: errors.New("bad wolf"),
	},
	expectValue: "Cannot get collections: bad wolf",
}}

func (s *statusSuite) TestMongoCollections(c *gc.C) {
	for i, test := range mongoCollectionsTests {
		c.Logf("test %d: %s", i, test.about)

		// Ensure a connection established is properly reported.
		check := debugstatus.MongoCollections(test.collector)
		key, result := check()
		c.Assert(key, gc.Equals, "mongo_collections")
		c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
			Name:   "MongoDB collections",
			Value:  test.expectValue,
			Passed: test.expectPassed,
		})
	}
}

// collector implements a debugstatus.Collector used for tests.
type collector struct {
	expected []string
	obtained []string
	err      error
}

func (c collector) CollectionNames() ([]string, error) {
	return c.obtained, c.err
}

func (c collector) Collections() []*mgo.Collection {
	collections := make([]*mgo.Collection, len(c.expected))
	for i, name := range c.expected {
		collections[i] = &mgo.Collection{Name: name}
	}
	return collections
}

var renameTests = []struct {
	about string
	key   string
	name  string
}{{
	about: "rename key",
	key:   "new-key",
}, {
	about: "rename name",
	name:  "new name",
}, {
	about: "rename both",
	key:   "another-key",
	name:  "another name",
}, {
	about: "do not rename",
}}

func (s *statusSuite) TestRename(c *gc.C) {
	check := makeCheckerFunc("old-key", "old name", "value", true)
	for i, test := range renameTests {
		c.Logf("test %d: %s", i, test.about)

		// Rename and run the check.
		key, result := debugstatus.Rename(test.key, test.name, check)()

		// Ensure the results are successfully renamed.
		expectKey := test.key
		if expectKey == "" {
			expectKey = "old-key"
		}
		expectName := test.name
		if expectName == "" {
			expectName = "old name"
		}
		c.Assert(key, gc.Equals, expectKey)
		c.Assert(result, jc.DeepEquals, debugstatus.CheckResult{
			Name:   expectName,
			Value:  "value",
			Passed: true,
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"lib/datastore"
	"lib/filelock"
	"lib/flannel"
	"lib/metrics"
	"lib/mutualtls"
	"lib/policy_client"
	"lib/rules"
	"lib/serial"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
	"vxlan-policy-agent/agent_metrics"
	"vxlan-policy-agent/config"
	"vxlan-policy-agent/enforcer"
	"vxlan-policy-agent/handlers"
	"vxlan-policy-agent/planner"
	"vxlan-policy-agent/poller"

	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/lager"
	"github.com/cloudfoundry/dropsonde"
	"github.com/coreos/go-iptables/iptables"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

const (
	dropsondeOrigin = "vxlan-policy-agent"
	emitInterval    = 30 * time.Second
)

func die(logger lager.Logger, action string, err error) {
	logger.Error(action, err)
	os.Exit(1)
}

func main() {
	conf := &config.VxlanPolicyAgent{}

	configFilePath := flag.String("config-file", "", "path to config file")
	flag.Parse()

	configBytes, err := ioutil.ReadFile(*configFilePath)
	if err != nil {
		log.Fatalf("error reading config: %s", err)
	}

	err = json.Unmarshal(configBytes, conf)
	if err != nil {
		log.Fatalf("error unmarshalling config: %s", err)
	}

	logger := lager.NewLogger("container-networking.vxlan-policy-agent")
	reconfigurableSink := initLoggerSink(logger, conf.LogLevel)
	logger.RegisterSink(reconfigurableSink)

	logger.Info("parsed-config", lager.Data{"config": conf})

	pollInterval := time.Duration(conf.PollInterval) * time.Second
	if pollInterval == 0 {
		pollInterval = time.Second
	}

	flannelInfoReader := &flannel.NetworkInfo{
		FlannelSubnetFilePath: conf.FlannelSubnetFile,
	}

	localSubnetCIDR, _, err := flannelInfoReader.DiscoverNetworkInfo()
	if err != nil {
		die(logger, "discovering network info", err)
	}

	clientTLSConfig, err := mutualtls.NewClientTLSConfig(conf.ClientCertFile, conf.ClientKeyFile, conf.ServerCACertFile)
	if err != nil {
		die(logger, "mutual tls config", err)
	}

	httpClient := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: clientTLSConfig,
		},
	}

	policyClient := policy_client.NewInternal(
		logger.Session("policy-client"),
		httpClient,
		conf.PolicyServerURL,
	)

	store := &datastore.Store{
		Serializer: &serial.Serial{},
		Locker: &filelock.Locker{
			Path: conf.Datastore,
		},
	}

	ipt, err := iptables.New()
	if err != nil {
		die(logger, "iptables-new", err)
	}

	iptLocker := &rules.IPTablesLocker{
		FileLocker: &filelock.Locker{Path: conf.IPTablesLockFile},
		Mutex:      &sync.Mutex{},
	}
	restorer := &rules.Restorer{}
	lockedIPTables := &rules.LockedIPTables{
		IPTables: ipt,
		Locker:   iptLocker,
		Restorer: restorer,
	}

	timeMetricsEmitter := &agent_metrics.TimeMetrics{
		Logger: logger.Session("time-metric-emitter"),
	}

	iptablesLoggingState := &planner.LoggingState{}
	dynamicPlanner := &planner.VxlanPolicyPlanner{
		Datastore:         store,
		PolicyClient:      policyClient,
		Logger:            logger.Session("rules-updater"),
		VNI:               conf.VNI,
		CollectionEmitter: timeMetricsEmitter,
		Chain: enforcer.Chain{
			Table:       "filter",
			ParentChain: "FORWARD",
			Prefix:      "vpa--",
		},
		LoggingState: iptablesLoggingState,
	}

	timestamper := &enforcer.Timestamper{}
	ruleEnforcer := enforcer.NewEnforcer(
		logger.Session("rules-enforcer"),
		timestamper,
		lockedIPTables,
	)

	vxlanDefaultLocalPlanner := &planner.VxlanDefaultLocalPlanner{
		Logger:      logger,
		LocalSubnet: localSubnetCIDR,
		Chain: enforcer.Chain{
			Table:       "filter",
			ParentChain: "FORWARD",
			Prefix:      "vpa--local-",
		},
		LoggingState: iptablesLoggingState,
	}

	vxlanDefaultRemotePlanner := &planner.VxlanDefaultRemotePlanner{
		Logger: logger,
		VNI:    conf.VNI,
		Chain: enforcer.Chain{
			Table:       "filter",
			ParentChain: "FORWARD",
			Prefix:      "vpa--remote-",
		},
		LoggingState: iptablesLoggingState,
	}

	err = dropsonde.Initialize(conf.MetronAddress, dropsondeOrigin)
	if err != nil {
		log.Fatalf("initializing dropsonde: %s", err)
	}

	uptimeSource := metrics.NewUptimeSource()
	metricsEmitter := metrics.NewMetricsEmitter(logger, emitInterval, uptimeSource)

	policyPoller := &poller.Poller{
		Logger:       logger,
		PollInterval: pollInterval,
		SingleCycleFunc: (&poller.SinglePollCycle{
			Planners: []poller.Planner{
				vxlanDefaultLocalPlanner,
				vxlanDefaultRemotePlanner,
				dynamicPlanner,
			},
			Enforcer:          ruleEnforcer,
			CollectionEmitter: timeMetricsEmitter,
		}).DoCycle,
	}

	debugServerAddress := fmt.Sprintf("%s:%d", conf.DebugServerHost, conf.DebugServerPort)
	debugServer := createCustomDebugServer(debugServerAddress, reconfigurableSink, iptablesLoggingState)
	members := grouper.Members{
		{"metrics_emitter", metricsEmitter},
		{"policy_poller", policyPoller},
		{"debug-server", debugServer},
	}

	monitor := ifrit.Invoke(sigmon.New(grouper.NewOrdered(os.Interrupt, members)))
	logger.Info("starting")
	err = <-monitor.Wait()
	if err != nil {
		die(logger, "ifrit monitor", err)
	}
}

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func initLoggerSink(logger lager.Logger, level string) *lager.ReconfigurableSink {
	var logLevel lager.LogLevel
	switch strings.ToLower(level) {
	case DEBUG:
		logLevel = lager.DEBUG
	case INFO:
		logLevel = lager.INFO
	case ERROR:
		logLevel = lager.ERROR
	case FATAL:
		logLevel = lager.FATAL
	default:
		logLevel = lager.INFO
	}
	w := lager.NewWriterSink(os.Stdout, lager.DEBUG)
	return lager.NewReconfigurableSink(w, logLevel)
}

func createCustomDebugServer(listenAddress string, sink *lager.ReconfigurableSink, iptablesLoggingState *planner.LoggingState) ifrit.Runner {
	mux := debugserver.Handler(sink).(*http.ServeMux)
	mux.Handle("/iptables_logging", &handlers.IPTablesLogging{
		LoggingState: iptablesLoggingState,
	})
	return http_server.New(listenAddress, mux)
}
<commit_msg>Add logger dependency to poll cycle<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"lib/datastore"
	"lib/filelock"
	"lib/flannel"
	"lib/metrics"
	"lib/mutualtls"
	"lib/policy_client"
	"lib/rules"
	"lib/serial"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
	"vxlan-policy-agent/agent_metrics"
	"vxlan-policy-agent/config"
	"vxlan-policy-agent/enforcer"
	"vxlan-policy-agent/handlers"
	"vxlan-policy-agent/planner"
	"vxlan-policy-agent/poller"

	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/lager"
	"github.com/cloudfoundry/dropsonde"
	"github.com/coreos/go-iptables/iptables"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

const (
	dropsondeOrigin = "vxlan-policy-agent"
	emitInterval    = 30 * time.Second
)

func die(logger lager.Logger, action string, err error) {
	logger.Error(action, err)
	os.Exit(1)
}

func main() {
	conf := &config.VxlanPolicyAgent{}

	configFilePath := flag.String("config-file", "", "path to config file")
	flag.Parse()

	configBytes, err := ioutil.ReadFile(*configFilePath)
	if err != nil {
		log.Fatalf("error reading config: %s", err)
	}

	err = json.Unmarshal(configBytes, conf)
	if err != nil {
		log.Fatalf("error unmarshalling config: %s", err)
	}

	logger := lager.NewLogger("container-networking.vxlan-policy-agent")
	reconfigurableSink := initLoggerSink(logger, conf.LogLevel)
	logger.RegisterSink(reconfigurableSink)

	logger.Info("parsed-config", lager.Data{"config": conf})

	pollInterval := time.Duration(conf.PollInterval) * time.Second
	if pollInterval == 0 {
		pollInterval = time.Second
	}

	flannelInfoReader := &flannel.NetworkInfo{
		FlannelSubnetFilePath: conf.FlannelSubnetFile,
	}

	localSubnetCIDR, _, err := flannelInfoReader.DiscoverNetworkInfo()
	if err != nil {
		die(logger, "discovering network info", err)
	}

	clientTLSConfig, err := mutualtls.NewClientTLSConfig(conf.ClientCertFile, conf.ClientKeyFile, conf.ServerCACertFile)
	if err != nil {
		die(logger, "mutual tls config", err)
	}

	httpClient := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: clientTLSConfig,
		},
	}

	policyClient := policy_client.NewInternal(
		logger.Session("policy-client"),
		httpClient,
		conf.PolicyServerURL,
	)

	store := &datastore.Store{
		Serializer: &serial.Serial{},
		Locker: &filelock.Locker{
			Path: conf.Datastore,
		},
	}

	ipt, err := iptables.New()
	if err != nil {
		die(logger, "iptables-new", err)
	}

	iptLocker := &rules.IPTablesLocker{
		FileLocker: &filelock.Locker{Path: conf.IPTablesLockFile},
		Mutex:      &sync.Mutex{},
	}
	restorer := &rules.Restorer{}
	lockedIPTables := &rules.LockedIPTables{
		IPTables: ipt,
		Locker:   iptLocker,
		Restorer: restorer,
	}

	timeMetricsEmitter := &agent_metrics.TimeMetrics{
		Logger: logger.Session("time-metric-emitter"),
	}

	iptablesLoggingState := &planner.LoggingState{}
	dynamicPlanner := &planner.VxlanPolicyPlanner{
		Datastore:         store,
		PolicyClient:      policyClient,
		Logger:            logger.Session("rules-updater"),
		VNI:               conf.VNI,
		CollectionEmitter: timeMetricsEmitter,
		Chain: enforcer.Chain{
			Table:       "filter",
			ParentChain: "FORWARD",
			Prefix:      "vpa--",
		},
		LoggingState: iptablesLoggingState,
	}

	timestamper := &enforcer.Timestamper{}
	ruleEnforcer := enforcer.NewEnforcer(
		logger.Session("rules-enforcer"),
		timestamper,
		lockedIPTables,
	)

	vxlanDefaultLocalPlanner := &planner.VxlanDefaultLocalPlanner{
		Logger:      logger,
		LocalSubnet: localSubnetCIDR,
		Chain: enforcer.Chain{
			Table:       "filter",
			ParentChain: "FORWARD",
			Prefix:      "vpa--local-",
		},
		LoggingState: iptablesLoggingState,
	}

	vxlanDefaultRemotePlanner := &planner.VxlanDefaultRemotePlanner{
		Logger: logger,
		VNI:    conf.VNI,
		Chain: enforcer.Chain{
			Table:       "filter",
			ParentChain: "FORWARD",
			Prefix:      "vpa--remote-",
		},
		LoggingState: iptablesLoggingState,
	}

	err = dropsonde.Initialize(conf.MetronAddress, dropsondeOrigin)
	if err != nil {
		log.Fatalf("initializing dropsonde: %s", err)
	}

	uptimeSource := metrics.NewUptimeSource()
	metricsEmitter := metrics.NewMetricsEmitter(logger, emitInterval, uptimeSource)

	policyPoller := &poller.Poller{
		Logger:       logger,
		PollInterval: pollInterval,
		SingleCycleFunc: (&poller.SinglePollCycle{
			Planners: []poller.Planner{
				vxlanDefaultLocalPlanner,
				vxlanDefaultRemotePlanner,
				dynamicPlanner,
			},
			Enforcer:          ruleEnforcer,
			CollectionEmitter: timeMetricsEmitter,
			Logger:            logger,
		}).DoCycle,
	}

	debugServerAddress := fmt.Sprintf("%s:%d", conf.DebugServerHost, conf.DebugServerPort)
	debugServer := createCustomDebugServer(debugServerAddress, reconfigurableSink, iptablesLoggingState)
	members := grouper.Members{
		{"metrics_emitter", metricsEmitter},
		{"policy_poller", policyPoller},
		{"debug-server", debugServer},
	}

	monitor := ifrit.Invoke(sigmon.New(grouper.NewOrdered(os.Interrupt, members)))
	logger.Info("starting")
	err = <-monitor.Wait()
	if err != nil {
		die(logger, "ifrit monitor", err)
	}
}

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func initLoggerSink(logger lager.Logger, level string) *lager.ReconfigurableSink {
	var logLevel lager.LogLevel
	switch strings.ToLower(level) {
	case DEBUG:
		logLevel = lager.DEBUG
	case INFO:
		logLevel = lager.INFO
	case ERROR:
		logLevel = lager.ERROR
	case FATAL:
		logLevel = lager.FATAL
	default:
		logLevel = lager.INFO
	}
	w := lager.NewWriterSink(os.Stdout, lager.DEBUG)
	return lager.NewReconfigurableSink(w, logLevel)
}

func createCustomDebugServer(listenAddress string, sink *lager.ReconfigurableSink, iptablesLoggingState *planner.LoggingState) ifrit.Runner {
	mux := debugserver.Handler(sink).(*http.ServeMux)
	mux.Handle("/iptables_logging", &handlers.IPTablesLogging{
		LoggingState: iptablesLoggingState,
	})
	return http_server.New(listenAddress, mux)
}
<|endoftext|>


<commit_before>package ec2

import (
	"bytes"
	"fmt"
	"github.com/megamsys/megamd/iaas"
	"github.com/megamsys/megamd/provisioner"
	"github.com/tsuru/config"
	"strings"
	"encoding/json"
)

func Init() {
	iaas.RegisterIaasProvider("ec2", &EC2IaaS{})
}

type EC2IaaS struct{}

func (i *EC2IaaS) DeleteMachine(string) error {

	return nil
}

func (i *EC2IaaS) CreateMachine(pdc *iaas.PredefClouds, assembly *provisioner.AssemblyResult) (string, error) {
	keys, err_keys := iaas.GetAccessKeys(pdc)
	if err_keys != nil {
		return "", err_keys
	}

	str, err := buildCommand(iaas.GetPlugins("ec2"), pdc, "create")
	if err != nil {
		return "", err
	}
	str = str + " -N " + assembly.Name + "." + assembly.Components[0].Inputs.Domain
	str = str + " -A " + keys.AccessKey
	str = str + " -K " + keys.SecretKey

	riak, err_riak := config.GetString("api:server")
	if err_riak != nil {
		return "", err_riak
	}

	recipe, err_recipe := config.GetString("knife:recipe")
	if err_recipe != nil {
		return "", err_recipe
	}

	str = str + " --run-list \"" + "recipe[" + recipe + "]" + "\""
	attributes := &iaas.Attributes{RiakHost: riak, AccountID: pdc.Accounts_id, AssemblyID: assembly.Id}
    b, aerr := json.Marshal(attributes)
    if aerr != nil {
        fmt.Println(aerr)
        return "", aerr
    }
	str = str + " --json-attributes '" + string(b) + "'"

	//strings.Replace(str,"-c","-c "+assembly.Name+"."+assembly.Components[0].Inputs.Domain,-1)
	knifePath, kerr := config.GetString("knife:path")
	if kerr != nil {
		return "", kerr
	}
	str = strings.Replace(str, "-c", "-c "+knifePath, -1)
	return str, nil
}

func buildCommand(plugin *iaas.Plugins, pdc *iaas.PredefClouds, command string) (string, error) {
	var buffer bytes.Buffer
	if len(plugin.Tool) > 0 {
		buffer.WriteString(plugin.Tool)
	} else {
		return "", fmt.Errorf("Plugin tool doesn't loaded")
	}

	if command == "create" {
		if len(plugin.Command.Create) > 0 {
			buffer.WriteString(" " + plugin.Command.Create)
		} else {
			return "", fmt.Errorf("Plugin commands doesn't loaded")
		}
	}

	if len(pdc.Spec.Groups) > 0 {
		buffer.WriteString(" -G " + pdc.Spec.Groups)
	} else {
		return "", fmt.Errorf("Groups doesn't loaded")
	}

	if len(pdc.Spec.Image) > 0 {
		buffer.WriteString(" -I " + pdc.Spec.Image)
	} else {
		return "", fmt.Errorf("Image doesn't loaded")
	}

	if len(pdc.Spec.Flavor) > 0 {
		buffer.WriteString(" -f " + pdc.Spec.Flavor)
	} else {
		return "", fmt.Errorf("Flavor doesn't loaded")
	}

	if len(pdc.Access.Sshkey) > 0 {
		buffer.WriteString(" -S " + pdc.Access.Sshkey)
	} else {
		return "", fmt.Errorf("Ssh key value doesn't loaded")
	}

	if len(pdc.Access.Sshuser) > 0 {
		buffer.WriteString(" -x " + pdc.Access.Sshuser)
	} else {
		return "", fmt.Errorf("Ssh user value doesn't loaded")
	}

	if len(pdc.Access.IdentityFile) > 0 {
		ifile, err := iaas.GetIdentityFileLocation(pdc.Access.IdentityFile)
		if err != nil {
			return "", fmt.Errorf("Identity file doesn't loaded")
		}
		buffer.WriteString(" --identity-file " + ifile)
	} else {
		return "", fmt.Errorf("Identity file doesn't loaded")
	}

	if len(pdc.Access.Region) > 0 {
		buffer.WriteString(" --region " + pdc.Access.Region)
	} else {
		return "", fmt.Errorf("Zone doesn't loaded")
	}

	return buffer.String(), nil
}
<commit_msg>test ec2 json-attribute change<commit_after>package ec2

import (
	"bytes"
	"fmt"
	"github.com/megamsys/megamd/iaas"
	"github.com/megamsys/megamd/provisioner"
	"github.com/tsuru/config"
	"strings"
	"encoding/json"
)

func Init() {
	iaas.RegisterIaasProvider("ec2", &EC2IaaS{})
}

type EC2IaaS struct{}

func (i *EC2IaaS) DeleteMachine(string) error {

	return nil
}

func (i *EC2IaaS) CreateMachine(pdc *iaas.PredefClouds, assembly *provisioner.AssemblyResult) (string, error) {
	keys, err_keys := iaas.GetAccessKeys(pdc)
	if err_keys != nil {
		return "", err_keys
	}

	str, err := buildCommand(iaas.GetPlugins("ec2"), pdc, "create")
	if err != nil {
		return "", err
	}
	str = str + " -N " + assembly.Name + "." + assembly.Components[0].Inputs.Domain
	str = str + " -A " + keys.AccessKey
	str = str + " -K " + keys.SecretKey

	riak, err_riak := config.GetString("api:server")
	if err_riak != nil {
		return "", err_riak
	}

	recipe, err_recipe := config.GetString("knife:recipe")
	if err_recipe != nil {
		return "", err_recipe
	}

	str = str + " --run-list \"" + "recipe[" + recipe + "]" + "\""
	attributes := &iaas.Attributes{RiakHost: riak, AccountID: pdc.Accounts_id, AssemblyID: assembly.Id}
    b, aerr := json.Marshal(attributes)
    if aerr != nil {
        fmt.Println(aerr)
        return "", aerr
    }
	str = str + " --json-attributes " + string(b)

	//strings.Replace(str,"-c","-c "+assembly.Name+"."+assembly.Components[0].Inputs.Domain,-1)
	knifePath, kerr := config.GetString("knife:path")
	if kerr != nil {
		return "", kerr
	}
	str = strings.Replace(str, "-c", "-c "+knifePath, -1)
	return str, nil
}

func buildCommand(plugin *iaas.Plugins, pdc *iaas.PredefClouds, command string) (string, error) {
	var buffer bytes.Buffer
	if len(plugin.Tool) > 0 {
		buffer.WriteString(plugin.Tool)
	} else {
		return "", fmt.Errorf("Plugin tool doesn't loaded")
	}

	if command == "create" {
		if len(plugin.Command.Create) > 0 {
			buffer.WriteString(" " + plugin.Command.Create)
		} else {
			return "", fmt.Errorf("Plugin commands doesn't loaded")
		}
	}

	if len(pdc.Spec.Groups) > 0 {
		buffer.WriteString(" -G " + pdc.Spec.Groups)
	} else {
		return "", fmt.Errorf("Groups doesn't loaded")
	}

	if len(pdc.Spec.Image) > 0 {
		buffer.WriteString(" -I " + pdc.Spec.Image)
	} else {
		return "", fmt.Errorf("Image doesn't loaded")
	}

	if len(pdc.Spec.Flavor) > 0 {
		buffer.WriteString(" -f " + pdc.Spec.Flavor)
	} else {
		return "", fmt.Errorf("Flavor doesn't loaded")
	}

	if len(pdc.Access.Sshkey) > 0 {
		buffer.WriteString(" -S " + pdc.Access.Sshkey)
	} else {
		return "", fmt.Errorf("Ssh key value doesn't loaded")
	}

	if len(pdc.Access.Sshuser) > 0 {
		buffer.WriteString(" -x " + pdc.Access.Sshuser)
	} else {
		return "", fmt.Errorf("Ssh user value doesn't loaded")
	}

	if len(pdc.Access.IdentityFile) > 0 {
		ifile, err := iaas.GetIdentityFileLocation(pdc.Access.IdentityFile)
		if err != nil {
			return "", fmt.Errorf("Identity file doesn't loaded")
		}
		buffer.WriteString(" --identity-file " + ifile)
	} else {
		return "", fmt.Errorf("Identity file doesn't loaded")
	}

	if len(pdc.Access.Region) > 0 {
		buffer.WriteString(" --region " + pdc.Access.Region)
	} else {
		return "", fmt.Errorf("Zone doesn't loaded")
	}

	return buffer.String(), nil
}
<|endoftext|>