<commit_before>package watcher

import (
	"os"
	"sync/atomic"

	"github.com/cloudfoundry-incubator/receptor"
	"github.com/cloudfoundry-incubator/route-emitter/cfroutes"
	"github.com/cloudfoundry-incubator/route-emitter/nats_emitter"
	"github.com/cloudfoundry-incubator/route-emitter/routing_table"
	"github.com/cloudfoundry-incubator/route-emitter/syncer"
	"github.com/cloudfoundry-incubator/runtime-schema/metric"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
)

var (
	routesTotal  = metric.Metric("RoutesTotal")
	routesSynced = metric.Counter("RoutesSynced")

	routeSyncDuration = metric.Duration("RouteEmitterSyncDuration")

	routesRegistered   = metric.Counter("RoutesRegistered")
	routesUnregistered = metric.Counter("RoutesUnregistered")
)

type Watcher struct {
	receptorClient receptor.Client
	clock          clock.Clock
	table          routing_table.RoutingTable
	emitter        nats_emitter.NATSEmitter
	syncEvents     syncer.Events
	logger         lager.Logger
}

type syncEndEvent struct {
	table    routing_table.RoutingTable
	callback func(routing_table.RoutingTable)

	logger lager.Logger
}

type set map[interface{}]struct{}

func (set set) contains(value interface{}) bool {
	_, found := set[value]
	return found
}

func (set set) add(value interface{}) {
	set[value] = struct{}{}
}

func NewWatcher(
	receptorClient receptor.Client,
	clock clock.Clock,
	table routing_table.RoutingTable,
	emitter nats_emitter.NATSEmitter,
	syncEvents syncer.Events,
	logger lager.Logger,
) *Watcher {
	return &Watcher{
		receptorClient: receptorClient,
		clock:          clock,
		table:          table,
		emitter:        emitter,
		syncEvents:     syncEvents,
		logger:         logger.Session("watcher"),
	}
}

func (watcher *Watcher) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	watcher.logger.Info("starting")

	close(ready)
	watcher.logger.Info("started")
	defer watcher.logger.Info("finished")

	var cachedEvents map[string]receptor.Event

	eventChan := make(chan receptor.Event)
	syncEndChan := make(chan syncEndEvent)

	syncing := false

	var eventSource atomic.Value
	var stopEventSource int32

	startEventSource := func() {
		go func() {
			var err error
			var es receptor.EventSource

			for {
				if atomic.LoadInt32(&stopEventSource) == 1 {
					return
				}

				es, err = watcher.receptorClient.SubscribeToEvents()
				if err != nil {
					watcher.logger.Error("failed-subscribing-to-events", err)
					continue
				}

				eventSource.Store(es)

				var event receptor.Event
				for {
					event, err = es.Next()
					if err != nil {
						watcher.logger.Error("failed-getting-next-event", err)
						break
					}

					if event != nil {
						eventChan <- event
					}
				}
			}
		}()
	}

	startedEventSource := false
	for {
		select {
		case <-watcher.syncEvents.Sync:
			if syncing == false {
				logger := watcher.logger.Session("sync")
				logger.Info("starting")
				syncing = true

				if !startedEventSource {
					startedEventSource = true
					startEventSource()
				}

				cachedEvents = make(map[string]receptor.Event)
				go watcher.sync(logger, syncEndChan)
			}

		case syncEnd := <-syncEndChan:
			watcher.completeSync(syncEnd, cachedEvents)
			cachedEvents = nil
			syncing = false
			syncEnd.logger.Info("complete")

		case <-watcher.syncEvents.Emit:
			logger := watcher.logger.Session("emit")
			watcher.emit(logger)

		case event := <-eventChan:
			if syncing {
				watcher.logger.Info("caching-event", lager.Data{
					"type": event.EventType(),
				})

				cachedEvents[event.Key()] = event
			} else {
				watcher.handleEvent(watcher.logger, event)
			}

		case <-signals:
			watcher.logger.Info("stopping")
			atomic.StoreInt32(&stopEventSource, 1)
			if es := eventSource.Load(); es != nil {
				err := es.(receptor.EventSource).Close()
				if err != nil {
					watcher.logger.Error("failed-closing-event-source", err)
				}
			}
			return nil
		}
	}
}

func (watcher *Watcher) emit(logger lager.Logger) {
	messagesToEmit := watcher.table.MessagesToEmit()

	logger.Info("emitting-messages", lager.Data{"messages": messagesToEmit})
	err := watcher.emitter.Emit(messagesToEmit)
	if err != nil {
		logger.Error("failed-to-emit-routes", err)
	}

	routesSynced.Add(messagesToEmit.RouteRegistrationCount())
	routesTotal.Send(watcher.table.RouteCount())
}

func (watcher *Watcher) sync(logger lager.Logger, syncEndChan chan syncEndEvent) {
	endEvent := syncEndEvent{logger: logger}
	defer func() {
		syncEndChan <- endEvent
	}()

	before := watcher.clock.Now()

	actualLRPResponses, err := watcher.receptorClient.ActualLRPs()
	if err != nil {
		logger.Error("failed-to-get-actual", err)
		return
	}

	desiredLRPResponses, err := watcher.receptorClient.DesiredLRPs()
	if err != nil {
		logger.Error("failed-to-get-desired", err)
		return
	}

	runningActualLRPs := make([]receptor.ActualLRPResponse, 0, len(actualLRPResponses))
	for _, actualLRPResponse := range actualLRPResponses {
		if actualLRPResponse.State == receptor.ActualLRPStateRunning {
			runningActualLRPs = append(runningActualLRPs, actualLRPResponse)
		}
	}

	desiredLRPs := make([]receptor.DesiredLRPResponse, 0, len(desiredLRPResponses))
	for _, desiredLRPResponse := range desiredLRPResponses {
		desiredLRPs = append(desiredLRPs, desiredLRPResponse)
	}

	newTable := routing_table.NewTempTable(
		routing_table.RoutesByRoutingKeyFromDesireds(desiredLRPs),
		routing_table.EndpointsByRoutingKeyFromActuals(runningActualLRPs),
	)

	endEvent.table = newTable
	endEvent.callback = func(table routing_table.RoutingTable) {
		after := watcher.clock.Now()
		routeSyncDuration.Send(after.Sub(before))
	}
}

func (watcher *Watcher) completeSync(syncEnd syncEndEvent, cachedEvents map[string]receptor.Event) {
	logger := syncEnd.logger

	if syncEnd.table == nil {
		// sync failed, process the events on the current table
		for _, e := range cachedEvents {
			watcher.handleEvent(logger, e)
		}

		return
	}

	emitter := watcher.emitter
	watcher.emitter = nil

	table := watcher.table
	watcher.table = syncEnd.table

	for _, e := range cachedEvents {
		watcher.handleEvent(logger, e)
	}

	watcher.table = table
	watcher.emitter = emitter

	messages := watcher.table.Swap(syncEnd.table)
	watcher.emitMessages(logger, messages)

	if syncEnd.callback != nil {
		syncEnd.callback(watcher.table)
	}
}

func (watcher *Watcher) handleEvent(logger lager.Logger, event receptor.Event) {
	switch event := event.(type) {
	case receptor.DesiredLRPCreatedEvent:
		watcher.handleDesiredCreate(logger, event.DesiredLRPResponse)
	case receptor.DesiredLRPChangedEvent:
		watcher.handleDesiredUpdate(logger, event.Before, event.After)
	case receptor.DesiredLRPRemovedEvent:
		watcher.handleDesiredDelete(logger, event.DesiredLRPResponse)
	case receptor.ActualLRPCreatedEvent:
		watcher.handleActualCreate(logger, event.ActualLRPResponse)
	case receptor.ActualLRPChangedEvent:
		watcher.handleActualUpdate(logger, event.Before, event.After)
	case receptor.ActualLRPRemovedEvent:
		watcher.handleActualDelete(logger, event.ActualLRPResponse)
	default:
		logger.Info("did-not-handle-unrecognizable-event", lager.Data{"event-type": event.EventType()})
	}
}

func (watcher *Watcher) handleDesiredCreate(logger lager.Logger, desiredLRP receptor.DesiredLRPResponse) {
	logger = logger.Session("handle-desired-create", desiredLRPData(desiredLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	watcher.setRoutesForDesired(logger, desiredLRP)
}

func (watcher *Watcher) handleDesiredUpdate(logger lager.Logger, before, after receptor.DesiredLRPResponse) {
	logger = logger.Session("handling-desired-update", lager.Data{
		"before": desiredLRPData(before),
		"after":  desiredLRPData(after),
	})
	logger.Info("starting")
	defer logger.Info("complete")

	afterKeysSet := watcher.setRoutesForDesired(logger, after)

	beforeRoutingKeys := routing_table.RoutingKeysFromDesired(before)
	afterRoutes, _ := cfroutes.CFRoutesFromRoutingInfo(after.Routes)

	afterContainerPorts := set{}
	for _, route := range afterRoutes {
		afterContainerPorts.add(route.Port)
	}

	for _, key := range beforeRoutingKeys {
		if !afterKeysSet.contains(key) || !afterContainerPorts.contains(key.ContainerPort) {
			messagesToEmit := watcher.table.RemoveRoutes(key, after.ModificationTag)
			watcher.emitMessages(logger, messagesToEmit)
		}
	}
}

func (watcher *Watcher) setRoutesForDesired(logger lager.Logger, desiredLRP receptor.DesiredLRPResponse) set {
	routingKeys := routing_table.RoutingKeysFromDesired(desiredLRP)
	routes, _ := cfroutes.CFRoutesFromRoutingInfo(desiredLRP.Routes)
	routingKeySet := set{}

	for _, key := range routingKeys {
		routingKeySet.add(key)
		for _, route := range routes {
			if key.ContainerPort == route.Port {
				messagesToEmit := watcher.table.SetRoutes(key, routing_table.Routes{
					Hostnames: route.Hostnames,
					LogGuid:   desiredLRP.LogGuid,
				})
				watcher.emitMessages(logger, messagesToEmit)
			}
		}
	}

	return routingKeySet
}

func (watcher *Watcher) handleDesiredDelete(logger lager.Logger, desiredLRP receptor.DesiredLRPResponse) {
	logger = logger.Session("handling-desired-delete", desiredLRPData(desiredLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	for _, key := range routing_table.RoutingKeysFromDesired(desiredLRP) {
		messagesToEmit := watcher.table.RemoveRoutes(key, desiredLRP.ModificationTag)

		watcher.emitMessages(logger, messagesToEmit)
	}
}

func (watcher *Watcher) handleActualCreate(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	logger = logger.Session("handling-actual-create", actualLRPData(actualLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	if actualLRP.State == receptor.ActualLRPStateRunning {
		watcher.addAndEmit(logger, actualLRP)
	}
}

func (watcher *Watcher) handleActualUpdate(logger lager.Logger, before, after receptor.ActualLRPResponse) {
	logger = logger.Session("handling-actual-update", lager.Data{
		"before": actualLRPData(before),
		"after":  actualLRPData(after),
	})
	logger.Info("starting")
	defer logger.Info("complete")

	switch {
	case after.State == receptor.ActualLRPStateRunning:
		watcher.addAndEmit(logger, after)
	case after.State != receptor.ActualLRPStateRunning && before.State == receptor.ActualLRPStateRunning:
		watcher.removeAndEmit(logger, before)
	}
}

func (watcher *Watcher) handleActualDelete(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	logger = logger.Session("handling-actual-delete", actualLRPData(actualLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	if actualLRP.State == receptor.ActualLRPStateRunning {
		watcher.removeAndEmit(logger, actualLRP)
	}
}

func (watcher *Watcher) addAndEmit(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	endpoints, err := routing_table.EndpointsFromActual(actualLRP)
	if err != nil {
		logger.Error("failed-to-extract-endpoint-from-actual", err)
		return
	}

	for _, key := range routing_table.RoutingKeysFromActual(actualLRP) {
		for _, endpoint := range endpoints {
			if key.ContainerPort == endpoint.ContainerPort {
				messagesToEmit := watcher.table.AddEndpoint(key, endpoint)
				watcher.emitMessages(logger, messagesToEmit)
			}
		}
	}
}

func (watcher *Watcher) removeAndEmit(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	endpoints, err := routing_table.EndpointsFromActual(actualLRP)
	if err != nil {
		logger.Error("failed-to-extract-endpoint-from-actual", err)
		return
	}

	for _, key := range routing_table.RoutingKeysFromActual(actualLRP) {
		for _, endpoint := range endpoints {
			if key.ContainerPort == endpoint.ContainerPort {
				messagesToEmit := watcher.table.RemoveEndpoint(key, endpoint)
				watcher.emitMessages(logger, messagesToEmit)
			}
		}
	}
}

func (watcher *Watcher) emitMessages(logger lager.Logger, messagesToEmit routing_table.MessagesToEmit) {
	if watcher.emitter != nil {
		logger.Info("emitting-messages", lager.Data{"messages": messagesToEmit})
		watcher.emitter.Emit(messagesToEmit)
		routesRegistered.Add(messagesToEmit.RouteRegistrationCount())
		routesUnregistered.Add(messagesToEmit.RouteUnregistrationCount())
	}
}

func desiredLRPData(lrp receptor.DesiredLRPResponse) lager.Data {
	return lager.Data{
		"process-guid": lrp.ProcessGuid,
		"routes":       lrp.Routes,
		"ports":        lrp.Ports,
	}
}

func actualLRPData(lrp receptor.ActualLRPResponse) lager.Data {
	return lager.Data{
		"process-guid":  lrp.ProcessGuid,
		"index":         lrp.Index,
		"domain":        lrp.Domain,
		"instance-guid": lrp.InstanceGuid,
		"cell-id":       lrp.CellID,
		"address":       lrp.Address,
		"ports":         lrp.Ports,
		"evacuating":    lrp.Evacuating,
	}
}
<commit_msg>Add logging to watcher [#90332636]<commit_after>package watcher

import (
	"os"
	"sync/atomic"

	"github.com/cloudfoundry-incubator/receptor"
	"github.com/cloudfoundry-incubator/route-emitter/cfroutes"
	"github.com/cloudfoundry-incubator/route-emitter/nats_emitter"
	"github.com/cloudfoundry-incubator/route-emitter/routing_table"
	"github.com/cloudfoundry-incubator/route-emitter/syncer"
	"github.com/cloudfoundry-incubator/runtime-schema/metric"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
)

var (
	routesTotal  = metric.Metric("RoutesTotal")
	routesSynced = metric.Counter("RoutesSynced")

	routeSyncDuration = metric.Duration("RouteEmitterSyncDuration")

	routesRegistered   = metric.Counter("RoutesRegistered")
	routesUnregistered = metric.Counter("RoutesUnregistered")
)

type Watcher struct {
	receptorClient receptor.Client
	clock          clock.Clock
	table          routing_table.RoutingTable
	emitter        nats_emitter.NATSEmitter
	syncEvents     syncer.Events
	logger         lager.Logger
}

type syncEndEvent struct {
	table    routing_table.RoutingTable
	callback func(routing_table.RoutingTable)

	logger lager.Logger
}

type set map[interface{}]struct{}

func (set set) contains(value interface{}) bool {
	_, found := set[value]
	return found
}

func (set set) add(value interface{}) {
	set[value] = struct{}{}
}

func NewWatcher(
	receptorClient receptor.Client,
	clock clock.Clock,
	table routing_table.RoutingTable,
	emitter nats_emitter.NATSEmitter,
	syncEvents syncer.Events,
	logger lager.Logger,
) *Watcher {
	return &Watcher{
		receptorClient: receptorClient,
		clock:          clock,
		table:          table,
		emitter:        emitter,
		syncEvents:     syncEvents,
		logger:         logger.Session("watcher"),
	}
}

func (watcher *Watcher) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	watcher.logger.Info("starting")

	close(ready)
	watcher.logger.Info("started")
	defer watcher.logger.Info("finished")

	var cachedEvents map[string]receptor.Event

	eventChan := make(chan receptor.Event)
	syncEndChan := make(chan syncEndEvent)

	syncing := false

	var eventSource atomic.Value
	var stopEventSource int32

	startEventSource := func() {
		go func() {
			var err error
			var es receptor.EventSource

			for {
				if atomic.LoadInt32(&stopEventSource) == 1 {
					return
				}

				es, err = watcher.receptorClient.SubscribeToEvents()
				if err != nil {
					watcher.logger.Error("failed-subscribing-to-events", err)
					continue
				}

				eventSource.Store(es)

				var event receptor.Event
				for {
					event, err = es.Next()
					if err != nil {
						watcher.logger.Error("failed-getting-next-event", err)
						break
					}

					if event != nil {
						eventChan <- event
					}
				}
			}
		}()
	}

	startedEventSource := false
	for {
		select {
		case <-watcher.syncEvents.Sync:
			if syncing == false {
				logger := watcher.logger.Session("sync")
				logger.Info("starting")
				syncing = true

				if !startedEventSource {
					startedEventSource = true
					startEventSource()
				}

				cachedEvents = make(map[string]receptor.Event)
				go watcher.sync(logger, syncEndChan)
			}

		case syncEnd := <-syncEndChan:
			watcher.completeSync(syncEnd, cachedEvents)
			cachedEvents = nil
			syncing = false
			syncEnd.logger.Info("complete")

		case <-watcher.syncEvents.Emit:
			logger := watcher.logger.Session("emit")
			watcher.emit(logger)

		case event := <-eventChan:
			if syncing {
				watcher.logger.Info("caching-event", lager.Data{
					"type": event.EventType(),
				})

				cachedEvents[event.Key()] = event
			} else {
				watcher.handleEvent(watcher.logger, event)
			}

		case <-signals:
			watcher.logger.Info("stopping")
			atomic.StoreInt32(&stopEventSource, 1)
			if es := eventSource.Load(); es != nil {
				err := es.(receptor.EventSource).Close()
				if err != nil {
					watcher.logger.Error("failed-closing-event-source", err)
				}
			}
			return nil
		}
	}
}

func (watcher *Watcher) emit(logger lager.Logger) {
	messagesToEmit := watcher.table.MessagesToEmit()

	logger.Info("emitting-messages", lager.Data{"messages": messagesToEmit})
	err := watcher.emitter.Emit(messagesToEmit)
	if err != nil {
		logger.Error("failed-to-emit-routes", err)
	}

	routesSynced.Add(messagesToEmit.RouteRegistrationCount())
	routesTotal.Send(watcher.table.RouteCount())
}

func (watcher *Watcher) sync(logger lager.Logger, syncEndChan chan syncEndEvent) {
	endEvent := syncEndEvent{logger: logger}
	defer func() {
		syncEndChan <- endEvent
	}()

	before := watcher.clock.Now()

	logger.Debug("getting-actual-lrps")
	actualLRPResponses, err := watcher.receptorClient.ActualLRPs()
	if err != nil {
		logger.Error("failed-getting-actual-lrps", err)
		return
	}
	logger.Debug("succeeded-getting-actual-lrps", lager.Data{"num-actual-responses": len(actualLRPResponses)})

	logger.Debug("getting-desired-lrps")
	desiredLRPResponses, err := watcher.receptorClient.DesiredLRPs()
	if err != nil {
		logger.Error("failed-getting-desired-lrps", err)
		return
	}
	logger.Debug("succeeded-getting-desired-lrps", lager.Data{"num-desired-responses": len(desiredLRPResponses)})

	runningActualLRPs := make([]receptor.ActualLRPResponse, 0, len(actualLRPResponses))
	for _, actualLRPResponse := range actualLRPResponses {
		if actualLRPResponse.State == receptor.ActualLRPStateRunning {
			runningActualLRPs = append(runningActualLRPs, actualLRPResponse)
		}
	}

	desiredLRPs := make([]receptor.DesiredLRPResponse, 0, len(desiredLRPResponses))
	for _, desiredLRPResponse := range desiredLRPResponses {
		desiredLRPs = append(desiredLRPs, desiredLRPResponse)
	}

	newTable := routing_table.NewTempTable(
		routing_table.RoutesByRoutingKeyFromDesireds(desiredLRPs),
		routing_table.EndpointsByRoutingKeyFromActuals(runningActualLRPs),
	)

	endEvent.table = newTable
	endEvent.callback = func(table routing_table.RoutingTable) {
		after := watcher.clock.Now()
		routeSyncDuration.Send(after.Sub(before))
	}
}

func (watcher *Watcher) completeSync(syncEnd syncEndEvent, cachedEvents map[string]receptor.Event) {
	logger := syncEnd.logger

	if syncEnd.table == nil {
		// sync failed, process the events on the current table
		logger.Debug("handling-events-from-failed-sync")
		for _, e := range cachedEvents {
			watcher.handleEvent(logger, e)
		}
		logger.Debug("done-handling-events-from-failed-sync")

		return
	}

	emitter := watcher.emitter
	watcher.emitter = nil

	table := watcher.table
	watcher.table = syncEnd.table

	logger.Debug("handling-cached-events")
	for _, e := range cachedEvents {
		watcher.handleEvent(logger, e)
	}
	logger.Debug("done-handling-cached-events")

	watcher.table = table
	watcher.emitter = emitter

	messages := watcher.table.Swap(syncEnd.table)
	logger.Debug("emitting-messages", lager.Data{
		"num-registration-messages":   len(messages.RegistrationMessages),
		"num-unregistration-messages": len(messages.UnregistrationMessages),
	})
	watcher.emitMessages(logger, messages)
	logger.Debug("done-emitting-messages", lager.Data{
		"num-registration-messages":   len(messages.RegistrationMessages),
		"num-unregistration-messages": len(messages.UnregistrationMessages),
	})

	if syncEnd.callback != nil {
		syncEnd.callback(watcher.table)
	}
}

func (watcher *Watcher) handleEvent(logger lager.Logger, event receptor.Event) {
	switch event := event.(type) {
	case receptor.DesiredLRPCreatedEvent:
		watcher.handleDesiredCreate(logger, event.DesiredLRPResponse)
	case receptor.DesiredLRPChangedEvent:
		watcher.handleDesiredUpdate(logger, event.Before, event.After)
	case receptor.DesiredLRPRemovedEvent:
		watcher.handleDesiredDelete(logger, event.DesiredLRPResponse)
	case receptor.ActualLRPCreatedEvent:
		watcher.handleActualCreate(logger, event.ActualLRPResponse)
	case receptor.ActualLRPChangedEvent:
		watcher.handleActualUpdate(logger, event.Before, event.After)
	case receptor.ActualLRPRemovedEvent:
		watcher.handleActualDelete(logger, event.ActualLRPResponse)
	default:
		logger.Info("did-not-handle-unrecognizable-event", lager.Data{"event-type": event.EventType()})
	}
}

func (watcher *Watcher) handleDesiredCreate(logger lager.Logger, desiredLRP receptor.DesiredLRPResponse) {
	logger = logger.Session("handle-desired-create", desiredLRPData(desiredLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	watcher.setRoutesForDesired(logger, desiredLRP)
}

func (watcher *Watcher) handleDesiredUpdate(logger lager.Logger, before, after receptor.DesiredLRPResponse) {
	logger = logger.Session("handling-desired-update", lager.Data{
		"before": desiredLRPData(before),
		"after":  desiredLRPData(after),
	})
	logger.Info("starting")
	defer logger.Info("complete")

	afterKeysSet := watcher.setRoutesForDesired(logger, after)

	beforeRoutingKeys := routing_table.RoutingKeysFromDesired(before)
	afterRoutes, _ := cfroutes.CFRoutesFromRoutingInfo(after.Routes)

	afterContainerPorts := set{}
	for _, route := range afterRoutes {
		afterContainerPorts.add(route.Port)
	}

	for _, key := range beforeRoutingKeys {
		if !afterKeysSet.contains(key) || !afterContainerPorts.contains(key.ContainerPort) {
			messagesToEmit := watcher.table.RemoveRoutes(key, after.ModificationTag)
			watcher.emitMessages(logger, messagesToEmit)
		}
	}
}

func (watcher *Watcher) setRoutesForDesired(logger lager.Logger, desiredLRP receptor.DesiredLRPResponse) set {
	routingKeys := routing_table.RoutingKeysFromDesired(desiredLRP)
	routes, _ := cfroutes.CFRoutesFromRoutingInfo(desiredLRP.Routes)
	routingKeySet := set{}

	for _, key := range routingKeys {
		routingKeySet.add(key)
		for _, route := range routes {
			if key.ContainerPort == route.Port {
				messagesToEmit := watcher.table.SetRoutes(key, routing_table.Routes{
					Hostnames: route.Hostnames,
					LogGuid:   desiredLRP.LogGuid,
				})
				watcher.emitMessages(logger, messagesToEmit)
			}
		}
	}

	return routingKeySet
}

func (watcher *Watcher) handleDesiredDelete(logger lager.Logger, desiredLRP receptor.DesiredLRPResponse) {
	logger = logger.Session("handling-desired-delete", desiredLRPData(desiredLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	for _, key := range routing_table.RoutingKeysFromDesired(desiredLRP) {
		messagesToEmit := watcher.table.RemoveRoutes(key, desiredLRP.ModificationTag)

		watcher.emitMessages(logger, messagesToEmit)
	}
}

func (watcher *Watcher) handleActualCreate(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	logger = logger.Session("handling-actual-create", actualLRPData(actualLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	if actualLRP.State == receptor.ActualLRPStateRunning {
		watcher.addAndEmit(logger, actualLRP)
	}
}

func (watcher *Watcher) handleActualUpdate(logger lager.Logger, before, after receptor.ActualLRPResponse) {
	logger = logger.Session("handling-actual-update", lager.Data{
		"before": actualLRPData(before),
		"after":  actualLRPData(after),
	})
	logger.Info("starting")
	defer logger.Info("complete")

	switch {
	case after.State == receptor.ActualLRPStateRunning:
		watcher.addAndEmit(logger, after)
	case after.State != receptor.ActualLRPStateRunning && before.State == receptor.ActualLRPStateRunning:
		watcher.removeAndEmit(logger, before)
	}
}

func (watcher *Watcher) handleActualDelete(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	logger = logger.Session("handling-actual-delete", actualLRPData(actualLRP))
	logger.Info("starting")
	defer logger.Info("complete")

	if actualLRP.State == receptor.ActualLRPStateRunning {
		watcher.removeAndEmit(logger, actualLRP)
	}
}

func (watcher *Watcher) addAndEmit(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	endpoints, err := routing_table.EndpointsFromActual(actualLRP)
	if err != nil {
		logger.Error("failed-to-extract-endpoint-from-actual", err)
		return
	}

	for _, key := range routing_table.RoutingKeysFromActual(actualLRP) {
		for _, endpoint := range endpoints {
			if key.ContainerPort == endpoint.ContainerPort {
				messagesToEmit := watcher.table.AddEndpoint(key, endpoint)
				watcher.emitMessages(logger, messagesToEmit)
			}
		}
	}
}

func (watcher *Watcher) removeAndEmit(logger lager.Logger, actualLRP receptor.ActualLRPResponse) {
	endpoints, err := routing_table.EndpointsFromActual(actualLRP)
	if err != nil {
		logger.Error("failed-to-extract-endpoint-from-actual", err)
		return
	}

	for _, key := range routing_table.RoutingKeysFromActual(actualLRP) {
		for _, endpoint := range endpoints {
			if key.ContainerPort == endpoint.ContainerPort {
				messagesToEmit := watcher.table.RemoveEndpoint(key, endpoint)
				watcher.emitMessages(logger, messagesToEmit)
			}
		}
	}
}

func (watcher *Watcher) emitMessages(logger lager.Logger, messagesToEmit routing_table.MessagesToEmit) {
	if watcher.emitter != nil {
		logger.Info("emitting-messages", lager.Data{"messages": messagesToEmit})
		watcher.emitter.Emit(messagesToEmit)
		routesRegistered.Add(messagesToEmit.RouteRegistrationCount())
		routesUnregistered.Add(messagesToEmit.RouteUnregistrationCount())
	}
}

func desiredLRPData(lrp receptor.DesiredLRPResponse) lager.Data {
	return lager.Data{
		"process-guid": lrp.ProcessGuid,
		"routes":       lrp.Routes,
		"ports":        lrp.Ports,
	}
}

func actualLRPData(lrp receptor.ActualLRPResponse) lager.Data {
	return lager.Data{
		"process-guid":  lrp.ProcessGuid,
		"index":         lrp.Index,
		"domain":        lrp.Domain,
		"instance-guid": lrp.InstanceGuid,
		"cell-id":       lrp.CellID,
		"address":       lrp.Address,
		"ports":         lrp.Ports,
		"evacuating":    lrp.Evacuating,
	}
}
<|endoftext|>


<commit_before>package natsx

import (
	"context"

	"github.com/altairsix/pkg/tracer"
	"github.com/altairsix/pkg/tracer/k"
	"github.com/savaki/nats-protobuf"
)

// Log provides the standardized logging service for nats queries
func Log(bc string) nats_protobuf.Filter {
	return func(fn nats_protobuf.HandlerFunc) nats_protobuf.HandlerFunc {
		return func(ctx context.Context, m *nats_protobuf.Message) (*nats_protobuf.Message, error) {
			segment, ctx := tracer.NewSegment(ctx, "nats.api",
				k.String("bc", bc),
				k.String("method", m.Method),
			)
			defer segment.Finish()

			out, err := fn(ctx, m)
			if err != nil {
				segment.LogFields(k.Err(err))
				return nil, err
			}

			return out, nil
		}
	}
}
<commit_msg>- log method is now provided the subject via the api<commit_after>package natsx

import (
	"context"

	"github.com/altairsix/pkg/tracer"
	"github.com/altairsix/pkg/tracer/k"
	"github.com/savaki/nats-protobuf"
)

// Log provides the standardized logging service for nats queries
func Log(bc string) nats_protobuf.Filter {
	return func(fn nats_protobuf.HandlerFunc) nats_protobuf.HandlerFunc {
		return func(ctx context.Context, subject string, m *nats_protobuf.Message) (*nats_protobuf.Message, error) {
			segment, ctx := tracer.NewSegment(ctx, "nats.api",
				k.String("subject", subject),
				k.String("bc", bc),
				k.String("method", m.Method),
			)
			defer segment.Finish()

			out, err := fn(ctx, subject, m)
			if err != nil {
				segment.LogFields(k.Err(err))
				return nil, err
			}

			return out, nil
		}
	}
}
<|endoftext|>


<commit_before>package org

import (
	"errors"
	"strings"

	"github.com/appcelerator/amp/api/rpc/account"
	"github.com/appcelerator/amp/cli"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// NewOrgRemoveCommand returns a new instance of the remove organization command.
func NewOrgRemoveCommand(c cli.Interface) *cobra.Command {
	return &cobra.Command{
		Use:     "rm ORGANIZATION(S)",
		Short:   "Remove organization",
		Aliases: []string{"remove"},
		PreRunE: cli.AtLeastArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return removeOrg(c, args)
		},
	}
}

func removeOrg(c cli.Interface, args []string) error {
	var errs []string
	conn := c.ClientConn()
	client := account.NewAccountClient(conn)
	for _, org := range args {
		request := &account.DeleteOrganizationRequest{
			Name: org,
		}
		if _, err := client.DeleteOrganization(context.Background(), request); err != nil {
			errs = append(errs, grpc.ErrorDesc(err))
			continue
		}
		c.Console().Println(org)
	}
	if len(errs) > 0 {
		return errors.New(strings.Join(errs, "\n"))
	}
	return nil
}
<commit_msg>org rm will switch user back to logged in account (#1420)<commit_after>package org

import (
	"errors"
	"fmt"
	"strings"

	"github.com/appcelerator/amp/api/rpc/account"
	"github.com/appcelerator/amp/api/auth"
	"github.com/appcelerator/amp/cli"
	"github.com/dgrijalva/jwt-go"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// NewOrgRemoveCommand returns a new instance of the remove organization command.
func NewOrgRemoveCommand(c cli.Interface) *cobra.Command {
	return &cobra.Command{
		Use:     "rm ORGANIZATION(S)",
		Short:   "Remove organization",
		Aliases: []string{"remove"},
		PreRunE: cli.AtLeastArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return removeOrg(c, args)
		},
	}
}

func removeOrg(c cli.Interface, args []string) error {
	var errs []string
	conn := c.ClientConn()
	client := account.NewAccountClient(conn)

	for _, org := range args {
		requestDelete := &account.DeleteOrganizationRequest{
			Name: org,
		}
		if _, err := client.DeleteOrganization(context.Background(), requestDelete); err != nil {
			errs = append(errs, grpc.ErrorDesc(err))
			continue
		}
		c.Console().Println(org)
		// Check if user is logged in on behalf of an org
		token, err := cli.ReadToken(c.Server())
		if err != nil {
			return err
		}
		pToken, _ := jwt.ParseWithClaims(token, &auth.AuthClaims{}, func(t *jwt.Token) (interface{}, error) {
			return []byte{}, nil
		})
		if claims, ok := pToken.Claims.(*auth.AuthClaims); ok {
			if claims.ActiveOrganization == org {
				// Switch back to the users account
				client := account.NewAccountClient(conn)
				requestSwitch := &account.SwitchRequest{
					Account: claims.AccountName,
				}
				headers := metadata.MD{}
				_, err := client.Switch(context.Background(), requestSwitch, grpc.Header(&headers))
				if err != nil {
					return fmt.Errorf("%s", grpc.ErrorDesc(err))
				}
				if err := cli.SaveToken(headers, c.Server()); err != nil {
					return fmt.Errorf("%s", grpc.ErrorDesc(err))
				}
				c.Console().Printf("Switched back to account: %s\n", claims.AccountName)
			}
		}
	}
	if len(errs) > 0 {
		return errors.New(strings.Join(errs, "\n"))
	}
	return nil
}
<|endoftext|>


<commit_before>//
// Copyright 2016 Gregory Trubetskoy. 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.

package cluster

import (
	"bytes"
	"compress/flate"
	"encoding/gob"
	"fmt"
	"github.com/hashicorp/memberlist"
	"log"
	"sort"
	"strings"
)

type Cluster struct {
	*memberlist.Memberlist
	d *delegate
}

func (c *Cluster) Create() error {
	return c.CreateBind("", 0)
}

func (c *Cluster) CreateBind(addr string, port int) (err error) {
	cfg := memberlist.DefaultLANConfig()
	cfg.BindAddr, cfg.AdvertiseAddr = addr, addr
	cfg.BindPort, cfg.AdvertisePort = port, port
	cfg.Name = fmt.Sprintf("%s:%d", addr, port)
	cfg.LogOutput = &logger{}
	c.d = &delegate{make([]chan *Msg, 0)}
	cfg.Delegate = c.d
	c.Memberlist, err = memberlist.Create(cfg)
	return err
}

func (c *Cluster) Join(existing []string) error {
	// NB: You can always join yourself
	if n, err := c.Memberlist.Join(existing); n <= 0 {
		return err
	}
	return nil
}

func (c *Cluster) SortedNodes() []*memberlist.Node {
	// Return in order of addr:port
	sn := sortableNodes{c.Members()}
	sort.Sort(sn)
	return sn.nl
}

type sortableNodes struct {
	nl []*memberlist.Node
}

func (sn sortableNodes) Len() int {
	return len(sn.nl)
}

func (sn sortableNodes) Less(i, j int) bool {
	istr := fmt.Sprintf("%s:%s", sn.nl[i].Addr, sn.nl[i].Port)
	jstr := fmt.Sprintf("%s:%s", sn.nl[j].Addr, sn.nl[i].Port)
	return istr < jstr
}

func (sn sortableNodes) Swap(i, j int) {
	sn.nl[i], sn.nl[j] = sn.nl[j], sn.nl[i]
}

func (c *Cluster) RegisterMsgType() (snd, rcv chan *Msg) {

	snd, rcv = make(chan *Msg, 16), make(chan *Msg, 16)

	d := c.d
	d.rcvChs = append(d.rcvChs, rcv)
	id := len(d.rcvChs) - 1

	go func(id int) {
		for {
			msg := <-snd
			msg.Src = c.LocalNode()
			msg.Id = id
			c.SendToTCP(msg.Dst, msg.bytes())
		}
	}(id)

	return snd, rcv
}

type Msg struct {
	Id       int
	Dst, Src *memberlist.Node
	Body     []byte
}

func NewMsg(dest *memberlist.Node, payload gob.GobEncoder) (*Msg, error) {
	var buf bytes.Buffer
	enc := gob.NewEncoder(&buf)
	if err := enc.Encode(payload); err != nil {
		return nil, err
	}
	return &Msg{Dst: dest, Body: buf.Bytes()}, nil
}

func (m *Msg) bytes() []byte {
	var buf bytes.Buffer
	z, _ := flate.NewWriter(&buf, -1)
	enc := gob.NewEncoder(z)
	if err := enc.Encode(m); err != nil {
		log.Printf("Msg.bytes(): Error encountered in encoding: %v", err)
		return nil
	}
	z.Close()
	return buf.Bytes()
}

func (m *Msg) Decode(dst interface{}) error {
	if err := gob.NewDecoder(bytes.NewBuffer(m.Body)).Decode(dst); err != nil {
		log.Printf("Msg.Decode() decoding error: %v", err)
		return err
	}
	return nil
}

type logger struct{}

// Ignore [DEBUG]
func (l *logger) Write(b []byte) (int, error) {
	s := string(b)
	if !strings.Contains(s, "[DEBUG]") {
		log.Printf(s)
	}
	return len(s), nil
}

// memberlist.Delegate interface
type delegate struct {
	rcvChs []chan *Msg
}

func (d *delegate) NodeMeta(limit int) []byte {
	return []byte{}
}

func (d *delegate) NotifyMsg(b []byte) {

	m := &Msg{}
	if err := gob.NewDecoder(flate.NewReader(bytes.NewBuffer(b))).Decode(m); err != nil {
		log.Printf("NotifyMsg(): error decoding: %#v", err)
	}

	if m.Id < len(d.rcvChs) {
		d.rcvChs[m.Id] <- m
	} else {
		log.Printf("NotifyMsg(): unknown msg Id: %d, dropping message", m.Id)
	}
}

func (d *delegate) GetBroadcasts(overhead, limit int) [][]byte {
	return [][]byte{}
}

func (d *delegate) LocalState(join bool) []byte {
	return []byte{}
}

func (d *delegate) MergeRemoteState(buf []byte, join bool) {}
<commit_msg>Obey the MLBIND var<commit_after>//
// Copyright 2016 Gregory Trubetskoy. 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.

package cluster

import (
	"bytes"
	"compress/flate"
	"encoding/gob"
	"fmt"
	"github.com/hashicorp/memberlist"
	"log"
	"os"
	"sort"
	"strings"
)

type Cluster struct {
	*memberlist.Memberlist
	d *delegate
}

func (c *Cluster) Create() error {
	return c.CreateBind(os.Getenv("MLBIND"), 0)
}

func (c *Cluster) CreateBind(addr string, port int) (err error) {
	cfg := memberlist.DefaultLANConfig()
	if addr != "" {
		cfg.BindAddr, cfg.AdvertiseAddr = addr, addr
	}
	if port != 0 {
		cfg.BindPort, cfg.AdvertisePort = port, port
	}
	cfg.Name = fmt.Sprintf("%s:%d", addr, port)
	cfg.LogOutput = &logger{}
	c.d = &delegate{make([]chan *Msg, 0)}
	cfg.Delegate = c.d
	c.Memberlist, err = memberlist.Create(cfg)
	return err
}

func (c *Cluster) Join(existing []string) error {
	// NB: You can always join yourself
	if n, err := c.Memberlist.Join(existing); n <= 0 {
		return err
	}
	return nil
}

func (c *Cluster) SortedNodes() []*memberlist.Node {
	// Return in order of addr:port
	sn := sortableNodes{c.Members()}
	sort.Sort(sn)
	return sn.nl
}

type sortableNodes struct {
	nl []*memberlist.Node
}

func (sn sortableNodes) Len() int {
	return len(sn.nl)
}

func (sn sortableNodes) Less(i, j int) bool {
	istr := fmt.Sprintf("%s:%s", sn.nl[i].Addr, sn.nl[i].Port)
	jstr := fmt.Sprintf("%s:%s", sn.nl[j].Addr, sn.nl[i].Port)
	return istr < jstr
}

func (sn sortableNodes) Swap(i, j int) {
	sn.nl[i], sn.nl[j] = sn.nl[j], sn.nl[i]
}

func (c *Cluster) RegisterMsgType() (snd, rcv chan *Msg) {

	snd, rcv = make(chan *Msg, 16), make(chan *Msg, 16)

	d := c.d
	d.rcvChs = append(d.rcvChs, rcv)
	id := len(d.rcvChs) - 1

	go func(id int) {
		for {
			msg := <-snd
			msg.Src = c.LocalNode()
			msg.Id = id
			c.SendToTCP(msg.Dst, msg.bytes())
		}
	}(id)

	return snd, rcv
}

type Msg struct {
	Id       int
	Dst, Src *memberlist.Node
	Body     []byte
}

func NewMsg(dest *memberlist.Node, payload gob.GobEncoder) (*Msg, error) {
	var buf bytes.Buffer
	enc := gob.NewEncoder(&buf)
	if err := enc.Encode(payload); err != nil {
		return nil, err
	}
	return &Msg{Dst: dest, Body: buf.Bytes()}, nil
}

func (m *Msg) bytes() []byte {
	var buf bytes.Buffer
	z, _ := flate.NewWriter(&buf, -1)
	enc := gob.NewEncoder(z)
	if err := enc.Encode(m); err != nil {
		log.Printf("Msg.bytes(): Error encountered in encoding: %v", err)
		return nil
	}
	z.Close()
	return buf.Bytes()
}

func (m *Msg) Decode(dst interface{}) error {
	if err := gob.NewDecoder(bytes.NewBuffer(m.Body)).Decode(dst); err != nil {
		log.Printf("Msg.Decode() decoding error: %v", err)
		return err
	}
	return nil
}

type logger struct{}

// Ignore [DEBUG]
func (l *logger) Write(b []byte) (int, error) {
	s := string(b)
	if !strings.Contains(s, "[DEBUG]") {
		log.Printf(s)
	}
	return len(s), nil
}

// memberlist.Delegate interface
type delegate struct {
	rcvChs []chan *Msg
}

func (d *delegate) NodeMeta(limit int) []byte {
	return []byte{}
}

func (d *delegate) NotifyMsg(b []byte) {

	m := &Msg{}
	if err := gob.NewDecoder(flate.NewReader(bytes.NewBuffer(b))).Decode(m); err != nil {
		log.Printf("NotifyMsg(): error decoding: %#v", err)
	}

	if m.Id < len(d.rcvChs) {
		d.rcvChs[m.Id] <- m
	} else {
		log.Printf("NotifyMsg(): unknown msg Id: %d, dropping message", m.Id)
	}
}

func (d *delegate) GetBroadcasts(overhead, limit int) [][]byte {
	return [][]byte{}
}

func (d *delegate) LocalState(join bool) []byte {
	return []byte{}
}

func (d *delegate) MergeRemoteState(buf []byte, join bool) {}
<|endoftext|>


<commit_before>package cluster

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

	"github.com/hashicorp/memberlist"
	"github.com/raintank/worldping-api/pkg/log"
)

type ClusterManager struct {
	sync.RWMutex
	Peers map[string]Node
	node  Node
	list  *memberlist.Memberlist
}

func (c *ClusterManager) SetList(list *memberlist.Memberlist) {
	c.Lock()
	c.list = list
	c.Unlock()
}

func (c *ClusterManager) ThisNode() Node {
	c.RLock()
	defer c.RUnlock()
	return c.node
}

func (c *ClusterManager) PeersList() []Node {
	c.RLock()
	list := make([]Node, len(c.Peers))
	i := 0
	for _, n := range c.Peers {
		list[i] = n
		i++
	}
	c.RUnlock()
	return list
}

func (c *ClusterManager) NotifyJoin(node *memberlist.Node) {
	c.Lock()
	defer c.Unlock()
	if len(node.Meta) == 0 {
		return
	}
	log.Info("Node %s with address %s has joined the cluster\n", node.Name, node.Addr.String())
	peer := Node{}
	err := json.Unmarshal(node.Meta, &peer)
	if err != nil {
		panic(err)
	}
	peer.RemoteAddr = node.Addr.String()
	if peer.Name == c.node.Name {
		peer.local = true
	}
	c.Peers[node.Name] = peer
}

func (c *ClusterManager) NotifyLeave(node *memberlist.Node) {
	c.Lock()
	defer c.Unlock()
	log.Info("Node %s has left the cluster\n", node.Name)
	delete(c.Peers, node.Name)
}

func (c *ClusterManager) NotifyUpdate(node *memberlist.Node) {
	c.Lock()
	defer c.Unlock()
	if len(node.Meta) == 0 {
		return
	}
	peer := Node{}
	err := json.Unmarshal(node.Meta, &peer)
	if err != nil {
		panic(err)
	}
	peer.RemoteAddr = node.Addr.String()
	if peer.Name == c.node.Name {
		peer.local = true
	}
	c.Peers[node.Name] = peer
	log.Info("Node %s at %s has been updated - %s\n", node.Name, node.Addr.String(), node.Meta)
}

func (c *ClusterManager) BroadcastUpdate() {
	if c.list != nil {
		go c.list.UpdateNode(time.Second)
	}
}

// NodeMeta is used to retrieve meta-data about the current node
// when broadcasting an alive message. It's length is limited to
// the given byte size. This metadata is available in the Node structure.
func (c *ClusterManager) NodeMeta(limit int) []byte {
	c.Lock()
	meta, err := json.Marshal(c.node)
	c.Unlock()
	if err != nil {
		log.Fatal(4, err.Error())
	}
	return meta
}

// NotifyMsg is called when a user-data message is received.
// Care should be taken that this method does not block, since doing
// so would block the entire UDP packet receive loop. Additionally, the byte
// slice may be modified after the call returns, so it should be copied if needed.
func (c *ClusterManager) NotifyMsg(buf []byte) {
	// we dont have any need for passing messages between nodes, other then
	// the NodeMeta sent with alive messages.
	return
}

// GetBroadcasts is called when user data messages can be broadcast.
// It can return a list of buffers to send. Each buffer should assume an
// overhead as provided with a limit on the total byte size allowed.
// The total byte size of the resulting data to send must not exceed
// the limit. Care should be taken that this method does not block,
// since doing so would block the entire UDP packet receive loop.
func (c *ClusterManager) GetBroadcasts(overhead, limit int) [][]byte {
	// we dont have any need for passing messages between nodes, other then
	// the NodeMeta sent with alive messages.
	return nil
}

// LocalState is used for a TCP Push/Pull. This is sent to
// the remote side in addition to the membership information. Any
// data can be sent here. See MergeRemoteState as well. The `join`
// boolean indicates this is for a join instead of a push/pull.
func (c *ClusterManager) LocalState(join bool) []byte {
	c.Lock()
	c.Peers[c.node.Name] = c.node
	meta, err := json.Marshal(c.Peers)
	c.Unlock()
	if err != nil {
		log.Fatal(4, err.Error())
	}
	return meta
}

func (c *ClusterManager) MergeRemoteState(buf []byte, join bool) {
	knownPeers := make(map[string]Node)
	err := json.Unmarshal(buf, &knownPeers)
	if err != nil {
		log.Fatal(4, err.Error())
	}
	c.Lock()
	for name, meta := range knownPeers {
		if existing, ok := c.Peers[name]; ok {
			if meta.Updated.After(existing.Updated) {
				log.Info("updated node meta found in state update for %s", meta.Name)
				c.Peers[name] = meta
			}
		} else {
			log.Info("new node found in state update. %s", meta.Name)
			c.Peers[name] = meta
		}
	}
	c.Unlock()
}

// Returns true if the node is a ready to accept requests
// from users.
func (c *ClusterManager) IsReady() bool {
	c.RLock()
	defer c.RUnlock()
	return c.node.IsReady()
}

func (c *ClusterManager) SetReady() {
	c.Lock()
	c.node.State = NodeReady
	c.node.Updated = time.Now()
	c.Unlock()
	c.BroadcastUpdate()
}

func (c *ClusterManager) SetReadyIn(t time.Duration) {
	go func() {
		// wait for warmupPeriod before marking ourselves
		// as ready.
		time.Sleep(t)
		c.SetReady()
	}()
}

// Returns true if the node is a set as a primary node that should write data to cassandra.
func (c *ClusterManager) IsPrimary() bool {
	c.RLock()
	defer c.RUnlock()
	return c.node.Primary
}

// SetPrimary sets the primary status.
// Note: since we set the primary metric here, this should only be called on ThisNode !
func (c *ClusterManager) SetPrimary(p bool) {
	c.Lock()
	c.node.Primary = p
	c.node.PrimaryChange = time.Now()
	c.node.Updated = time.Now()
	c.Unlock()
	clusterPrimary.Set(p)
	c.BroadcastUpdate()
}

func (c *ClusterManager) SetPartitions(part []int32) {
	c.Lock()
	c.node.Partitions = part
	c.node.Updated = time.Now()
	c.Unlock()
	c.BroadcastUpdate()
}

func (c *ClusterManager) GetPartitions() []int32 {
	c.RLock()
	defer c.RUnlock()
	return c.node.Partitions
}
<commit_msg>only update primary/state if they have changed. fixes #430<commit_after>package cluster

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

	"github.com/hashicorp/memberlist"
	"github.com/raintank/worldping-api/pkg/log"
)

type ClusterManager struct {
	sync.RWMutex
	Peers map[string]Node
	node  Node
	list  *memberlist.Memberlist
}

func (c *ClusterManager) SetList(list *memberlist.Memberlist) {
	c.Lock()
	c.list = list
	c.Unlock()
}

func (c *ClusterManager) ThisNode() Node {
	c.RLock()
	defer c.RUnlock()
	return c.node
}

func (c *ClusterManager) PeersList() []Node {
	c.RLock()
	list := make([]Node, len(c.Peers))
	i := 0
	for _, n := range c.Peers {
		list[i] = n
		i++
	}
	c.RUnlock()
	return list
}

func (c *ClusterManager) NotifyJoin(node *memberlist.Node) {
	c.Lock()
	defer c.Unlock()
	if len(node.Meta) == 0 {
		return
	}
	log.Info("Node %s with address %s has joined the cluster\n", node.Name, node.Addr.String())
	peer := Node{}
	err := json.Unmarshal(node.Meta, &peer)
	if err != nil {
		panic(err)
	}
	peer.RemoteAddr = node.Addr.String()
	if peer.Name == c.node.Name {
		peer.local = true
	}
	c.Peers[node.Name] = peer
}

func (c *ClusterManager) NotifyLeave(node *memberlist.Node) {
	c.Lock()
	defer c.Unlock()
	log.Info("Node %s has left the cluster\n", node.Name)
	delete(c.Peers, node.Name)
}

func (c *ClusterManager) NotifyUpdate(node *memberlist.Node) {
	c.Lock()
	defer c.Unlock()
	if len(node.Meta) == 0 {
		return
	}
	peer := Node{}
	err := json.Unmarshal(node.Meta, &peer)
	if err != nil {
		panic(err)
	}
	peer.RemoteAddr = node.Addr.String()
	if peer.Name == c.node.Name {
		peer.local = true
	}
	c.Peers[node.Name] = peer
	log.Info("Node %s at %s has been updated - %s\n", node.Name, node.Addr.String(), node.Meta)
}

func (c *ClusterManager) BroadcastUpdate() {
	if c.list != nil {
		go c.list.UpdateNode(time.Second)
	}
}

// NodeMeta is used to retrieve meta-data about the current node
// when broadcasting an alive message. It's length is limited to
// the given byte size. This metadata is available in the Node structure.
func (c *ClusterManager) NodeMeta(limit int) []byte {
	c.Lock()
	meta, err := json.Marshal(c.node)
	c.Unlock()
	if err != nil {
		log.Fatal(4, err.Error())
	}
	return meta
}

// NotifyMsg is called when a user-data message is received.
// Care should be taken that this method does not block, since doing
// so would block the entire UDP packet receive loop. Additionally, the byte
// slice may be modified after the call returns, so it should be copied if needed.
func (c *ClusterManager) NotifyMsg(buf []byte) {
	// we dont have any need for passing messages between nodes, other then
	// the NodeMeta sent with alive messages.
	return
}

// GetBroadcasts is called when user data messages can be broadcast.
// It can return a list of buffers to send. Each buffer should assume an
// overhead as provided with a limit on the total byte size allowed.
// The total byte size of the resulting data to send must not exceed
// the limit. Care should be taken that this method does not block,
// since doing so would block the entire UDP packet receive loop.
func (c *ClusterManager) GetBroadcasts(overhead, limit int) [][]byte {
	// we dont have any need for passing messages between nodes, other then
	// the NodeMeta sent with alive messages.
	return nil
}

// LocalState is used for a TCP Push/Pull. This is sent to
// the remote side in addition to the membership information. Any
// data can be sent here. See MergeRemoteState as well. The `join`
// boolean indicates this is for a join instead of a push/pull.
func (c *ClusterManager) LocalState(join bool) []byte {
	c.Lock()
	c.Peers[c.node.Name] = c.node
	meta, err := json.Marshal(c.Peers)
	c.Unlock()
	if err != nil {
		log.Fatal(4, err.Error())
	}
	return meta
}

func (c *ClusterManager) MergeRemoteState(buf []byte, join bool) {
	knownPeers := make(map[string]Node)
	err := json.Unmarshal(buf, &knownPeers)
	if err != nil {
		log.Fatal(4, err.Error())
	}
	c.Lock()
	for name, meta := range knownPeers {
		if existing, ok := c.Peers[name]; ok {
			if meta.Updated.After(existing.Updated) {
				log.Info("updated node meta found in state update for %s", meta.Name)
				c.Peers[name] = meta
			}
		} else {
			log.Info("new node found in state update. %s", meta.Name)
			c.Peers[name] = meta
		}
	}
	c.Unlock()
}

// Returns true if the node is a ready to accept requests
// from users.
func (c *ClusterManager) IsReady() bool {
	c.RLock()
	defer c.RUnlock()
	return c.node.IsReady()
}

func (c *ClusterManager) SetReady() {
	c.Lock()
	if c.node.State == NodeReady {
		c.Unlock()
		return
	}
	c.node.State = NodeReady
	c.node.Updated = time.Now()
	c.Unlock()
	c.BroadcastUpdate()
}

func (c *ClusterManager) SetReadyIn(t time.Duration) {
	go func() {
		// wait for warmupPeriod before marking ourselves
		// as ready.
		time.Sleep(t)
		c.SetReady()
	}()
}

// Returns true if the node is a set as a primary node that should write data to cassandra.
func (c *ClusterManager) IsPrimary() bool {
	c.RLock()
	defer c.RUnlock()
	return c.node.Primary
}

// SetPrimary sets the primary status.
// Note: since we set the primary metric here, this should only be called on ThisNode !
func (c *ClusterManager) SetPrimary(p bool) {
	c.Lock()
	if c.node.Primary == p {
		c.Unlock()
		return
	}
	c.node.Primary = p
	c.node.PrimaryChange = time.Now()
	c.node.Updated = time.Now()
	c.Unlock()
	clusterPrimary.Set(p)
	c.BroadcastUpdate()
}

func (c *ClusterManager) SetPartitions(part []int32) {
	c.Lock()
	c.node.Partitions = part
	c.node.Updated = time.Now()
	c.Unlock()
	c.BroadcastUpdate()
}

func (c *ClusterManager) GetPartitions() []int32 {
	c.RLock()
	defer c.RUnlock()
	return c.node.Partitions
}
<|endoftext|>


<commit_before>package cluster

import (
	"strings"

	"github.com/docker/docker/pkg/stringid"
	"github.com/samalba/dockerclient"
)

// Network is exported
type Network struct {
	dockerclient.NetworkResource

	Engine *Engine
}

// Networks represents an array of networks
type Networks []*Network

// Uniq returns all uniq networks
func (networks Networks) Uniq() Networks {
	tmp := make(map[string]*Network)
	for _, network := range networks {
		if _, ok := tmp[network.ID]; ok {
			for id, endpoint := range network.Containers {
				tmp[network.ID].Containers[id] = endpoint
			}
		} else {
			tmp[network.ID] = network
		}
	}
	uniq := Networks{}
	for _, network := range tmp {
		uniq = append(uniq, network)
	}
	return uniq
}

// Filter returns networks filtered by names or ids
func (networks Networks) Filter(names []string, ids []string) Networks {
	if len(names) == 0 && len(ids) == 0 {
		return networks.Uniq()
	}

	out := Networks{}
	for _, idOrName := range append(names, ids...) {
		if network := networks.Get(idOrName); network != nil {
			out = append(out, network)
		}
	}
	return out
}

// Get returns a network using it's ID or Name
func (networks Networks) Get(IDOrName string) *Network {
	// Abort immediately if the name is empty.
	if len(IDOrName) == 0 {
		return nil
	}

	// Match exact or short Network ID.
	for _, network := range networks {
		if network.ID == IDOrName || stringid.TruncateID(network.ID) == IDOrName {
			return network
		}
	}

	candidates := []*Network{}

	// Match name or engine/name.
	for _, network := range networks {
		if network.Name == IDOrName || network.Engine.ID+"/"+network.Name == IDOrName || network.Engine.Name+"/"+network.Name == IDOrName {
			candidates = append(candidates, network)
		}
	}

	// Return if we found a unique match.
	if size := len(candidates); size == 1 {
		return candidates[0]
	} else if size > 1 {
		return nil
	}

	// Match /name and return as soon as we find one.
	for _, network := range networks {
		if network.Name == "/"+IDOrName {
			return network
		}
	}

	// Match Network ID prefix.
	for _, network := range networks {
		if strings.HasPrefix(network.ID, IDOrName) {
			candidates = append(candidates, network)
		}
	}

	if len(candidates) == 1 {
		return candidates[0]
	}

	return nil
}
<commit_msg>fix network filter with global network<commit_after>package cluster

import (
	"strings"

	"github.com/docker/docker/pkg/stringid"
	"github.com/samalba/dockerclient"
)

// Network is exported
type Network struct {
	dockerclient.NetworkResource

	Engine *Engine
}

// Networks represents an array of networks
type Networks []*Network

// Uniq returns all uniq networks
func (networks Networks) Uniq() Networks {
	tmp := make(map[string]*Network)
	for _, network := range networks {
		if _, ok := tmp[network.ID]; ok {
			for id, endpoint := range network.Containers {
				tmp[network.ID].Containers[id] = endpoint
			}
		} else {
			tmp[network.ID] = network
		}
	}
	uniq := Networks{}
	for _, network := range tmp {
		uniq = append(uniq, network)
	}
	return uniq
}

// Filter returns networks filtered by names or ids
func (networks Networks) Filter(names []string, ids []string) Networks {
	if len(names) == 0 && len(ids) == 0 {
		return networks.Uniq()
	}

	out := Networks{}
	for _, idOrName := range append(names, ids...) {
		if network := networks.Get(idOrName); network != nil {
			out = append(out, network)
		}
	}
	return out
}

// Get returns a network using it's ID or Name
func (networks Networks) Get(IDOrName string) *Network {
	// Abort immediately if the name is empty.
	if len(IDOrName) == 0 {
		return nil
	}

	// Match exact or short Network ID.
	for _, network := range networks {
		if network.ID == IDOrName || stringid.TruncateID(network.ID) == IDOrName {
			return network
		}
	}

	candidates := Networks{}

	// Match name or engine/name.
	for _, network := range networks {
		if network.Name == IDOrName || network.Engine.ID+"/"+network.Name == IDOrName || network.Engine.Name+"/"+network.Name == IDOrName {
			candidates = append(candidates, network)
		}
	}

	// Return if we found a unique match.
	if size := len(candidates.Uniq()); size == 1 {
		return candidates[0]
	} else if size > 1 {
		return nil
	}

	// Match /name and return as soon as we find one.
	for _, network := range networks {
		if network.Name == "/"+IDOrName {
			return network
		}
	}

	// Match Network ID prefix.
	for _, network := range networks {
		if strings.HasPrefix(network.ID, IDOrName) {
			candidates = append(candidates, network)
		}
	}

	if len(candidates.Uniq()) == 1 {
		return candidates[0]
	}

	return nil
}
<|endoftext|>


<commit_before>// Package app provides application support for context and MongoDB access.
// Current Status Codes:
//		200 OK           : StatusOK                  : Call is success and returning data.
//		204 No Content   : StatusNoContent           : Call is success and returns no data.
//		400 Bad Request  : StatusBadRequest          : Invalid post data (syntax or semantics).
//		401 Unauthorized : StatusUnauthorized        : Authentication failure.
//		404 Not Found    : StatusNotFound            : Invalid URL or identifier.
//		500 Internal     : StatusInternalServerError : Application specific beyond scope of user.
package app

import (
	"encoding/json"
	"io"
	"net/http"
	"time"

	"github.com/ardanlabs/kit/log"
)

// Invalid describes a validation error belonging to a specific field.
type Invalid struct {
	Fld string `json:"field_name"`
	Err string `json:"error"`
}

// jsonError is the response for errors that occur within the API.
type jsonError struct {
	Error  string    `json:"error"`
	Fields []Invalid `json:"fields,omitempty"`
}

//==============================================================================

// Context contains data associated with a single request.
type Context struct {
	http.ResponseWriter
	Request   *http.Request
	Now       time.Time
	Params    map[string]string
	SessionID string
	Status    int
	Ctx       map[string]interface{}
}

// Error handles all error responses for the API.
func (c *Context) Error(err error) {
	switch err {
	case ErrNotFound:
		c.RespondError(err.Error(), http.StatusNotFound)
	case ErrInvalidID:
		c.RespondError(err.Error(), http.StatusBadRequest)
	case ErrValidation:
		c.RespondError(err.Error(), http.StatusBadRequest)
	case ErrNotAuthorized:
		c.RespondError(err.Error(), http.StatusUnauthorized)
	default:
		c.RespondError(err.Error(), http.StatusInternalServerError)
	}
}

// Respond sends JSON to the client.
// If code is StatusNoContent, v is expected to be nil.
func (c *Context) Respond(data interface{}, code int) {
	log.User(c.SessionID, "api : Respond", "Started : Code[%d]", code)

	c.Status = code

	if code == http.StatusNoContent {
		c.WriteHeader(http.StatusNoContent)
		return
	}

	// Set application default header values.
	c.Header().Set("Content-Type", "application/json")

	// Load any user defined header values.
	if app.userHeaders != nil {
		for key, value := range app.userHeaders {
			log.User("startup", "Init", "Setting user headers : %s:%s", key, value)
			c.Header().Set(key, value)
		}
	}

	c.WriteHeader(code)

	// Marshal the data into a JSON string.
	jsonData, err := json.Marshal(data)
	if err != nil {
		jsonData = []byte("{}")
	}

	// Look for a JSONP marker
	if cb := c.Request.URL.Query().Get("callback"); cb != "" {

		// We need to wrap the result in a function call.
		// callback_value({"data_1": "hello world", "data_2": ["the","sun","is","shining"]});
		io.WriteString(c, cb+"("+string(jsonData)+")")

		log.User(c.SessionID, "api : Respond", "Completed")
		return
	}

	// We can send the result straight through.
	io.WriteString(c, string(jsonData))

	log.User(c.SessionID, "api : Respond", "Completed")
}

// RespondInvalid sends JSON describing field validation errors.
func (c *Context) RespondInvalid(fields []Invalid) {
	v := jsonError{
		Error:  "field validation failure",
		Fields: fields,
	}
	c.Respond(v, http.StatusBadRequest)
}

// RespondError sends JSON describing the error
func (c *Context) RespondError(error string, code int) {
	c.Respond(jsonError{Error: error}, code)
}
<commit_msg>Fixed bug with not sending headers on StatusNoContent.<commit_after>// Package app provides application support for context and MongoDB access.
// Current Status Codes:
//		200 OK           : StatusOK                  : Call is success and returning data.
//		204 No Content   : StatusNoContent           : Call is success and returns no data.
//		400 Bad Request  : StatusBadRequest          : Invalid post data (syntax or semantics).
//		401 Unauthorized : StatusUnauthorized        : Authentication failure.
//		404 Not Found    : StatusNotFound            : Invalid URL or identifier.
//		500 Internal     : StatusInternalServerError : Application specific beyond scope of user.
package app

import (
	"encoding/json"
	"io"
	"net/http"
	"time"

	"github.com/ardanlabs/kit/log"
)

// Invalid describes a validation error belonging to a specific field.
type Invalid struct {
	Fld string `json:"field_name"`
	Err string `json:"error"`
}

// jsonError is the response for errors that occur within the API.
type jsonError struct {
	Error  string    `json:"error"`
	Fields []Invalid `json:"fields,omitempty"`
}

//==============================================================================

// Context contains data associated with a single request.
type Context struct {
	http.ResponseWriter
	Request   *http.Request
	Now       time.Time
	Params    map[string]string
	SessionID string
	Status    int
	Ctx       map[string]interface{}
}

// Error handles all error responses for the API.
func (c *Context) Error(err error) {
	switch err {
	case ErrNotFound:
		c.RespondError(err.Error(), http.StatusNotFound)
	case ErrInvalidID:
		c.RespondError(err.Error(), http.StatusBadRequest)
	case ErrValidation:
		c.RespondError(err.Error(), http.StatusBadRequest)
	case ErrNotAuthorized:
		c.RespondError(err.Error(), http.StatusUnauthorized)
	default:
		c.RespondError(err.Error(), http.StatusInternalServerError)
	}
}

// Respond sends JSON to the client.
// If code is StatusNoContent, v is expected to be nil.
func (c *Context) Respond(data interface{}, code int) {
	log.User(c.SessionID, "api : Respond", "Started : Code[%d]", code)

	c.Status = code

	// Load any user defined header values.
	if app.userHeaders != nil {
		for key, value := range app.userHeaders {
			log.User("startup", "Init", "Setting user headers : %s:%s", key, value)
			c.Header().Set(key, value)
		}
	}

	// Just set the status code and we are done.
	if code == http.StatusNoContent {
		c.WriteHeader(code)
		return
	}

	// Set the content type.
	c.Header().Set("Content-Type", "application/json")

	// Marshal the data into a JSON string.
	jsonData, err := json.Marshal(data)
	if err != nil {
		jsonData = []byte("{}")
	}

	// Look for a JSONP marker
	if cb := c.Request.URL.Query().Get("callback"); cb != "" {

		// We need to wrap the result in a function call.
		// callback_value({"data_1": "hello world", "data_2": ["the","sun","is","shining"]});
		io.WriteString(c, cb+"("+string(jsonData)+")")

		log.User(c.SessionID, "api : Respond", "Completed")
		return
	}

	// We can send the result straight through.
	io.WriteString(c, string(jsonData))

	log.User(c.SessionID, "api : Respond", "Completed")
}

// RespondInvalid sends JSON describing field validation errors.
func (c *Context) RespondInvalid(fields []Invalid) {
	v := jsonError{
		Error:  "field validation failure",
		Fields: fields,
	}
	c.Respond(v, http.StatusBadRequest)
}

// RespondError sends JSON describing the error
func (c *Context) RespondError(error string, code int) {
	c.Respond(jsonError{Error: error}, code)
}
<|endoftext|>


<commit_before>package driver

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/env"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutils "github.com/hashicorp/nomad/client/testutil"
)

func TestRktVersionRegex(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	input_rkt := "rkt version 0.8.1"
	input_appc := "appc version 1.2.0"
	expected_rkt := "0.8.1"
	expected_appc := "1.2.0"
	rktMatches := reRktVersion.FindStringSubmatch(input_rkt)
	appcMatches := reAppcVersion.FindStringSubmatch(input_appc)
	if rktMatches[1] != expected_rkt {
		fmt.Printf("Test failed; got %q; want %q\n", rktMatches[1], expected_rkt)
	}
	if appcMatches[1] != expected_appc {
		fmt.Printf("Test failed; got %q; want %q\n", appcMatches[1], expected_appc)
	}
}

// The fingerprinter test should always pass, even if rkt is not installed.
func TestRktDriver_Fingerprint(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	driverCtx, _ := testDriverContexts(&structs.Task{Name: "foo"})
	d := NewRktDriver(driverCtx)
	node := &structs.Node{
		Attributes: make(map[string]string),
	}
	apply, err := d.Fingerprint(&config.Config{}, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !apply {
		t.Fatalf("should apply")
	}
	if node.Attributes["driver.rkt"] != "1" {
		t.Fatalf("Missing Rkt driver")
	}
	if node.Attributes["driver.rkt.version"] == "" {
		t.Fatalf("Missing Rkt driver version")
	}
	if node.Attributes["driver.rkt.appc.version"] == "" {
		t.Fatalf("Missing appc version for the Rkt driver")
	}
}

func TestRktDriver_Start_DNS(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	// TODO: use test server to load from a fixture
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"trust_prefix":       "coreos.com/etcd",
			"image":              "coreos.com/etcd:v2.0.4",
			"command":            "/etcd",
			"dns_servers":        []string{"8.8.8.8", "8.8.4.4"},
			"dns_search_domains": []string{"example.com", "example.org", "example.net"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()

	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	// Attempt to open
	handle2, err := d.Open(execCtx, handle.ID())
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle2 == nil {
		t.Fatalf("missing handle")
	}
}

func TestRktDriver_Start_Wait(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"trust_prefix": "coreos.com/etcd",
			"image":        "coreos.com/etcd:v2.0.4",
			"command":      "/etcd",
			"args":         []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	// Update should be a no-op
	err = handle.Update(task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestRktDriver_Start_Wait_Skip_Trust(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"image":   "coreos.com/etcd:v2.0.4",
			"command": "/etcd",
			"args":    []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	// Update should be a no-op
	err = handle.Update(task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestRktDriver_Start_Wait_AllocDir(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)

	exp := []byte{'w', 'i', 'n'}
	file := "output.txt"

	task := &structs.Task{
		Name: "alpine",
		Config: map[string]interface{}{
			"image":   "docker://alpine",
			"command": "/bin/sh",
			"args": []string{
				"-c",
				fmt.Sprintf(`echo -n %s > ${%s}/%s`, string(exp), env.AllocDir, file),
			},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}

	// Check that data was written to the shared alloc directory.
	outputFile := filepath.Join(execCtx.AllocDir.SharedDir, file)
	act, err := ioutil.ReadFile(outputFile)
	if err != nil {
		t.Fatalf("Couldn't read expected output: %v", err)
	}

	if !reflect.DeepEqual(act, exp) {
		t.Fatalf("Command output is %v; expected %v", act, exp)
	}
}

func TestRktDriverUser(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		User: "alice",
		Config: map[string]interface{}{
			"trust_prefix": "coreos.com/etcd",
			"image":        "coreos.com/etcd:v2.0.4",
			"command":      "/etcd",
			"args":         []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err == nil {
		handle.Kill()
		t.Fatalf("Should've failed")
	}
	msg := "unknown user alice"
	if !strings.Contains(err.Error(), msg) {
		t.Fatalf("Expecting '%v' in '%v'", msg, err)
	}
}
<commit_msg>Add TestRktTrustPrefix<commit_after>package driver

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/env"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutils "github.com/hashicorp/nomad/client/testutil"
)

func TestRktVersionRegex(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	input_rkt := "rkt version 0.8.1"
	input_appc := "appc version 1.2.0"
	expected_rkt := "0.8.1"
	expected_appc := "1.2.0"
	rktMatches := reRktVersion.FindStringSubmatch(input_rkt)
	appcMatches := reAppcVersion.FindStringSubmatch(input_appc)
	if rktMatches[1] != expected_rkt {
		fmt.Printf("Test failed; got %q; want %q\n", rktMatches[1], expected_rkt)
	}
	if appcMatches[1] != expected_appc {
		fmt.Printf("Test failed; got %q; want %q\n", appcMatches[1], expected_appc)
	}
}

// The fingerprinter test should always pass, even if rkt is not installed.
func TestRktDriver_Fingerprint(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	driverCtx, _ := testDriverContexts(&structs.Task{Name: "foo"})
	d := NewRktDriver(driverCtx)
	node := &structs.Node{
		Attributes: make(map[string]string),
	}
	apply, err := d.Fingerprint(&config.Config{}, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !apply {
		t.Fatalf("should apply")
	}
	if node.Attributes["driver.rkt"] != "1" {
		t.Fatalf("Missing Rkt driver")
	}
	if node.Attributes["driver.rkt.version"] == "" {
		t.Fatalf("Missing Rkt driver version")
	}
	if node.Attributes["driver.rkt.appc.version"] == "" {
		t.Fatalf("Missing appc version for the Rkt driver")
	}
}

func TestRktDriver_Start_DNS(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	// TODO: use test server to load from a fixture
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"trust_prefix":       "coreos.com/etcd",
			"image":              "coreos.com/etcd:v2.0.4",
			"command":            "/etcd",
			"dns_servers":        []string{"8.8.8.8", "8.8.4.4"},
			"dns_search_domains": []string{"example.com", "example.org", "example.net"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()

	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	// Attempt to open
	handle2, err := d.Open(execCtx, handle.ID())
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle2 == nil {
		t.Fatalf("missing handle")
	}
}

func TestRktDriver_Start_Wait(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"trust_prefix": "coreos.com/etcd",
			"image":        "coreos.com/etcd:v2.0.4",
			"command":      "/etcd",
			"args":         []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	// Update should be a no-op
	err = handle.Update(task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestRktDriver_Start_Wait_Skip_Trust(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"image":   "coreos.com/etcd:v2.0.4",
			"command": "/etcd",
			"args":    []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	// Update should be a no-op
	err = handle.Update(task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestRktDriver_Start_Wait_AllocDir(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)

	exp := []byte{'w', 'i', 'n'}
	file := "output.txt"

	task := &structs.Task{
		Name: "alpine",
		Config: map[string]interface{}{
			"image":   "docker://alpine",
			"command": "/bin/sh",
			"args": []string{
				"-c",
				fmt.Sprintf(`echo -n %s > ${%s}/%s`, string(exp), env.AllocDir, file),
			},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}
	defer handle.Kill()

	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}

	// Check that data was written to the shared alloc directory.
	outputFile := filepath.Join(execCtx.AllocDir.SharedDir, file)
	act, err := ioutil.ReadFile(outputFile)
	if err != nil {
		t.Fatalf("Couldn't read expected output: %v", err)
	}

	if !reflect.DeepEqual(act, exp) {
		t.Fatalf("Command output is %v; expected %v", act, exp)
	}
}

func TestRktDriverUser(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}

	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		User: "alice",
		Config: map[string]interface{}{
			"trust_prefix": "coreos.com/etcd",
			"image":        "coreos.com/etcd:v2.0.4",
			"command":      "/etcd",
			"args":         []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}

	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()
	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err == nil {
		handle.Kill()
		t.Fatalf("Should've failed")
	}
	msg := "unknown user alice"
	if !strings.Contains(err.Error(), msg) {
		t.Fatalf("Expecting '%v' in '%v'", msg, err)
	}
}

func TestRktTrustPrefix(t *testing.T) {
	if os.Getenv("NOMAD_TEST_RKT") == "" {
		t.Skip("skipping rkt tests")
	}
	ctestutils.RktCompatible(t)
	task := &structs.Task{
		Name: "etcd",
		Config: map[string]interface{}{
			"trust_prefix": "example.com/invalid",
			"image":        "coreos.com/etcd:v2.0.4",
			"command":      "/etcd",
			"args":         []string{"--version"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: &structs.Resources{
			MemoryMB: 128,
			CPU:      100,
		},
	}
	driverCtx, execCtx := testDriverContexts(task)
	defer execCtx.AllocDir.Destroy()

	d := NewRktDriver(driverCtx)

	handle, err := d.Start(execCtx, task)
	if err == nil {
		handle.Kill()
		t.Fatalf("Should've failed")
	}
	msg := "Error running rkt trust"
	if !strings.Contains(err.Error(), msg) {
		t.Fatalf("Expecting '%v' in '%v'", msg, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"

	"github.com/mattn/go-colorable"
	"github.com/mattn/go-isatty"
	"gopkg.in/alecthomas/kingpin.v3-unstable"

	"github.com/alecthomas/chroma"
	"github.com/alecthomas/chroma/formatters"
	"github.com/alecthomas/chroma/formatters/html"
	"github.com/alecthomas/chroma/lexers"
	"github.com/alecthomas/chroma/styles"
)

var (
	// Populated by goreleaser.
	version = "?"
	commit  = "?"
	date    = "?"

	profileFlag    = kingpin.Flag("profile", "Enable profiling to file.").Hidden().String()
	listFlag       = kingpin.Flag("list", "List lexers, styles and formatters.").Bool()
	unbufferedFlag = kingpin.Flag("unbuffered", "Do not buffer output.").Bool()
	traceFlag      = kingpin.Flag("trace", "Trace lexer states as they are traversed.").Bool()
	checkFlag      = kingpin.Flag("check", "Do not format, check for tokenization errors instead.").Bool()

	lexerFlag     = kingpin.Flag("lexer", "Lexer to use when formatting.").PlaceHolder("autodetect").Short('l').Enum(lexers.Names(true)...)
	styleFlag     = kingpin.Flag("style", "Style to use for formatting.").Short('s').Default("swapoff").Enum(styles.Names()...)
	formatterFlag = kingpin.Flag("formatter", "Formatter to use.").Default("terminal").Short('f').Enum(formatters.Names()...)

	htmlFlag               = kingpin.Flag("html", "Enable HTML mode (equivalent to '--formatter html').").Bool()
	htmlPrefixFlag         = kingpin.Flag("html-prefix", "HTML CSS class prefix.").PlaceHolder("PREFIX").String()
	htmlStylesFlag         = kingpin.Flag("html-styles", "Output HTML CSS styles.").Bool()
	htmlOnlyFlag           = kingpin.Flag("html-only", "Output HTML fragment.").Bool()
	htmlInlineStyleFlag    = kingpin.Flag("html-inline-styles", "Output HTML with inline styles (no classes).").Bool()
	htmlTabWidthFlag       = kingpin.Flag("html-tab-width", "Set the HTML tab width.").Default("8").Int()
	htmlLinesFlag          = kingpin.Flag("html-lines", "Include line numbers in output.").Bool()
	htmlLinesStyleFlag     = kingpin.Flag("html-lines-style", "Style for line numbers.").String()
	htmlHighlightFlag      = kingpin.Flag("html-highlight", "Highlight these lines.").PlaceHolder("N[:M][,...]").String()
	htmlHighlightStyleFlag = kingpin.Flag("html-highlight-style", "Style used for highlighting lines.").String()
	htmlBaseLineFlag       = kingpin.Flag("html-base-line", "Base line number.").Default("1").Int()

	filesArgs = kingpin.Arg("files", "Files to highlight.").ExistingFiles()
)

type flushableWriter interface {
	io.Writer
	Flush() error
}

type nopFlushableWriter struct{ io.Writer }

func (n *nopFlushableWriter) Flush() error { return nil }

func main() {
	kingpin.CommandLine.Version(fmt.Sprintf("%s-%s-%s", version, commit, date))
	kingpin.CommandLine.Help = `
Chroma is a general purpose syntax highlighting library and corresponding
command, for Go.
`
	kingpin.Parse()
	if *listFlag {
		listAll()
		return
	}
	if *profileFlag != "" {
		f, err := os.Create(*profileFlag)
		kingpin.FatalIfError(err, "")
		pprof.StartCPUProfile(f)
		signals := make(chan os.Signal, 1)
		signal.Notify(signals, os.Interrupt)
		go func() {
			<-signals
			pprof.StopCPUProfile()
			os.Exit(128 + 3)
		}()
		defer pprof.StopCPUProfile()
	}

	var out io.Writer = os.Stdout
	if runtime.GOOS == "windows" && isatty.IsTerminal(os.Stdout.Fd()) {
		out = colorable.NewColorableStdout()
	}
	var w flushableWriter
	if *unbufferedFlag {
		w = &nopFlushableWriter{out}
	} else {
		w = bufio.NewWriterSize(out, 16384)
	}
	defer w.Flush()

	if *htmlFlag {
		*formatterFlag = "html"
	}

	// Retrieve user-specified style, clone it, and add some overrides.
	builder := styles.Get(*styleFlag).Builder()
	if *htmlHighlightStyleFlag != "" {
		builder.Add(chroma.LineHighlight, *htmlHighlightStyleFlag)
	}
	if *htmlLinesStyleFlag != "" {
		builder.Add(chroma.LineNumbers, *htmlLinesStyleFlag)
	}
	style, err := builder.Build()
	kingpin.FatalIfError(err, "")

	if *formatterFlag == "html" {
		options := []html.Option{
			html.TabWidth(*htmlTabWidthFlag),
			html.BaseLineNumber(*htmlBaseLineFlag),
		}
		if *htmlPrefixFlag != "" {
			options = append(options, html.ClassPrefix(*htmlPrefixFlag))
		}

		// Dump styles.
		if *htmlStylesFlag {
			formatter := html.New(html.WithClasses())
			formatter.WriteCSS(w, style)
			return
		}
		if !*htmlInlineStyleFlag {
			options = append(options, html.WithClasses())
		}
		if !*htmlOnlyFlag {
			options = append(options, html.Standalone())
		}
		if *htmlLinesFlag {
			options = append(options, html.WithLineNumbers())
		}
		if len(*htmlHighlightFlag) > 0 {
			ranges := [][2]int{}
			for _, span := range strings.Split(*htmlHighlightFlag, ",") {
				parts := strings.Split(span, ":")
				if len(parts) > 2 {
					kingpin.Fatalf("range should be N[:M], not %q", span)
				}
				start, err := strconv.ParseInt(parts[0], 10, 64)
				kingpin.FatalIfError(err, "min value of range should be integer not %q", parts[0])
				end := start
				if len(parts) == 2 {
					end, err = strconv.ParseInt(parts[1], 10, 64)
					kingpin.FatalIfError(err, "max value of range should be integer not %q", parts[1])
				}
				ranges = append(ranges, [2]int{int(start), int(end)})
			}
			options = append(options, html.HighlightLines(ranges))
		}
		formatters.Register("html", html.New(options...))
	}
	if len(*filesArgs) == 0 {
		contents, err := ioutil.ReadAll(os.Stdin)
		kingpin.FatalIfError(err, "")
		format(w, style, lex("", string(contents)))
	} else {
		for _, filename := range *filesArgs {
			contents, err := ioutil.ReadFile(filename)
			kingpin.FatalIfError(err, "")
			if *checkFlag {
				check(filename, lex(filename, string(contents)))
			} else {
				format(w, style, lex(filename, string(contents)))
			}
		}
	}
}

func listAll() {
	fmt.Println("lexers:")
	sort.Sort(lexers.Registry.Lexers)
	for _, l := range lexers.Registry.Lexers {
		config := l.Config()
		fmt.Printf("  %s\n", config.Name)
		filenames := []string{}
		filenames = append(filenames, config.Filenames...)
		filenames = append(filenames, config.AliasFilenames...)
		if len(config.Aliases) > 0 {
			fmt.Printf("    aliases: %s\n", strings.Join(config.Aliases, " "))
		}
		if len(filenames) > 0 {
			fmt.Printf("    filenames: %s\n", strings.Join(filenames, " "))
		}
		if len(config.MimeTypes) > 0 {
			fmt.Printf("    mimetypes: %s\n", strings.Join(config.MimeTypes, " "))
		}
	}
	fmt.Println()
	fmt.Printf("styles:")
	for name := range styles.Registry {
		fmt.Printf(" %s", name)
	}
	fmt.Println()
	fmt.Printf("formatters:")
	for name := range formatters.Registry {
		fmt.Printf(" %s", name)
	}
	fmt.Println()
}

func lex(path string, contents string) chroma.Iterator {
	lexer := selexer(path, contents)
	if lexer == nil {
		lexer = lexers.Fallback
	}
	if rel, ok := lexer.(*chroma.RegexLexer); ok {
		rel.Trace(*traceFlag)
	}
	lexer = chroma.Coalesce(lexer)
	it, err := lexer.Tokenise(nil, string(contents))
	kingpin.FatalIfError(err, "")
	return it
}

func selexer(path, contents string) (lexer chroma.Lexer) {
	if *lexerFlag != "" {
		return lexers.Get(*lexerFlag)
	}
	if path != "" {
		lexer := lexers.Match(path)
		if lexer != nil {
			return lexer
		}
	}
	return lexers.Analyse(contents)
}

func format(w io.Writer, style *chroma.Style, it chroma.Iterator) {
	formatter := formatters.Get(*formatterFlag)
	err := formatter.Format(w, style, it)
	kingpin.FatalIfError(err, "")
}

func check(filename string, it chroma.Iterator) {
	line, col := 1, 0
	for token := it(); token != nil; token = it() {
		if token.Type == chroma.Error {
			fmt.Printf("%s:%d:%d %q\n", filename, line, col, token.String())
		}
		for _, c := range token.String() {
			col++
			if c == '\n' {
				line, col = line+1, 0
			}
		}
	}
}
<commit_msg>Fix order output of chroma cli tool<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"

	"github.com/mattn/go-colorable"
	"github.com/mattn/go-isatty"
	"gopkg.in/alecthomas/kingpin.v3-unstable"

	"github.com/alecthomas/chroma"
	"github.com/alecthomas/chroma/formatters"
	"github.com/alecthomas/chroma/formatters/html"
	"github.com/alecthomas/chroma/lexers"
	"github.com/alecthomas/chroma/styles"
)

var (
	// Populated by goreleaser.
	version = "?"
	commit  = "?"
	date    = "?"

	profileFlag    = kingpin.Flag("profile", "Enable profiling to file.").Hidden().String()
	listFlag       = kingpin.Flag("list", "List lexers, styles and formatters.").Bool()
	unbufferedFlag = kingpin.Flag("unbuffered", "Do not buffer output.").Bool()
	traceFlag      = kingpin.Flag("trace", "Trace lexer states as they are traversed.").Bool()
	checkFlag      = kingpin.Flag("check", "Do not format, check for tokenization errors instead.").Bool()

	lexerFlag     = kingpin.Flag("lexer", "Lexer to use when formatting.").PlaceHolder("autodetect").Short('l').Enum(lexers.Names(true)...)
	styleFlag     = kingpin.Flag("style", "Style to use for formatting.").Short('s').Default("swapoff").Enum(styles.Names()...)
	formatterFlag = kingpin.Flag("formatter", "Formatter to use.").Default("terminal").Short('f').Enum(formatters.Names()...)

	htmlFlag               = kingpin.Flag("html", "Enable HTML mode (equivalent to '--formatter html').").Bool()
	htmlPrefixFlag         = kingpin.Flag("html-prefix", "HTML CSS class prefix.").PlaceHolder("PREFIX").String()
	htmlStylesFlag         = kingpin.Flag("html-styles", "Output HTML CSS styles.").Bool()
	htmlOnlyFlag           = kingpin.Flag("html-only", "Output HTML fragment.").Bool()
	htmlInlineStyleFlag    = kingpin.Flag("html-inline-styles", "Output HTML with inline styles (no classes).").Bool()
	htmlTabWidthFlag       = kingpin.Flag("html-tab-width", "Set the HTML tab width.").Default("8").Int()
	htmlLinesFlag          = kingpin.Flag("html-lines", "Include line numbers in output.").Bool()
	htmlLinesStyleFlag     = kingpin.Flag("html-lines-style", "Style for line numbers.").String()
	htmlHighlightFlag      = kingpin.Flag("html-highlight", "Highlight these lines.").PlaceHolder("N[:M][,...]").String()
	htmlHighlightStyleFlag = kingpin.Flag("html-highlight-style", "Style used for highlighting lines.").String()
	htmlBaseLineFlag       = kingpin.Flag("html-base-line", "Base line number.").Default("1").Int()

	filesArgs = kingpin.Arg("files", "Files to highlight.").ExistingFiles()
)

type flushableWriter interface {
	io.Writer
	Flush() error
}

type nopFlushableWriter struct{ io.Writer }

func (n *nopFlushableWriter) Flush() error { return nil }

func main() {
	kingpin.CommandLine.Version(fmt.Sprintf("%s-%s-%s", version, commit, date))
	kingpin.CommandLine.Help = `
Chroma is a general purpose syntax highlighting library and corresponding
command, for Go.
`
	kingpin.Parse()
	if *listFlag {
		listAll()
		return
	}
	if *profileFlag != "" {
		f, err := os.Create(*profileFlag)
		kingpin.FatalIfError(err, "")
		pprof.StartCPUProfile(f)
		signals := make(chan os.Signal, 1)
		signal.Notify(signals, os.Interrupt)
		go func() {
			<-signals
			pprof.StopCPUProfile()
			os.Exit(128 + 3)
		}()
		defer pprof.StopCPUProfile()
	}

	var out io.Writer = os.Stdout
	if runtime.GOOS == "windows" && isatty.IsTerminal(os.Stdout.Fd()) {
		out = colorable.NewColorableStdout()
	}
	var w flushableWriter
	if *unbufferedFlag {
		w = &nopFlushableWriter{out}
	} else {
		w = bufio.NewWriterSize(out, 16384)
	}
	defer w.Flush()

	if *htmlFlag {
		*formatterFlag = "html"
	}

	// Retrieve user-specified style, clone it, and add some overrides.
	builder := styles.Get(*styleFlag).Builder()
	if *htmlHighlightStyleFlag != "" {
		builder.Add(chroma.LineHighlight, *htmlHighlightStyleFlag)
	}
	if *htmlLinesStyleFlag != "" {
		builder.Add(chroma.LineNumbers, *htmlLinesStyleFlag)
	}
	style, err := builder.Build()
	kingpin.FatalIfError(err, "")

	if *formatterFlag == "html" {
		options := []html.Option{
			html.TabWidth(*htmlTabWidthFlag),
			html.BaseLineNumber(*htmlBaseLineFlag),
		}
		if *htmlPrefixFlag != "" {
			options = append(options, html.ClassPrefix(*htmlPrefixFlag))
		}

		// Dump styles.
		if *htmlStylesFlag {
			formatter := html.New(html.WithClasses())
			formatter.WriteCSS(w, style)
			return
		}
		if !*htmlInlineStyleFlag {
			options = append(options, html.WithClasses())
		}
		if !*htmlOnlyFlag {
			options = append(options, html.Standalone())
		}
		if *htmlLinesFlag {
			options = append(options, html.WithLineNumbers())
		}
		if len(*htmlHighlightFlag) > 0 {
			ranges := [][2]int{}
			for _, span := range strings.Split(*htmlHighlightFlag, ",") {
				parts := strings.Split(span, ":")
				if len(parts) > 2 {
					kingpin.Fatalf("range should be N[:M], not %q", span)
				}
				start, err := strconv.ParseInt(parts[0], 10, 64)
				kingpin.FatalIfError(err, "min value of range should be integer not %q", parts[0])
				end := start
				if len(parts) == 2 {
					end, err = strconv.ParseInt(parts[1], 10, 64)
					kingpin.FatalIfError(err, "max value of range should be integer not %q", parts[1])
				}
				ranges = append(ranges, [2]int{int(start), int(end)})
			}
			options = append(options, html.HighlightLines(ranges))
		}
		formatters.Register("html", html.New(options...))
	}
	if len(*filesArgs) == 0 {
		contents, err := ioutil.ReadAll(os.Stdin)
		kingpin.FatalIfError(err, "")
		format(w, style, lex("", string(contents)))
	} else {
		for _, filename := range *filesArgs {
			contents, err := ioutil.ReadFile(filename)
			kingpin.FatalIfError(err, "")
			if *checkFlag {
				check(filename, lex(filename, string(contents)))
			} else {
				format(w, style, lex(filename, string(contents)))
			}
		}
	}
}

func listAll() {
	fmt.Println("lexers:")
	sort.Sort(lexers.Registry.Lexers)
	for _, l := range lexers.Registry.Lexers {
		config := l.Config()
		fmt.Printf("  %s\n", config.Name)
		filenames := []string{}
		filenames = append(filenames, config.Filenames...)
		filenames = append(filenames, config.AliasFilenames...)
		if len(config.Aliases) > 0 {
			fmt.Printf("    aliases: %s\n", strings.Join(config.Aliases, " "))
		}
		if len(filenames) > 0 {
			fmt.Printf("    filenames: %s\n", strings.Join(filenames, " "))
		}
		if len(config.MimeTypes) > 0 {
			fmt.Printf("    mimetypes: %s\n", strings.Join(config.MimeTypes, " "))
		}
	}
	fmt.Println()
	fmt.Printf("styles:")
	for _, name := range styles.Names() {
		fmt.Printf(" %s", name)
	}
	fmt.Println()
	fmt.Printf("formatters:")
	for _, name := range formatters.Names() {
		fmt.Printf(" %s", name)
	}
	fmt.Println()
}

func lex(path string, contents string) chroma.Iterator {
	lexer := selexer(path, contents)
	if lexer == nil {
		lexer = lexers.Fallback
	}
	if rel, ok := lexer.(*chroma.RegexLexer); ok {
		rel.Trace(*traceFlag)
	}
	lexer = chroma.Coalesce(lexer)
	it, err := lexer.Tokenise(nil, string(contents))
	kingpin.FatalIfError(err, "")
	return it
}

func selexer(path, contents string) (lexer chroma.Lexer) {
	if *lexerFlag != "" {
		return lexers.Get(*lexerFlag)
	}
	if path != "" {
		lexer := lexers.Match(path)
		if lexer != nil {
			return lexer
		}
	}
	return lexers.Analyse(contents)
}

func format(w io.Writer, style *chroma.Style, it chroma.Iterator) {
	formatter := formatters.Get(*formatterFlag)
	err := formatter.Format(w, style, it)
	kingpin.FatalIfError(err, "")
}

func check(filename string, it chroma.Iterator) {
	line, col := 1, 0
	for token := it(); token != nil; token = it() {
		if token.Type == chroma.Error {
			fmt.Printf("%s:%d:%d %q\n", filename, line, col, token.String())
		}
		for _, c := range token.String() {
			col++
			if c == '\n' {
				line, col = line+1, 0
			}
		}
	}
}
<|endoftext|>