<commit_before>// Copyright 2015-present Oursky Ltd.
//
// 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 plugin

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"github.com/robfig/cron"
	"github.com/skygeario/skygear-server/pkg/server/logging"
	"github.com/skygeario/skygear-server/pkg/server/plugin/hook"
	"github.com/skygeario/skygear-server/pkg/server/plugin/provider"
	"github.com/skygeario/skygear-server/pkg/server/router"
	"github.com/skygeario/skygear-server/pkg/server/skyconfig"
)

var log = logging.LoggerEntry("plugin")

const (
	// PluginInitMaxRetryCount defines the maximum retries for plugin initialization
	PluginInitMaxRetryCount = 100
)

// Plugin represents a collection of handlers, hooks and lambda functions
// that extends or modifies functionality provided by skygear.
type Plugin struct {
	initRetryCount int
	transport      Transport
	gatewayMap     map[string]*router.Gateway
}

type pluginHandlerInfo struct {
	AuthRequired bool     `json:"auth_required"`
	Name         string   `json:"name"`
	Methods      []string `json:"methods"`
	KeyRequired  bool     `json:"key_required"`
	UserRequired bool     `json:"user_required"`
}

type pluginHookInfo struct {
	Async   bool   `json:"async"`   // execute hook asynchronously
	Trigger string `json:"trigger"` // before_save etc.
	Type    string `json:"type"`    // record type
	Name    string `json:"name"`    // hook name
}

type timerInfo struct {
	Name string `json:"name"`
	Spec string `json:"spec"`
}

type providerInfo struct {
	Type string `json:"type"`
	Name string `json:"id"`
}

type registrationInfo struct {
	Handlers  []pluginHandlerInfo      `json:"handler"`
	Hooks     []pluginHookInfo         `json:"hook"`
	Lambdas   []map[string]interface{} `json:"op"`
	Timers    []timerInfo              `json:"timer"`
	Providers []providerInfo           `json:"provider"`
}

var transportFactories = map[string]TransportFactory{}

// RegisterTransport registers a transport factory by name.
func RegisterTransport(name string, transport TransportFactory) {
	transportFactories[name] = transport
}

// SupportedTransports tells all supported transport names
func SupportedTransports() []string {
	var transports []string
	for name := range transportFactories {
		transports = append(transports, name)
	}
	return transports
}

func unregisterAllTransports() {
	transportFactories = map[string]TransportFactory{}
}

// NewPlugin creates an instance of Plugin by transport and configuration.
func NewPlugin(name string, path string, args []string, config skyconfig.Configuration) Plugin {
	factory := transportFactories[name]
	if factory == nil {
		panic(fmt.Errorf("unable to find plugin transport '%v'", name))
	}
	p := Plugin{
		transport:  factory.Open(path, args, config),
		gatewayMap: map[string]*router.Gateway{},
	}
	return p
}

// Context contains reference to structs that will be initialized by plugin.
type Context struct {
	plugins          []*Plugin
	Router           *router.Router
	Mux              *http.ServeMux
	HandlerInjector  router.HandlerInjector
	HookRegistry     *hook.Registry
	ProviderRegistry *provider.Registry
	Scheduler        *cron.Cron
	Config           skyconfig.Configuration
}

// AddPluginConfiguration creates and appends a plugin
func (c *Context) AddPluginConfiguration(name string, path string, args []string) *Plugin {
	plug := NewPlugin(name, path, args, c.Config)
	c.plugins = append(c.plugins, &plug)
	return &plug
}

func (c *Context) getInitPayload() ([]byte, error) {
	payload := struct {
		Config skyconfig.Configuration `json:"config"`
	}{c.Config}

	return json.Marshal(payload)
}

// InitPlugins initializes all plugins registered
func (c *Context) InitPlugins() {
	wg := sync.WaitGroup{}
	for _, eachPlugin := range c.plugins {
		wg.Add(1)
		go func(plug *Plugin) {
			defer wg.Done()
			plug.Init(c)
		}(eachPlugin)
	}

	// make a goroutine to wait for all the plugins to be ready
	go func() {
		log.
			WithField("count", len(c.plugins)).
			Info("Wait for all plugin configurations")
		wg.Wait()

		data, err := c.getInitPayload()
		if err != nil {
			log.WithError(err).Warning("Fail to get init payload")
			data = []byte{}
		}
		c.SendEvent("before-plugins-ready", data, false)

		for _, eachPlugin := range c.plugins {
			eachPlugin.transport.SetState(TransportStateReady)
		}

		c.SendEvent("after-plugins-ready", data, false)
		c.SendEvent("server-ready", data, false)
	}()
}

// IsInitialized returns true if all the plugins have been initialized
func (c *Context) IsInitialized() bool {
	for _, eachPlugin := range c.plugins {
		if !eachPlugin.IsInitialized() {
			return false
		}
	}
	return true
}

// IsReady returns true if all the plugins are ready for client requests
func (c *Context) IsReady() bool {
	for _, eachPlugin := range c.plugins {
		if !eachPlugin.IsReady() {
			return false
		}
	}
	return true
}

// SendEvent sends event to all plugins
//
// SendEvent accepts `async` flag. Setting `async` to `false` means that
// an event will be sent to a plugin after another.
func (c *Context) SendEvent(name string, data []byte, async bool) {
	sendEventFunc := func(plugin *Plugin, name string, data []byte) {
		plugin.transport.SendEvent(name, data)
	}

	for _, eachPlugin := range c.plugins {
		if async {
			go sendEventFunc(eachPlugin, name, data)
		} else {
			sendEventFunc(eachPlugin, name, data)
		}
	}
}

// Init instantiates a plugin. This sets up hooks and handlers.
func (p *Plugin) Init(context *Context) {
	data, err := context.getInitPayload()
	if err != nil {
		log.WithError(err).Panic("Fail to get init payload")
	}

	transport := p.transport
	if bidirectional, ok := transport.(BidirectionalTransport); ok {
		bidirectional.SetRouter(context.Router)
	}

	transport.SendEvent("before-config", data)
	for {
		log.
			WithField("retry", p.initRetryCount).
			Info("Sending init event to plugin")

		transport.SetState(TransportStateUninitialized)
		regInfo, err := p.requestInit(data)
		if err != nil {
			transport.SetState(TransportStateError)

			p.initRetryCount++
			if p.initRetryCount >= PluginInitMaxRetryCount {
				log.Panic("Fail to initialize plugin")
			}
			time.Sleep(2 * time.Second)
			continue
		}

		p.processRegistrationInfo(context, regInfo)
		transport.SetState(TransportStateInitialized)

		break
	}
	transport.SendEvent("after-config", data)
}

func (p *Plugin) requestInit(data []byte) (regInfo registrationInfo, initErr error) {
	out, err := p.transport.SendEvent("init", data)
	log.WithFields(logrus.Fields{
		"out":    string(out),
		"err":    err,
		"plugin": p,
	}).Info("Get response from init")

	if err != nil {
		initErr = fmt.Errorf("Cannot encode plugin initialization payload. Error: %v", err)
		return
	}

	if err := json.Unmarshal(out, &regInfo); err != nil {
		initErr = fmt.Errorf("Unable to decode plugin initialization info. Error: %v", err)
		return
	}

	return
}

// IsInitialized returns true if the plugin has been initialized
func (p *Plugin) IsInitialized() bool {
	transportState := p.transport.State()
	return transportState == TransportStateInitialized ||
		transportState == TransportStateReady
}

// IsReady returns true if the plugin is ready for client request
func (p *Plugin) IsReady() bool {
	return p.transport.State() == TransportStateReady
}

func (p *Plugin) processRegistrationInfo(context *Context, regInfo registrationInfo) {
	log.WithFields(logrus.Fields{
		"regInfo":   regInfo,
		"transport": p.transport,
	}).Debugln("Got configuration from plugin, registering")
	p.initHandler(context.Mux, context.HandlerInjector, regInfo.Handlers, context.Config)
	p.initLambda(context.Router, context.HandlerInjector, regInfo.Lambdas)
	p.initHook(context.HookRegistry, regInfo.Hooks)
	if context.Scheduler != nil {
		p.initTimer(context.Scheduler, regInfo.Timers)
	} else {
		log.Info("Ignoring scheduled cron jobs because server is in slave mode.")
	}
	p.initProvider(context.ProviderRegistry, regInfo.Providers)
}

func (p *Plugin) initHandler(mux *http.ServeMux, injector router.HandlerInjector, handlers []pluginHandlerInfo, config skyconfig.Configuration) {
	for _, handler := range handlers {
		h := NewPluginHandler(handler, p)
		injector.Inject(h)
		h.Setup()
		name := h.Name
		name = strings.Replace(name, ":", "/", -1)
		if !strings.HasPrefix(name, "/") {
			name = "/" + name
		}
		var handlerGateway *router.Gateway
		handlerGateway, ok := p.gatewayMap[name]
		if !ok {
			handlerGateway = router.NewGateway("", name, "plugin", mux)
			handlerGateway.ResponseTimeout = time.Duration(config.App.ResponseTimeout) * time.Second
			p.gatewayMap[name] = handlerGateway
		}
		for _, method := range handler.Methods {
			handlerGateway.Handle(method, h)
		}
		log.Debugf(`Registered handler "%s" with serveMux at path "%s"`, h.Name, name)
	}
}

func (p *Plugin) initLambda(r *router.Router, injector router.HandlerInjector, lambdas []map[string]interface{}) {
	for _, lambda := range lambdas {
		handler := NewLambdaHandler(lambda, p)
		injector.Inject(handler)
		handler.Setup()
		r.Map(handler.Name, "plugin", handler)
		log.Debugf(`Registered lambda "%s" with router.`, handler.Name)
	}
}

func (p *Plugin) initHook(registry *hook.Registry, hookInfos []pluginHookInfo) {
	for _, hookInfo := range hookInfos {
		kind := hook.Kind(hookInfo.Trigger)
		recordType := hookInfo.Type

		registry.Register(kind, recordType, CreateHookFunc(p, hookInfo))
	}
}

func (p *Plugin) initTimer(c *cron.Cron, timerInfos []timerInfo) {
	for _, timerInfo := range timerInfos {
		timerName := timerInfo.Name
		err := c.AddFunc(timerInfo.Spec, func() {
			output, _ := p.transport.RunTimer(timerName, []byte{})
			log.Debugf("Executed a timer{%v} with result: %s", timerName, output)
		})

		if err != nil {
			panic(fmt.Errorf(`unable to add timer for "%s": %s`, timerName, err))
		}
	}
}

func (p *Plugin) initProvider(registry *provider.Registry, providerInfos []providerInfo) {
	for _, providerInfo := range providerInfos {
		provider := NewAuthProvider(providerInfo.Name, p)
		registry.RegisterAuthProvider(providerInfo.Name, provider)
	}
}
<commit_msg>Fix fatal error, concurrent map iteration and map write when init plugins<commit_after>// Copyright 2015-present Oursky Ltd.
//
// 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 plugin

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"github.com/robfig/cron"
	"github.com/skygeario/skygear-server/pkg/server/logging"
	"github.com/skygeario/skygear-server/pkg/server/plugin/hook"
	"github.com/skygeario/skygear-server/pkg/server/plugin/provider"
	"github.com/skygeario/skygear-server/pkg/server/router"
	"github.com/skygeario/skygear-server/pkg/server/skyconfig"
)

var log = logging.LoggerEntry("plugin")

const (
	// PluginInitMaxRetryCount defines the maximum retries for plugin initialization
	PluginInitMaxRetryCount = 100
)

// Plugin represents a collection of handlers, hooks and lambda functions
// that extends or modifies functionality provided by skygear.
type Plugin struct {
	initRetryCount int
	transport      Transport
	gatewayMap     map[string]*router.Gateway
}

type pluginHandlerInfo struct {
	AuthRequired bool     `json:"auth_required"`
	Name         string   `json:"name"`
	Methods      []string `json:"methods"`
	KeyRequired  bool     `json:"key_required"`
	UserRequired bool     `json:"user_required"`
}

type pluginHookInfo struct {
	Async   bool   `json:"async"`   // execute hook asynchronously
	Trigger string `json:"trigger"` // before_save etc.
	Type    string `json:"type"`    // record type
	Name    string `json:"name"`    // hook name
}

type timerInfo struct {
	Name string `json:"name"`
	Spec string `json:"spec"`
}

type providerInfo struct {
	Type string `json:"type"`
	Name string `json:"id"`
}

type registrationInfo struct {
	Handlers  []pluginHandlerInfo      `json:"handler"`
	Hooks     []pluginHookInfo         `json:"hook"`
	Lambdas   []map[string]interface{} `json:"op"`
	Timers    []timerInfo              `json:"timer"`
	Providers []providerInfo           `json:"provider"`
}

var transportFactories = map[string]TransportFactory{}

// RegisterTransport registers a transport factory by name.
func RegisterTransport(name string, transport TransportFactory) {
	transportFactories[name] = transport
}

// SupportedTransports tells all supported transport names
func SupportedTransports() []string {
	var transports []string
	for name := range transportFactories {
		transports = append(transports, name)
	}
	return transports
}

func unregisterAllTransports() {
	transportFactories = map[string]TransportFactory{}
}

// NewPlugin creates an instance of Plugin by transport and configuration.
func NewPlugin(name string, path string, args []string, config skyconfig.Configuration) Plugin {
	factory := transportFactories[name]
	if factory == nil {
		panic(fmt.Errorf("unable to find plugin transport '%v'", name))
	}
	p := Plugin{
		transport:  factory.Open(path, args, config),
		gatewayMap: map[string]*router.Gateway{},
	}
	return p
}

// Context contains reference to structs that will be initialized by plugin.
type Context struct {
	plugins          []*Plugin
	Router           *router.Router
	Mux              *http.ServeMux
	HandlerInjector  router.HandlerInjector
	HookRegistry     *hook.Registry
	ProviderRegistry *provider.Registry
	Scheduler        *cron.Cron
	Config           skyconfig.Configuration
	sync.Mutex
}

// AddPluginConfiguration creates and appends a plugin
func (c *Context) AddPluginConfiguration(name string, path string, args []string) *Plugin {
	plug := NewPlugin(name, path, args, c.Config)
	c.plugins = append(c.plugins, &plug)
	return &plug
}

func (c *Context) getInitPayload() ([]byte, error) {
	payload := struct {
		Config skyconfig.Configuration `json:"config"`
	}{c.Config}

	return json.Marshal(payload)
}

// InitPlugins initializes all plugins registered
func (c *Context) InitPlugins() {
	wg := sync.WaitGroup{}
	for _, eachPlugin := range c.plugins {
		wg.Add(1)
		go func(plug *Plugin) {
			defer wg.Done()
			plug.Init(c)
		}(eachPlugin)
	}

	// make a goroutine to wait for all the plugins to be ready
	go func() {
		log.
			WithField("count", len(c.plugins)).
			Info("Wait for all plugin configurations")
		wg.Wait()

		data, err := c.getInitPayload()
		if err != nil {
			log.WithError(err).Warning("Fail to get init payload")
			data = []byte{}
		}
		c.SendEvent("before-plugins-ready", data, false)

		for _, eachPlugin := range c.plugins {
			eachPlugin.transport.SetState(TransportStateReady)
		}

		c.SendEvent("after-plugins-ready", data, false)
		c.SendEvent("server-ready", data, false)
	}()
}

// IsInitialized returns true if all the plugins have been initialized
func (c *Context) IsInitialized() bool {
	for _, eachPlugin := range c.plugins {
		if !eachPlugin.IsInitialized() {
			return false
		}
	}
	return true
}

// IsReady returns true if all the plugins are ready for client requests
func (c *Context) IsReady() bool {
	for _, eachPlugin := range c.plugins {
		if !eachPlugin.IsReady() {
			return false
		}
	}
	return true
}

// SendEvent sends event to all plugins
//
// SendEvent accepts `async` flag. Setting `async` to `false` means that
// an event will be sent to a plugin after another.
func (c *Context) SendEvent(name string, data []byte, async bool) {
	sendEventFunc := func(plugin *Plugin, name string, data []byte) {
		plugin.transport.SendEvent(name, data)
	}

	for _, eachPlugin := range c.plugins {
		if async {
			go sendEventFunc(eachPlugin, name, data)
		} else {
			sendEventFunc(eachPlugin, name, data)
		}
	}
}

// Init instantiates a plugin. This sets up hooks and handlers.
func (p *Plugin) Init(context *Context) {
	data, err := context.getInitPayload()
	if err != nil {
		log.WithError(err).Panic("Fail to get init payload")
	}

	transport := p.transport
	if bidirectional, ok := transport.(BidirectionalTransport); ok {
		bidirectional.SetRouter(context.Router)
	}

	transport.SendEvent("before-config", data)
	for {
		log.
			WithField("retry", p.initRetryCount).
			Info("Sending init event to plugin")

		transport.SetState(TransportStateUninitialized)
		regInfo, err := p.requestInit(data)
		if err != nil {
			transport.SetState(TransportStateError)

			p.initRetryCount++
			if p.initRetryCount >= PluginInitMaxRetryCount {
				log.Panic("Fail to initialize plugin")
			}
			time.Sleep(2 * time.Second)
			continue
		}

		p.processRegistrationInfo(context, regInfo)
		transport.SetState(TransportStateInitialized)

		break
	}
	transport.SendEvent("after-config", data)
}

func (p *Plugin) requestInit(data []byte) (regInfo registrationInfo, initErr error) {
	out, err := p.transport.SendEvent("init", data)
	log.WithFields(logrus.Fields{
		"out":    string(out),
		"err":    err,
		"plugin": p,
	}).Info("Get response from init")

	if err != nil {
		initErr = fmt.Errorf("Cannot encode plugin initialization payload. Error: %v", err)
		return
	}

	if err := json.Unmarshal(out, &regInfo); err != nil {
		initErr = fmt.Errorf("Unable to decode plugin initialization info. Error: %v", err)
		return
	}

	return
}

// IsInitialized returns true if the plugin has been initialized
func (p *Plugin) IsInitialized() bool {
	transportState := p.transport.State()
	return transportState == TransportStateInitialized ||
		transportState == TransportStateReady
}

// IsReady returns true if the plugin is ready for client request
func (p *Plugin) IsReady() bool {
	return p.transport.State() == TransportStateReady
}

func (p *Plugin) processRegistrationInfo(context *Context, regInfo registrationInfo) {
	context.Lock()
	defer context.Unlock()
	log.WithFields(logrus.Fields{
		"regInfo":   regInfo,
		"transport": p.transport,
	}).Debugln("Got configuration from plugin, registering")
	p.initHandler(context.Mux, context.HandlerInjector, regInfo.Handlers, context.Config)
	p.initLambda(context.Router, context.HandlerInjector, regInfo.Lambdas)
	p.initHook(context.HookRegistry, regInfo.Hooks)
	if context.Scheduler != nil {
		p.initTimer(context.Scheduler, regInfo.Timers)
	} else {
		log.Info("Ignoring scheduled cron jobs because server is in slave mode.")
	}
	p.initProvider(context.ProviderRegistry, regInfo.Providers)
}

func (p *Plugin) initHandler(mux *http.ServeMux, injector router.HandlerInjector, handlers []pluginHandlerInfo, config skyconfig.Configuration) {
	for _, handler := range handlers {
		h := NewPluginHandler(handler, p)
		injector.Inject(h)
		h.Setup()
		name := h.Name
		name = strings.Replace(name, ":", "/", -1)
		if !strings.HasPrefix(name, "/") {
			name = "/" + name
		}
		var handlerGateway *router.Gateway
		handlerGateway, ok := p.gatewayMap[name]
		if !ok {
			handlerGateway = router.NewGateway("", name, "plugin", mux)
			handlerGateway.ResponseTimeout = time.Duration(config.App.ResponseTimeout) * time.Second
			p.gatewayMap[name] = handlerGateway
		}
		for _, method := range handler.Methods {
			handlerGateway.Handle(method, h)
		}
		log.Debugf(`Registered handler "%s" with serveMux at path "%s"`, h.Name, name)
	}
}

func (p *Plugin) initLambda(r *router.Router, injector router.HandlerInjector, lambdas []map[string]interface{}) {
	for _, lambda := range lambdas {
		handler := NewLambdaHandler(lambda, p)
		injector.Inject(handler)
		handler.Setup()
		r.Map(handler.Name, "plugin", handler)
		log.Debugf(`Registered lambda "%s" with router.`, handler.Name)
	}
}

func (p *Plugin) initHook(registry *hook.Registry, hookInfos []pluginHookInfo) {
	for _, hookInfo := range hookInfos {
		kind := hook.Kind(hookInfo.Trigger)
		recordType := hookInfo.Type

		registry.Register(kind, recordType, CreateHookFunc(p, hookInfo))
	}
}

func (p *Plugin) initTimer(c *cron.Cron, timerInfos []timerInfo) {
	for _, timerInfo := range timerInfos {
		timerName := timerInfo.Name
		err := c.AddFunc(timerInfo.Spec, func() {
			output, _ := p.transport.RunTimer(timerName, []byte{})
			log.Debugf("Executed a timer{%v} with result: %s", timerName, output)
		})

		if err != nil {
			panic(fmt.Errorf(`unable to add timer for "%s": %s`, timerName, err))
		}
	}
}

func (p *Plugin) initProvider(registry *provider.Registry, providerInfos []providerInfo) {
	for _, providerInfo := range providerInfos {
		provider := NewAuthProvider(providerInfo.Name, p)
		registry.RegisterAuthProvider(providerInfo.Name, provider)
	}
}
<|endoftext|>


<commit_before>package store_test

import (
	"encoding/json"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/store"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"time"
)

func (s *StoreSuite) prepareServer(c *C) (*store.Server, *charm.URL) {
	curl := charm.MustParseURL("cs:oneiric/wordpress")
	pub, err := s.store.CharmPublisher([]*charm.URL{curl}, "some-digest")
	c.Assert(err, IsNil)
	err = pub.Publish(&FakeCharmDir{})
	c.Assert(err, IsNil)

	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	return server, curl
}

func (s *StoreSuite) TestServerCharmInfo(c *C) {
	server, curl := s.prepareServer(c)
	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)

	var tests = []struct{ url, sha, err string }{
		{curl.String(), fakeRevZeroSha, ""},
		{"cs:oneiric/non-existent", "", "entry not found"},
		{"cs:bad", "", `charm URL without series: "cs:bad"`},
	}

	for _, t := range tests {
		req.Form = url.Values{"charms": []string{t.url}}
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		expected := make(map[string]interface{})
		if t.sha != "" {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"sha256": t.sha,
			}
		} else {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"errors": []interface{}{t.err},
			}
		}
		obtained := map[string]interface{}{}
		err = json.NewDecoder(rec.Body).Decode(&obtained)
		c.Assert(err, IsNil)
		c.Assert(obtained, DeepEquals, expected)
		c.Assert(rec.Header().Get("Content-Type"), Equals, "application/json")
	}

	s.checkCounterSum(c, []string{"charm-info", curl.Series, curl.Name}, false, 1)
	s.checkCounterSum(c, []string{"charm-missing", "oneiric", "non-existent"}, false, 1)
}

// checkCounterSum checks that statistics are properly collected.
// It retries a few times as they are generally collected in background.
func (s *StoreSuite) checkCounterSum(c *C, key []string, prefix bool, expected int64) {
	var sum int64
	var err error
	for retry := 0; retry < 10; retry++ {
		time.Sleep(1e8)
		sum, err = s.store.SumCounter(key, prefix)
		c.Assert(err, IsNil)
		if sum == expected {
			if expected == 0 && retry < 5 {
				continue // Wait a bit to make sure.
			}
			return
		}
	}
	c.Errorf("counter sum for %#v is %d, want %d", key, sum, expected)
}

func (s *StoreSuite) TestCharmStreaming(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "charm-revision-0")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "application/octet-stream")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "16")

	// Check that it was accounted for in statistics.
	s.checkCounterSum(c, []string{"charm-bundle", curl.Series, curl.Name}, false, 1)
}

func (s *StoreSuite) TestDisableStats(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"charms": []string{curl.String()}, "stats": []string{"0"}}
	server.ServeHTTP(httptest.NewRecorder(), req)

	req, err = http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"stats": []string{"0"}}
	server.ServeHTTP(httptest.NewRecorder(), req)

	// No statistics should have been collected given the use of stats=0.
	for _, prefix := range []string{"charm-info", "charm-bundle", "charm-missing"} {
		s.checkCounterSum(c, []string{prefix}, true, 0)
	}
}

func (s *StoreSuite) TestServerStatus(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	tests := []struct {
		path string
		code int
	}{
		{"/charm-info/any", 404},
		{"/charm/bad-url", 404},
		{"/charm/bad-series/wordpress", 404},
		{"/stats/counter/", 403},
		{"/stats/counter/*", 403},
		{"/stats/counter/any/", 404},
		{"/stats/", 404},
		{"/stats/any", 404},
	}
	for _, test := range tests {
		req, err := http.NewRequest("GET", test.path, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)
		c.Assert(rec.Code, Equals, test.code, Commentf("Path: %s", test.path))
	}
}

func (s *StoreSuite) TestRootRedirect(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	req, err := http.NewRequest("GET", "/", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)
	c.Assert(rec.Code, Equals, 303)
	c.Assert(rec.Header().Get("Location"), Equals, "https://juju.ubuntu.com")
}

func (s *StoreSuite) TestStatsCounter(c *C) {
	for _, key := range [][]string{{"a", "b"}, {"a", "b"}, {"a"}} {
		err := s.store.IncCounter(key)
		c.Assert(err, IsNil)
	}

	server, _ := s.prepareServer(c)

	expected := map[string]string{
		"a:b": "2",
		"a:*": "3",
		"a":   "1",
	}

	for counter, n := range expected {
		req, err := http.NewRequest("GET", "/stats/counter/" + counter, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		data, err := ioutil.ReadAll(rec.Body)
		c.Assert(string(data), Equals, n)

		c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
		c.Assert(rec.Header().Get("Content-Length"), Equals, strconv.Itoa(len(n)))
	}
}

func (s *StoreSuite) TestBlitzKey(c *C) {
	server, _ := s.prepareServer(c)

	// This is just a validation key to allow blitz.io to run
	// performance tests against the site.
	req, err := http.NewRequest("GET", "/mu-35700a31-6bf320ca-a800b670-05f845ee", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "42")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "2")
}
<commit_msg>Reduce wait for absence of counter sums to 200 milliseconds.<commit_after>package store_test

import (
	"encoding/json"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/store"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"time"
)

func (s *StoreSuite) prepareServer(c *C) (*store.Server, *charm.URL) {
	curl := charm.MustParseURL("cs:oneiric/wordpress")
	pub, err := s.store.CharmPublisher([]*charm.URL{curl}, "some-digest")
	c.Assert(err, IsNil)
	err = pub.Publish(&FakeCharmDir{})
	c.Assert(err, IsNil)

	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	return server, curl
}

func (s *StoreSuite) TestServerCharmInfo(c *C) {
	server, curl := s.prepareServer(c)
	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)

	var tests = []struct{ url, sha, err string }{
		{curl.String(), fakeRevZeroSha, ""},
		{"cs:oneiric/non-existent", "", "entry not found"},
		{"cs:bad", "", `charm URL without series: "cs:bad"`},
	}

	for _, t := range tests {
		req.Form = url.Values{"charms": []string{t.url}}
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		expected := make(map[string]interface{})
		if t.sha != "" {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"sha256": t.sha,
			}
		} else {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"errors": []interface{}{t.err},
			}
		}
		obtained := map[string]interface{}{}
		err = json.NewDecoder(rec.Body).Decode(&obtained)
		c.Assert(err, IsNil)
		c.Assert(obtained, DeepEquals, expected)
		c.Assert(rec.Header().Get("Content-Type"), Equals, "application/json")
	}

	s.checkCounterSum(c, []string{"charm-info", curl.Series, curl.Name}, false, 1)
	s.checkCounterSum(c, []string{"charm-missing", "oneiric", "non-existent"}, false, 1)
}

// checkCounterSum checks that statistics are properly collected.
// It retries a few times as they are generally collected in background.
func (s *StoreSuite) checkCounterSum(c *C, key []string, prefix bool, expected int64) {
	var sum int64
	var err error
	for retry := 0; retry < 10; retry++ {
		time.Sleep(1e8)
		sum, err = s.store.SumCounter(key, prefix)
		c.Assert(err, IsNil)
		if sum == expected {
			if expected == 0 && retry < 2 {
				continue // Wait a bit to make sure.
			}
			return
		}
	}
	c.Errorf("counter sum for %#v is %d, want %d", key, sum, expected)
}

func (s *StoreSuite) TestCharmStreaming(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "charm-revision-0")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "application/octet-stream")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "16")

	// Check that it was accounted for in statistics.
	s.checkCounterSum(c, []string{"charm-bundle", curl.Series, curl.Name}, false, 1)
}

func (s *StoreSuite) TestDisableStats(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"charms": []string{curl.String()}, "stats": []string{"0"}}
	server.ServeHTTP(httptest.NewRecorder(), req)

	req, err = http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"stats": []string{"0"}}
	server.ServeHTTP(httptest.NewRecorder(), req)

	// No statistics should have been collected given the use of stats=0.
	for _, prefix := range []string{"charm-info", "charm-bundle", "charm-missing"} {
		s.checkCounterSum(c, []string{prefix}, true, 0)
	}
}

func (s *StoreSuite) TestServerStatus(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	tests := []struct {
		path string
		code int
	}{
		{"/charm-info/any", 404},
		{"/charm/bad-url", 404},
		{"/charm/bad-series/wordpress", 404},
		{"/stats/counter/", 403},
		{"/stats/counter/*", 403},
		{"/stats/counter/any/", 404},
		{"/stats/", 404},
		{"/stats/any", 404},
	}
	for _, test := range tests {
		req, err := http.NewRequest("GET", test.path, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)
		c.Assert(rec.Code, Equals, test.code, Commentf("Path: %s", test.path))
	}
}

func (s *StoreSuite) TestRootRedirect(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	req, err := http.NewRequest("GET", "/", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)
	c.Assert(rec.Code, Equals, 303)
	c.Assert(rec.Header().Get("Location"), Equals, "https://juju.ubuntu.com")
}

func (s *StoreSuite) TestStatsCounter(c *C) {
	for _, key := range [][]string{{"a", "b"}, {"a", "b"}, {"a"}} {
		err := s.store.IncCounter(key)
		c.Assert(err, IsNil)
	}

	server, _ := s.prepareServer(c)

	expected := map[string]string{
		"a:b": "2",
		"a:*": "3",
		"a":   "1",
	}

	for counter, n := range expected {
		req, err := http.NewRequest("GET", "/stats/counter/" + counter, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		data, err := ioutil.ReadAll(rec.Body)
		c.Assert(string(data), Equals, n)

		c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
		c.Assert(rec.Header().Get("Content-Length"), Equals, strconv.Itoa(len(n)))
	}
}

func (s *StoreSuite) TestBlitzKey(c *C) {
	server, _ := s.prepareServer(c)

	// This is just a validation key to allow blitz.io to run
	// performance tests against the site.
	req, err := http.NewRequest("GET", "/mu-35700a31-6bf320ca-a800b670-05f845ee", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "42")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "2")
}
<|endoftext|>


<commit_before>package nntp

import (
	"testing"

	"github.com/diocles/nntp/nntptest"
)

func TestServerHandler(t *testing.T) {
	handler := serverHandler{}

	w := nntptest.NewRecorder()

	c := Command{Keyword: "HELP"}
	handler.ServeNNTP(w, &c)

	if w.Code != 100 {
		t.Errorf("HELP should return status 100")
	}
	if w.Body.String() == "" {
		t.Errorf("HELP body should not be empty")
	}
}
<commit_msg>Add license header<commit_after>// Copyright (C) 2015 Tim Retout <tim@retout.co.uk>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package nntp

import (
	"testing"

	"github.com/diocles/nntp/nntptest"
)

func TestServerHandler(t *testing.T) {
	handler := serverHandler{}

	w := nntptest.NewRecorder()

	c := Command{Keyword: "HELP"}
	handler.ServeNNTP(w, &c)

	if w.Code != 100 {
		t.Errorf("HELP should return status 100")
	}
	if w.Body.String() == "" {
		t.Errorf("HELP body should not be empty")
	}
}
<|endoftext|>


<commit_before>package store

import (
	"sync"
	"testing"

	"github.com/hellofresh/janus/pkg/notifier"
	"github.com/stretchr/testify/assert"
)

const (
	testKey   = "key"
	testValue = "value"
)

func TestInMemoryStore(t *testing.T) {
	t.Parallel()

	tests := []struct {
		scenario string
		function func(*testing.T, *InMemoryStore)
	}{
		{
			scenario: "new in memory store",
			function: testNewInMemoryStore,
		},
		{
			scenario: "set in memory store",
			function: testInMemoryStoreSet,
		},
		{
			scenario: "check exists from in memory store",
			function: testInMemoryStoreExists,
		},
		{
			scenario: "get from in memory store",
			function: testInMemoryStoreGet,
		},
		{
			scenario: "remove from in memory store",
			function: testInMemoryStoreRemove,
		},
		{
			scenario: "publish subscribe into memory store",
			function: testInMemoryStorePublishSubscribe,
		},
	}

	for _, test := range tests {
		t.Run(test.scenario, func(t *testing.T) {
			instance := NewInMemoryStore()
			test.function(t, instance)
		})
	}
}

func testNewInMemoryStore(t *testing.T, instance *InMemoryStore) {
	assert.IsType(t, &InMemoryStore{}, instance)
	assert.Implements(t, (*Store)(nil), instance)
}

func testInMemoryStoreSet(t *testing.T, instance *InMemoryStore) {
	err := instance.Set(testKey, testValue, 0)
	assert.Nil(t, err)
}

func testInMemoryStoreExists(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	val, err := instance.Exists("foo")
	assert.Nil(t, err)
	assert.False(t, val)

	val, err = instance.Exists(testKey)
	assert.Nil(t, err)
	assert.True(t, val)
}

func testInMemoryStoreGet(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	val, err := instance.Get(testKey)
	assert.Nil(t, err)
	assert.Equal(t, testValue, val)

	val, err = instance.Get("foo")
	assert.Nil(t, err)
	assert.Empty(t, val)
}

func testInMemoryStoreRemove(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	val, err := instance.Get(testKey)
	assert.Nil(t, err)
	assert.Equal(t, testValue, val)

	instance.Remove(testKey)

	val, err = instance.Get(testKey)
	assert.Nil(t, err)
	assert.Empty(t, val)
}

func testInMemoryStorePublishSubscribe(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	var (
		wg     sync.WaitGroup
		called bool
	)

	wg.Add(1)

	go func() {
		instance.Subscribe("test", func(notifier.Notification) {
			called = true
			wg.Done()
		})
	}()

	instance.Publish("test", []byte(`foo`))

	wg.Wait()

	assert.True(t, called)
}
<commit_msg>Fix test<commit_after>package store

import (
	"sync"
	"testing"

	"github.com/hellofresh/janus/pkg/notifier"
	"github.com/stretchr/testify/assert"
)

const (
	testKey   = "key"
	testValue = "value"
)

func TestInMemoryStore(t *testing.T) {
	t.Parallel()

	tests := []struct {
		scenario string
		function func(*testing.T, *InMemoryStore)
	}{
		{
			scenario: "new in memory store",
			function: testNewInMemoryStore,
		},
		{
			scenario: "set in memory store",
			function: testInMemoryStoreSet,
		},
		{
			scenario: "check exists from in memory store",
			function: testInMemoryStoreExists,
		},
		{
			scenario: "get from in memory store",
			function: testInMemoryStoreGet,
		},
		{
			scenario: "remove from in memory store",
			function: testInMemoryStoreRemove,
		},
		{
			scenario: "publish subscribe into memory store",
			function: testInMemoryStorePublishSubscribe,
		},
	}

	for _, test := range tests {
		t.Run(test.scenario, func(t *testing.T) {
			instance := NewInMemoryStore()
			test.function(t, instance)
		})
	}
}

func testNewInMemoryStore(t *testing.T, instance *InMemoryStore) {
	assert.IsType(t, &InMemoryStore{}, instance)
	assert.Implements(t, (*Store)(nil), instance)
}

func testInMemoryStoreSet(t *testing.T, instance *InMemoryStore) {
	err := instance.Set(testKey, testValue, 0)
	assert.Nil(t, err)
}

func testInMemoryStoreExists(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	val, err := instance.Exists("foo")
	assert.Nil(t, err)
	assert.False(t, val)

	val, err = instance.Exists(testKey)
	assert.Nil(t, err)
	assert.True(t, val)
}

func testInMemoryStoreGet(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	val, err := instance.Get(testKey)
	assert.Nil(t, err)
	assert.Equal(t, testValue, val)

	val, err = instance.Get("foo")
	assert.Nil(t, err)
	assert.Empty(t, val)
}

func testInMemoryStoreRemove(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	val, err := instance.Get(testKey)
	assert.Nil(t, err)
	assert.Equal(t, testValue, val)

	instance.Remove(testKey)

	val, err = instance.Get(testKey)
	assert.Nil(t, err)
	assert.Empty(t, val)
}

func testInMemoryStorePublishSubscribe(t *testing.T, instance *InMemoryStore) {
	instance.Set(testKey, testValue, 0)

	var (
		wg     sync.WaitGroup
		called bool
	)

	wg.Add(1)

	go func() {
		instance.Subscribe("test", func(notifier.Notification) {
			called = true
			wg.Done()
		})
	}()

	instance.Publish("test", []byte(`{"foo": "bar"}`))

	wg.Wait()

	assert.True(t, called)
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC 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 stream

import (
	"archive/tar"
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"io/ioutil"
	"strings"
	"testing"

	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
)

func TestStreamVsBuffer(t *testing.T) {
	var n, wantSize int64 = 10000, 49
	newBlob := func() io.ReadCloser { return ioutil.NopCloser(bytes.NewReader(bytes.Repeat([]byte{'a'}, int(n)))) }
	wantDigest := "sha256:3d7c465be28d9e1ed810c42aeb0e747b44441424f566722ba635dc93c947f30e"
	wantDiffID := "sha256:27dd1f61b867b6a0f6e9d8a41c43231de52107e53ae424de8f847b821db4b711"

	// Check that streaming some content results in the expected digest/diffID/size.
	l := NewLayer(newBlob())
	if c, err := l.Compressed(); err != nil {
		t.Errorf("Compressed: %v", err)
	} else {
		if _, err := io.Copy(ioutil.Discard, c); err != nil {
			t.Errorf("error reading Compressed: %v", err)
		}
		if err := c.Close(); err != nil {
			t.Errorf("Close: %v", err)
		}
	}
	if d, err := l.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if d.String() != wantDigest {
		t.Errorf("stream Digest got %q, want %q", d.String(), wantDigest)
	}
	if d, err := l.DiffID(); err != nil {
		t.Errorf("DiffID: %v", err)
	} else if d.String() != wantDiffID {
		t.Errorf("stream DiffID got %q, want %q", d.String(), wantDiffID)
	}
	if s, err := l.Size(); err != nil {
		t.Errorf("Size: %v", err)
	} else if s != wantSize {
		t.Errorf("stream Size got %d, want %d", s, wantSize)
	}

	// Test that buffering the same contents and using
	// tarball.LayerFromOpener results in the same digest/diffID/size.
	tl, err := tarball.LayerFromOpener(func() (io.ReadCloser, error) { return newBlob(), nil })
	if err != nil {
		t.Fatalf("LayerFromOpener: %v", err)
	}
	if d, err := tl.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if d.String() != wantDigest {
		t.Errorf("tarball Digest got %q, want %q", d.String(), wantDigest)
	}
	if d, err := tl.DiffID(); err != nil {
		t.Errorf("DiffID: %v", err)
	} else if d.String() != wantDiffID {
		t.Errorf("tarball DiffID got %q, want %q", d.String(), wantDiffID)
	}
	if s, err := tl.Size(); err != nil {
		t.Errorf("Size: %v", err)
	} else if s != wantSize {
		t.Errorf("stream Size got %d, want %d", s, wantSize)
	}
}

func TestLargeStream(t *testing.T) {
	var n, wantSize int64 = 100000000, 100007653 // "Compressing" n random bytes results in this many bytes.
	sl := NewLayer(ioutil.NopCloser(io.LimitReader(rand.Reader, n)))
	rc, err := sl.Compressed()
	if err != nil {
		t.Fatalf("Uncompressed: %v", err)
	}
	if _, err := io.Copy(ioutil.Discard, rc); err != nil {
		t.Fatalf("Reading layer: %v", err)
	}
	if err := rc.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	if dig, err := sl.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if dig.String() == (v1.Hash{}).String() {
		t.Errorf("Digest got %q, want anything else", (v1.Hash{}).String())
	}
	if diffID, err := sl.DiffID(); err != nil {
		t.Errorf("DiffID: %v", err)
	} else if diffID.String() == (v1.Hash{}).String() {
		t.Errorf("DiffID got %q, want anything else", (v1.Hash{}).String())
	}
	if size, err := sl.Size(); err != nil {
		t.Errorf("Size: %v", err)
	} else if size != wantSize {
		t.Errorf("Size got %d, want %d", size, n)
	}
}

func TestStreamableLayerFromTarball(t *testing.T) {
	pr, pw := io.Pipe()
	tw := tar.NewWriter(pw)
	go func() {
		// "Stream" a bunch of files into the layer.
		pw.CloseWithError(func() error {
			for i := 0; i < 1000; i++ {
				name := fmt.Sprintf("file-%d.txt", i)
				body := fmt.Sprintf("i am file number %d", i)
				if err := tw.WriteHeader(&tar.Header{
					Name:     name,
					Mode:     0600,
					Size:     int64(len(body)),
					Typeflag: tar.TypeReg,
				}); err != nil {
					return err
				}
				if _, err := tw.Write([]byte(body)); err != nil {
					return err
				}
			}
			return nil
		}())
	}()

	l := NewLayer(pr)
	rc, err := l.Compressed()
	if err != nil {
		t.Fatalf("Compressed: %v", err)
	}
	if _, err := io.Copy(ioutil.Discard, rc); err != nil {
		t.Fatalf("Copy: %v", err)
	}
	if err := rc.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	wantDigest := "sha256:f53d6a164ab476294212843f267740bd12f79e00abd8050c24ce8a9bceaa36b0"
	if got, err := l.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if got.String() != wantDigest {
		t.Errorf("Digest: got %q, want %q", got.String(), wantDigest)
	}
}

// TestNotComputed tests that Digest/DiffID/Size return ErrNotComputed before
// the stream has been consumed.
func TestNotComputed(t *testing.T) {
	l := NewLayer(ioutil.NopCloser(bytes.NewBufferString("hi")))

	// All methods should return ErrNotComputed until the stream has been
	// consumed and closed.
	if _, err := l.Size(); err != ErrNotComputed {
		t.Errorf("Size: got %v, want %v", err, ErrNotComputed)
	}
	if _, err := l.Digest(); err == nil {
		t.Errorf("Digest: got %v, want %v", err, ErrNotComputed)
	}
	if _, err := l.DiffID(); err == nil {
		t.Errorf("DiffID: got %v, want %v", err, ErrNotComputed)
	}
}

// TestConsumed tests that Compressed returns ErrConsumed when the stream has
// already been consumed.
func TestConsumed(t *testing.T) {
	l := NewLayer(ioutil.NopCloser(strings.NewReader("hello")))
	rc, err := l.Compressed()
	if err != nil {
		t.Errorf("Compressed: %v", err)
	}
	if _, err := io.Copy(ioutil.Discard, rc); err != nil {
		t.Errorf("Error reading contents: %v", err)
	}
	if err := rc.Close(); err != nil {
		t.Errorf("Close: %v", err)
	}

	if _, err := l.Compressed(); err != ErrConsumed {
		t.Errorf("Compressed() after consuming; got %v, want %v", err, ErrConsumed)
	}
}
<commit_msg>Speed up TestLargeStream (#362)<commit_after>// Copyright 2018 Google LLC 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 stream

import (
	"archive/tar"
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"io/ioutil"
	"strings"
	"testing"

	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
)

func TestStreamVsBuffer(t *testing.T) {
	var n, wantSize int64 = 10000, 49
	newBlob := func() io.ReadCloser { return ioutil.NopCloser(bytes.NewReader(bytes.Repeat([]byte{'a'}, int(n)))) }
	wantDigest := "sha256:3d7c465be28d9e1ed810c42aeb0e747b44441424f566722ba635dc93c947f30e"
	wantDiffID := "sha256:27dd1f61b867b6a0f6e9d8a41c43231de52107e53ae424de8f847b821db4b711"

	// Check that streaming some content results in the expected digest/diffID/size.
	l := NewLayer(newBlob())
	if c, err := l.Compressed(); err != nil {
		t.Errorf("Compressed: %v", err)
	} else {
		if _, err := io.Copy(ioutil.Discard, c); err != nil {
			t.Errorf("error reading Compressed: %v", err)
		}
		if err := c.Close(); err != nil {
			t.Errorf("Close: %v", err)
		}
	}
	if d, err := l.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if d.String() != wantDigest {
		t.Errorf("stream Digest got %q, want %q", d.String(), wantDigest)
	}
	if d, err := l.DiffID(); err != nil {
		t.Errorf("DiffID: %v", err)
	} else if d.String() != wantDiffID {
		t.Errorf("stream DiffID got %q, want %q", d.String(), wantDiffID)
	}
	if s, err := l.Size(); err != nil {
		t.Errorf("Size: %v", err)
	} else if s != wantSize {
		t.Errorf("stream Size got %d, want %d", s, wantSize)
	}

	// Test that buffering the same contents and using
	// tarball.LayerFromOpener results in the same digest/diffID/size.
	tl, err := tarball.LayerFromOpener(func() (io.ReadCloser, error) { return newBlob(), nil })
	if err != nil {
		t.Fatalf("LayerFromOpener: %v", err)
	}
	if d, err := tl.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if d.String() != wantDigest {
		t.Errorf("tarball Digest got %q, want %q", d.String(), wantDigest)
	}
	if d, err := tl.DiffID(); err != nil {
		t.Errorf("DiffID: %v", err)
	} else if d.String() != wantDiffID {
		t.Errorf("tarball DiffID got %q, want %q", d.String(), wantDiffID)
	}
	if s, err := tl.Size(); err != nil {
		t.Errorf("Size: %v", err)
	} else if s != wantSize {
		t.Errorf("stream Size got %d, want %d", s, wantSize)
	}
}

func TestLargeStream(t *testing.T) {
	var n, wantSize int64 = 10000000, 10000788 // "Compressing" n random bytes results in this many bytes.
	sl := NewLayer(ioutil.NopCloser(io.LimitReader(rand.Reader, n)))
	rc, err := sl.Compressed()
	if err != nil {
		t.Fatalf("Uncompressed: %v", err)
	}
	if _, err := io.Copy(ioutil.Discard, rc); err != nil {
		t.Fatalf("Reading layer: %v", err)
	}
	if err := rc.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	if dig, err := sl.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if dig.String() == (v1.Hash{}).String() {
		t.Errorf("Digest got %q, want anything else", (v1.Hash{}).String())
	}
	if diffID, err := sl.DiffID(); err != nil {
		t.Errorf("DiffID: %v", err)
	} else if diffID.String() == (v1.Hash{}).String() {
		t.Errorf("DiffID got %q, want anything else", (v1.Hash{}).String())
	}
	if size, err := sl.Size(); err != nil {
		t.Errorf("Size: %v", err)
	} else if size != wantSize {
		t.Errorf("Size got %d, want %d", size, wantSize)
	}
}

func TestStreamableLayerFromTarball(t *testing.T) {
	pr, pw := io.Pipe()
	tw := tar.NewWriter(pw)
	go func() {
		// "Stream" a bunch of files into the layer.
		pw.CloseWithError(func() error {
			for i := 0; i < 1000; i++ {
				name := fmt.Sprintf("file-%d.txt", i)
				body := fmt.Sprintf("i am file number %d", i)
				if err := tw.WriteHeader(&tar.Header{
					Name:     name,
					Mode:     0600,
					Size:     int64(len(body)),
					Typeflag: tar.TypeReg,
				}); err != nil {
					return err
				}
				if _, err := tw.Write([]byte(body)); err != nil {
					return err
				}
			}
			return nil
		}())
	}()

	l := NewLayer(pr)
	rc, err := l.Compressed()
	if err != nil {
		t.Fatalf("Compressed: %v", err)
	}
	if _, err := io.Copy(ioutil.Discard, rc); err != nil {
		t.Fatalf("Copy: %v", err)
	}
	if err := rc.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	wantDigest := "sha256:f53d6a164ab476294212843f267740bd12f79e00abd8050c24ce8a9bceaa36b0"
	if got, err := l.Digest(); err != nil {
		t.Errorf("Digest: %v", err)
	} else if got.String() != wantDigest {
		t.Errorf("Digest: got %q, want %q", got.String(), wantDigest)
	}
}

// TestNotComputed tests that Digest/DiffID/Size return ErrNotComputed before
// the stream has been consumed.
func TestNotComputed(t *testing.T) {
	l := NewLayer(ioutil.NopCloser(bytes.NewBufferString("hi")))

	// All methods should return ErrNotComputed until the stream has been
	// consumed and closed.
	if _, err := l.Size(); err != ErrNotComputed {
		t.Errorf("Size: got %v, want %v", err, ErrNotComputed)
	}
	if _, err := l.Digest(); err == nil {
		t.Errorf("Digest: got %v, want %v", err, ErrNotComputed)
	}
	if _, err := l.DiffID(); err == nil {
		t.Errorf("DiffID: got %v, want %v", err, ErrNotComputed)
	}
}

// TestConsumed tests that Compressed returns ErrConsumed when the stream has
// already been consumed.
func TestConsumed(t *testing.T) {
	l := NewLayer(ioutil.NopCloser(strings.NewReader("hello")))
	rc, err := l.Compressed()
	if err != nil {
		t.Errorf("Compressed: %v", err)
	}
	if _, err := io.Copy(ioutil.Discard, rc); err != nil {
		t.Errorf("Error reading contents: %v", err)
	}
	if err := rc.Close(); err != nil {
		t.Errorf("Close: %v", err)
	}

	if _, err := l.Compressed(); err != ErrConsumed {
		t.Errorf("Compressed() after consuming; got %v, want %v", err, ErrConsumed)
	}
}
<|endoftext|>


<commit_before>package veneur

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stripe/veneur/tdigest"
)

const ε = .00002

var DebugMode bool

func TestMain(m *testing.M) {
	flag.Parse()
	DebugMode = flag.Lookup("test.v").Value.(flag.Getter).Get().(bool)
	os.Exit(m.Run())
}

// On the CI server, we can't be guaranteed that the port will be
// released immediately after the server is shut down. Instead, use
// a unique port for each test. As long as we don't have an insane number
// of integration tests, we should be fine.
var HttpAddrPort = 8127

var localHandler http.Handler

// set up a boilerplate local config for later use
func localConfig() Config {
  port := HttpAddrPort
  HttpAddrPort++

	return Config{
		APIHostname: "http://localhost",
		Debug:       DebugMode,
		Hostname:    "localhost",

		// Use a shorter interval for tests
		Interval:            50 * time.Millisecond,
		Key:                 "",
		MetricMaxLength:     4096,
		Percentiles:         []float64{.5, .75, .99},
		ReadBufferSizeBytes: 2097152,
		UDPAddr:             "localhost:8126",
		HTTPAddr:            fmt.Sprintf("localhost:%d", port),
		ForwardAddr:         "http://localhost",
		NumWorkers:          96,

		// Use only one reader, so that we can run tests
		// on platforms which do not support SO_REUSEPORT
		NumReaders: 1,

		// Currently this points nowhere, which is intentional.
		// We don't need internal metrics for the tests, and they make testing
		// more complicated.
		StatsAddr:  "localhost:8125",
		Tags:       []string{},
		SentryDSN:  "",
		FlushLimit: 1024,
	}
}

// assertMetrics checks that all expected metrics are present
// and have the correct value
func assertMetrics(t *testing.T, metrics DDMetricsRequest, expectedMetrics map[string]float64) {
	// it doesn't count as accidentally quadratic if it's intentional
	for metricName, expectedValue := range expectedMetrics {
		assertMetric(t, metrics, metricName, expectedValue)
	}
}

func assertMetric(t *testing.T, metrics DDMetricsRequest, metricName string, value float64) {
	defer func() {
		if r := recover(); r != nil {
			assert.Fail(t, "error extracting metrics", r)
		}
	}()
	for _, metric := range metrics.Series {
		if metric.Name == metricName {
			assert.Equal(t, int(value+.5), int(metric.Value[0][1]+.5), fmt.Sprintf("Incorrect value for metric %s", metricName))
			return
		}
	}
	assert.Fail(t, "did not find expected metric", metricName)
}

// setupLocalServer creates a local server from the specified config
// and starts listening for requests. It returns the server for inspection.
func setupLocalServer(t *testing.T, config Config) Server {
	server, err := NewFromConfig(config)
	if err != nil {
		t.Fatal(err)
	}

	packetPool := &sync.Pool{
		New: func() interface{} {
			return make([]byte, config.MetricMaxLength)
		},
	}

	for i := 0; i < config.NumReaders; i++ {
		go func() {
			defer func() {
				if r := recover(); r != nil {
					assert.Fail(t, "reader panicked while reading from socket", err)
				}
			}()
			server.ReadSocket(packetPool, config.NumReaders != 1)
		}()
	}

	go server.HTTPServe()
	return server
}

// DDMetricsRequest represents the body of the POST request
// for sending metrics data to Datadog
type DDMetricsRequest struct {
	Series []DDMetric
}

func TestLocalServerUnaggregatedMetrics(t *testing.T) {
	expectedMetrics := map[string]float64{
		"a.b.c.max": 100,
		"a.b.c.min": 1,

		// Count is normalized by second
		// so 5 values/50ms = 100 values/s
		"a.b.c.count": 100,

		// tdigest approximation causes this to be off by 1
		"a.b.c.50percentile": 6,
		"a.b.c.75percentile": 42,
		"a.b.c.99percentile": 98,
	}

	remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		zr, err := zlib.NewReader(r.Body)
		if err != nil {
			t.Fatal(err)
		}

		var ddmetrics DDMetricsRequest
		err = json.NewDecoder(zr).Decode(&ddmetrics)
		if err != nil {
			t.Fatal(err)
		}
		assert.Equal(t, 6, len(ddmetrics.Series), "incorrect number of elements in the flushed series on the remote server")
		assertMetrics(t, ddmetrics, expectedMetrics)
		w.WriteHeader(http.StatusAccepted)
	}))
	config := localConfig()
	config.APIHostname = remoteServer.URL

	server := setupLocalServer(t, config)
	defer server.Shutdown()

	metricValues := []float64{1.0, 2.0, 7.0, 8.0, 100.0}

	for _, value := range metricValues {
		server.Workers[0].ProcessMetric(&UDPMetric{
			MetricKey: MetricKey{
				Name: "a.b.c",
				Type: "histogram",
			},
			Value:      value,
			Digest:     12345,
			SampleRate: 1.0,
			LocalOnly:  true,
		})
	}
	server.Flush(config.Interval, config.FlushLimit)
}

func TestLocalServerMixedMetrics(t *testing.T) {

	// We can't be guaranteed that we will receive this exact stream
	// but we should expect that they represent the same underlying tdigest
	const ExpectedGobStream = "\r\xff\x87\x02\x01\x02\xff\x88\x00\x01\xff\x84\x00\x007\xff\x83\x03\x01\x01\bCentroid\x01\xff\x84\x00\x01\x03\x01\x04Mean\x01\b\x00\x01\x06Weight\x01\b\x00\x01\aSamples\x01\xff\x86\x00\x00\x00\x17\xff\x85\x02\x01\x01\t[]float64\x01\xff\x86\x00\x01\b\x00\x00/\xff\x88\x00\x05\x01\xfe\xf0?\x01\xfe\xf0?\x00\x01@\x01\xfe\xf0?\x00\x01\xfe\x1c@\x01\xfe\xf0?\x00\x01\xfe @\x01\xfe\xf0?\x00\x01\xfeY@\x01\xfe\xf0?\x00\x05\b\x00\xfeY@\x05\b\x00\xfe\xf0?\x05\b\x00\xfeY@"

	const ABCCountRaw = 5
	const ABCCountNormalized = 100

	expectedMetrics := map[string]float64{
		"x.y.z":     800,
		"a.b.c.max": 100,
		"a.b.c.min": 1,

		// Count is normalized by second
		// so 5 values/50ms = 100 values/s
		"a.b.c.count": ABCCountNormalized,
	}

	// Set up a remote server (the API that we're sending the data to)
	// e.g. Datadog
	remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		zr, err := zlib.NewReader(r.Body)
		if err != nil {
			t.Fatal(err)
		}

		var ddmetrics DDMetricsRequest

		err = json.NewDecoder(zr).Decode(&ddmetrics)
		if err != nil {
			t.Fatal(err)
		}
		assert.Equal(t, len(expectedMetrics), len(ddmetrics.Series), "incorrect number of elements in the flushed series on the remote server")
		assertMetrics(t, ddmetrics, expectedMetrics)
		w.WriteHeader(http.StatusAccepted)
	}))

	globalVeneur := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		zr, err := zlib.NewReader(r.Body)
		if err != nil {
			t.Fatal(err)
		}

		type requestItem struct {
			Name      string      `json:"name"`
			Tags      interface{} `json:"tags"`
			Tagstring string      `json:"tagstring"`
			Type      string      `json:"type"`
			Value     []byte      `json:"value"`
		}

		var metrics []requestItem

		err = json.NewDecoder(zr).Decode(&metrics)
		if err != nil {
			t.Fatal(err)
		}

		assert.Equal(t, 1, len(metrics), "incorrect number of elements in the flushed series")

		tdExpected := tdigest.NewMerging(100, false)
		err = tdExpected.GobDecode([]byte(ExpectedGobStream))
		assert.NoError(t, err, "Should not have encountered error in decoding expected gob stream")

		td := tdigest.NewMerging(100, false)
		err = td.GobDecode(metrics[0].Value)
		assert.NoError(t, err, "Should not have encountered error in decoding gob stream")

		assert.Equal(t, expectedMetrics["a.b.c.min"], td.Min(), "Minimum value is incorrect")
		assert.Equal(t, expectedMetrics["a.b.c.max"], td.Max(), "Maximum value is incorrect")

		// The remote server receives the raw count, *not* the normalized count
		assert.InEpsilon(t, ABCCountRaw, td.Count(), ε)
		assert.Equal(t, tdExpected, td, "Underlying tdigest structure is incorrect")

		w.WriteHeader(http.StatusAccepted)
	}))

	config := localConfig()
	config.APIHostname = remoteServer.URL
	config.ForwardAddr = globalVeneur.URL
	config.NumWorkers = 1

	server := setupLocalServer(t, config)
	defer server.Shutdown()

	histogramValues := []float64{1.0, 2.0, 7.0, 8.0, 100.0}

	// Create non-local metrics that should be passed to the global veneur instance
	for _, value := range histogramValues {
		server.Workers[0].ProcessMetric(&UDPMetric{
			MetricKey: MetricKey{
				Name: "a.b.c",
				Type: "histogram",
			},
			Value:      value,
			Digest:     12345,
			SampleRate: 1.0,
			LocalOnly:  false,
		})
	}

	const CountSize = 40
	for i := 0; i < CountSize; i++ {
		server.Workers[0].ProcessMetric(&UDPMetric{
			MetricKey: MetricKey{
				Name: "x.y.z",
				Type: "counter",
			},
			Value:      1.0,
			Digest:     12345,
			SampleRate: 1.0,
			LocalOnly:  true,
		})
	}

	server.Flush(config.Interval, config.FlushLimit)
}

func TestSplitBytes(t *testing.T) {
	rand.Seed(time.Now().Unix())
	buf := make([]byte, 1000)

	for i := 0; i < 1000; i++ {
		// we construct a string of random length which is approximately 1/3rd A
		// and the other 2/3rds B
		buf = buf[:rand.Intn(cap(buf))]
		for i := range buf {
			if rand.Intn(3) == 0 {
				buf[i] = 'A'
			} else {
				buf[i] = 'B'
			}
		}
		checkBufferSplit(t, buf)
		buf = buf[:cap(buf)]
	}

	// also test pathological cases that the fuzz is unlikely to find
	checkBufferSplit(t, nil)
	checkBufferSplit(t, []byte{})
}

func checkBufferSplit(t *testing.T, buf []byte) {
	var testSplit [][]byte
	sb := NewSplitBytes(buf, 'A')
	for sb.Next() {
		testSplit = append(testSplit, sb.Chunk())
	}

	// now compare our split to the "real" implementation of split
	assert.EqualValues(t, bytes.Split(buf, []byte{'A'}), testSplit, "should have split %s correctly", buf)
}
<commit_msg>Ensure that test server responses terminate before tests<commit_after>package veneur

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stripe/veneur/tdigest"
)

const ε = .00002

const DefaultServerTimeout = 100 * time.Millisecond

var DebugMode bool

func TestMain(m *testing.M) {
	flag.Parse()
	DebugMode = flag.Lookup("test.v").Value.(flag.Getter).Get().(bool)
	os.Exit(m.Run())
}

// On the CI server, we can't be guaranteed that the port will be
// released immediately after the server is shut down. Instead, use
// a unique port for each test. As long as we don't have an insane number
// of integration tests, we should be fine.
var HttpAddrPort = 8127

// set up a boilerplate local config for later use
func localConfig() Config {
	port := HttpAddrPort
	HttpAddrPort++

	return Config{
		APIHostname: "http://localhost",
		Debug:       DebugMode,
		Hostname:    "localhost",

		// Use a shorter interval for tests
		Interval:            50 * time.Millisecond,
		Key:                 "",
		MetricMaxLength:     4096,
		Percentiles:         []float64{.5, .75, .99},
		ReadBufferSizeBytes: 2097152,
		UDPAddr:             "localhost:8126",
		HTTPAddr:            fmt.Sprintf("localhost:%d", port),
		ForwardAddr:         "http://localhost",
		NumWorkers:          96,

		// Use only one reader, so that we can run tests
		// on platforms which do not support SO_REUSEPORT
		NumReaders: 1,

		// Currently this points nowhere, which is intentional.
		// We don't need internal metrics for the tests, and they make testing
		// more complicated.
		StatsAddr:  "localhost:8125",
		Tags:       []string{},
		SentryDSN:  "",
		FlushLimit: 1024,
	}
}

// assertMetrics checks that all expected metrics are present
// and have the correct value
func assertMetrics(t *testing.T, metrics DDMetricsRequest, expectedMetrics map[string]float64) {
	// it doesn't count as accidentally quadratic if it's intentional
	for metricName, expectedValue := range expectedMetrics {
		assertMetric(t, metrics, metricName, expectedValue)
	}
}

func assertMetric(t *testing.T, metrics DDMetricsRequest, metricName string, value float64) {
	defer func() {
		if r := recover(); r != nil {
			assert.Fail(t, "error extracting metrics", r)
		}
	}()
	for _, metric := range metrics.Series {
		if metric.Name == metricName {
			assert.Equal(t, int(value+.5), int(metric.Value[0][1]+.5), fmt.Sprintf("Incorrect value for metric %s", metricName))
			return
		}
	}
	assert.Fail(t, "did not find expected metric", metricName)
}

// setupLocalServer creates a local server from the specified config
// and starts listening for requests. It returns the server for inspection.
func setupLocalServer(t *testing.T, config Config) Server {
	server, err := NewFromConfig(config)
	if err != nil {
		t.Fatal(err)
	}

	packetPool := &sync.Pool{
		New: func() interface{} {
			return make([]byte, config.MetricMaxLength)
		},
	}

	for i := 0; i < config.NumReaders; i++ {
		go func() {
			defer func() {
				if r := recover(); r != nil {
					assert.Fail(t, "reader panicked while reading from socket", err)
				}
			}()
			server.ReadSocket(packetPool, config.NumReaders != 1)
		}()
	}

	go server.HTTPServe()
	return server
}

// DDMetricsRequest represents the body of the POST request
// for sending metrics data to Datadog
// Eventually we'll want to define this symmetrically.
type DDMetricsRequest struct {
	Series []DDMetric
}

// TestLocalServerUnaggregatedMetrics tests the behavior of
// the veneur client when operating without a global veneur
// instance (ie, when sending data directly to the remote server)
func TestLocalServerUnaggregatedMetrics(t *testing.T) {
	// Since we are asserting inside a separate goroutine, we should
	// always assert that we get to the end of the goroutine.
	// In the future, if the key call (in this case, server.Flush)
	// is made asynchronous, we will know if we are ever encountering
	// a race condition in which our assertions are simply not executing.
	// (Failures after a test has finished executing will actually cause
	// the test suite to fail, as long as it is still running, but this is
	// a failsafe).
	var RemoteResponses int
	defer func() {
		assert.Equal(t, 1, RemoteResponses, "Server did not complete all responses before test terminated!")
	}()

	expectedMetrics := map[string]float64{
		"a.b.c.max": 100,
		"a.b.c.min": 1,

		// Count is normalized by second
		// so 5 values/50ms = 100 values/s
		"a.b.c.count": 100,

		// tdigest approximation causes this to be off by 1
		"a.b.c.50percentile": 6,
		"a.b.c.75percentile": 42,
		"a.b.c.99percentile": 98,
	}

	remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		zr, err := zlib.NewReader(r.Body)
		if err != nil {
			t.Fatal(err)
		}

		var ddmetrics DDMetricsRequest
		err = json.NewDecoder(zr).Decode(&ddmetrics)
		if err != nil {
			t.Fatal(err)
		}
		assert.Equal(t, 6, len(ddmetrics.Series), "incorrect number of elements in the flushed series on the remote server")
		assertMetrics(t, ddmetrics, expectedMetrics)
		w.WriteHeader(http.StatusAccepted)
		RemoteResponses++
	}))

	config := localConfig()
	config.APIHostname = remoteServer.URL

	server := setupLocalServer(t, config)
	defer server.Shutdown()

	metricValues := []float64{1.0, 2.0, 7.0, 8.0, 100.0}

	for _, value := range metricValues {
		server.Workers[0].ProcessMetric(&UDPMetric{
			MetricKey: MetricKey{
				Name: "a.b.c",
				Type: "histogram",
			},
			Value:      value,
			Digest:     12345,
			SampleRate: 1.0,
			LocalOnly:  true,
		})
	}
	server.Flush(config.Interval, config.FlushLimit)
}

func TestLocalServerMixedMetrics(t *testing.T) {

	// Same as in TestLocalServerUnaggregatedMetrics
	var RemoteResponses int
	defer func() {
		assert.Equal(t, 1, RemoteResponses, "Server did not complete all responses before test terminated!")
	}()

	// Unlike flushing to to the remote server, flushing forward
	// (ie, flushing to global veneur) is asynchronous.
	// This means it is quite likely that the function will return before
	// the global server has actually responded, unless we introduce a small timeout
	GlobalResponseChan := make(chan struct{})
	defer func() {
		select {
		case <-GlobalResponseChan:
			// all is safe
		case <-time.After(DefaultServerTimeout):
			assert.Fail(t, "Global server did not complete all responses before test terminated!")
		}
	}()

	// We can't be guaranteed that we will receive this exact stream
	// but we should expect that they represent the same underlying tdigest
	const ExpectedGobStream = "\r\xff\x87\x02\x01\x02\xff\x88\x00\x01\xff\x84\x00\x007\xff\x83\x03\x01\x01\bCentroid\x01\xff\x84\x00\x01\x03\x01\x04Mean\x01\b\x00\x01\x06Weight\x01\b\x00\x01\aSamples\x01\xff\x86\x00\x00\x00\x17\xff\x85\x02\x01\x01\t[]float64\x01\xff\x86\x00\x01\b\x00\x00/\xff\x88\x00\x05\x01\xfe\xf0?\x01\xfe\xf0?\x00\x01@\x01\xfe\xf0?\x00\x01\xfe\x1c@\x01\xfe\xf0?\x00\x01\xfe @\x01\xfe\xf0?\x00\x01\xfeY@\x01\xfe\xf0?\x00\x05\b\x00\xfeY@\x05\b\x00\xfe\xf0?\x05\b\x00\xfeY@"

	var HistogramValues = []float64{1.0, 2.0, 7.0, 8.0, 100.0}
	const ABCCountRaw = 5
	const ABCCountNormalized = 100

	const XYZCountSize = 40
	expectedMetrics := map[string]float64{
		// 40 events/50ms = 800 events/s
		"x.y.z":     800,
		"a.b.c.max": 100,
		"a.b.c.min": 1,

		// Count is normalized by second
		// so 5 values/50ms = 100 values/s
		"a.b.c.count": ABCCountNormalized,
	}

	// Set up a remote server (the API that we're sending the data to)
	// (e.g. Datadog)
	remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		zr, err := zlib.NewReader(r.Body)
		if err != nil {
			t.Fatal(err)
		}

		var ddmetrics DDMetricsRequest

		err = json.NewDecoder(zr).Decode(&ddmetrics)
		if err != nil {
			t.Fatal(err)
		}
		assert.Equal(t, len(expectedMetrics), len(ddmetrics.Series), "incorrect number of elements in the flushed series on the remote server")
		assertMetrics(t, ddmetrics, expectedMetrics)

		RemoteResponses++
		w.WriteHeader(http.StatusAccepted)
	}))

	// This represents the global veneur instance, which receives request from
	// the local veneur instances, aggregates the data, and sends it to the remote API
	// (e.g. Datadog)
	globalVeneur := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		zr, err := zlib.NewReader(r.Body)
		if err != nil {
			t.Fatal(err)
		}

		type requestItem struct {
			Name      string      `json:"name"`
			Tags      interface{} `json:"tags"`
			Tagstring string      `json:"tagstring"`
			Type      string      `json:"type"`
			Value     []byte      `json:"value"`
		}

		var metrics []requestItem

		err = json.NewDecoder(zr).Decode(&metrics)
		if err != nil {
			t.Fatal(err)
		}

		assert.Equal(t, 1, len(metrics), "incorrect number of elements in the flushed series")

		tdExpected := tdigest.NewMerging(100, false)
		err = tdExpected.GobDecode([]byte(ExpectedGobStream))
		assert.NoError(t, err, "Should not have encountered error in decoding expected gob stream")

		td := tdigest.NewMerging(100, false)
		err = td.GobDecode(metrics[0].Value)
		assert.NoError(t, err, "Should not have encountered error in decoding gob stream")

		assert.Equal(t, expectedMetrics["a.b.c.min"], td.Min(), "Minimum value is incorrect")
		assert.Equal(t, expectedMetrics["a.b.c.max"], td.Max(), "Maximum value is incorrect")

		// The remote server receives the raw count, *not* the normalized count
		assert.InEpsilon(t, ABCCountRaw, td.Count(), ε)
		assert.Equal(t, tdExpected, td, "Underlying tdigest structure is incorrect")

		GlobalResponseChan <- struct{}{}
		w.WriteHeader(http.StatusAccepted)
	}))

	config := localConfig()
	config.APIHostname = remoteServer.URL
	config.ForwardAddr = globalVeneur.URL
	config.NumWorkers = 1

	server := setupLocalServer(t, config)
	defer server.Shutdown()

	// Create non-local metrics that should be passed to the global veneur instance
	for _, value := range HistogramValues {
		server.Workers[0].ProcessMetric(&UDPMetric{
			MetricKey: MetricKey{
				Name: "a.b.c",
				Type: "histogram",
			},
			Value:      value,
			Digest:     12345,
			SampleRate: 1.0,
			LocalOnly:  false,
		})
	}

	// Create local-only metrics that should be passed directly to the remote API
	for i := 0; i < XYZCountSize; i++ {
		server.Workers[0].ProcessMetric(&UDPMetric{
			MetricKey: MetricKey{
				Name: "x.y.z",
				Type: "counter",
			},
			Value:      1.0,
			Digest:     12345,
			SampleRate: 1.0,
			LocalOnly:  true,
		})
	}

	server.Flush(config.Interval, config.FlushLimit)
}

func TestSplitBytes(t *testing.T) {
	rand.Seed(time.Now().Unix())
	buf := make([]byte, 1000)

	for i := 0; i < 1000; i++ {
		// we construct a string of random length which is approximately 1/3rd A
		// and the other 2/3rds B
		buf = buf[:rand.Intn(cap(buf))]
		for i := range buf {
			if rand.Intn(3) == 0 {
				buf[i] = 'A'
			} else {
				buf[i] = 'B'
			}
		}
		checkBufferSplit(t, buf)
		buf = buf[:cap(buf)]
	}

	// also test pathological cases that the fuzz is unlikely to find
	checkBufferSplit(t, nil)
	checkBufferSplit(t, []byte{})
}

func checkBufferSplit(t *testing.T, buf []byte) {
	var testSplit [][]byte
	sb := NewSplitBytes(buf, 'A')
	for sb.Next() {
		testSplit = append(testSplit, sb.Chunk())
	}

	// now compare our split to the "real" implementation of split
	assert.EqualValues(t, bytes.Split(buf, []byte{'A'}), testSplit, "should have split %s correctly", buf)
}
<|endoftext|>


<commit_before>package core

import (
	"google.golang.org/grpc"
	"log"
	"os"
	msg "qpm.io/common/messages"
	"google.golang.org/grpc/credentials"
	"fmt"
	"runtime"
)

const (
	Version       = "0.0.1"
	PackageFile   = "qpm.json"
	SignatureFile = "qpm.asc"
	Vendor        = "vendor"
	Address       = "pkg.qpm.io:7000"
	LicenseFile   = "LICENSE"
)

var UA = fmt.Sprintf("qpm/%v (%s; %s)", Version, runtime.GOOS, runtime.GOARCH)

type Context struct {
	Log    *log.Logger
	Client msg.QpmClient
}

func NewContext() *Context {
	log := log.New(os.Stderr, "QPM: ", log.LstdFlags)

	creds := credentials.NewClientTLSFromCert(nil, "")
	address := os.Getenv("SERVER")
	if address == "" {
		address = Address
	}
	conn, err := grpc.Dial(address, grpc.WithTransportCredentials(creds), grpc.WithUserAgent(UA))
	if err != nil {
		log.Fatalf("did not connect: %v", err)
	}

	return &Context{
		Log:    log,
		Client: msg.NewQpmClient(conn),
	}
}
<commit_msg>Bump the version to 0.9.0.<commit_after>package core

import (
	"google.golang.org/grpc"
	"log"
	"os"
	msg "qpm.io/common/messages"
	"google.golang.org/grpc/credentials"
	"fmt"
	"runtime"
)

const (
	Version       = "0.9.0"
	PackageFile   = "qpm.json"
	SignatureFile = "qpm.asc"
	Vendor        = "vendor"
	Address       = "pkg.qpm.io:7000"
	LicenseFile   = "LICENSE"
)

var UA = fmt.Sprintf("qpm/%v (%s; %s)", Version, runtime.GOOS, runtime.GOARCH)

type Context struct {
	Log    *log.Logger
	Client msg.QpmClient
}

func NewContext() *Context {
	log := log.New(os.Stderr, "QPM: ", log.LstdFlags)

	creds := credentials.NewClientTLSFromCert(nil, "")
	address := os.Getenv("SERVER")
	if address == "" {
		address = Address
	}
	conn, err := grpc.Dial(address, grpc.WithTransportCredentials(creds), grpc.WithUserAgent(UA))
	if err != nil {
		log.Fatalf("did not connect: %v", err)
	}

	return &Context{
		Log:    log,
		Client: msg.NewQpmClient(conn),
	}
}
<|endoftext|>


<commit_before>package serverstats

import (
	"fmt"
	"os"
	"time"

	"github.com/aleasoluciones/goaleasoluciones/scheduledtask"
	"github.com/cloudfoundry/gosigar"
)

const (
	MEMSTATS_PERIODE           = 10 * time.Second
	LOADAVG_AND_UPTIME_PERIODE = 10 * time.Second
	CPU_PERIODE                = 5 * time.Second
)

type Metric struct {
	Timestamp int64  `json:"timestamp"`
	Name      string `json:"name"`
	Value     string `json:"value"`
	Unit      string `json:"unit,omitempty"`
}

func loadavg(cs chan Metric, hostname string) {
	now := time.Now().Unix()
	concreteSigar := sigar.ConcreteSigar{}

	avg, err := concreteSigar.GetLoadAverage()
	if err != nil {
		fmt.Printf("Failed to get load average")
		return
	}

	cs <- createMetric(now, hostname, "loadavg.one", fmt.Sprintf("%.2f", avg.One), "")
	cs <- createMetric(now, hostname, "loadavg.five", fmt.Sprintf("%.2f", avg.Five), "")
	cs <- createMetric(now, hostname, "loadavg.Fifteen", fmt.Sprintf("%.2f", avg.Five), "")
}

func memStats(cs chan Metric, hostname string) {
	now := time.Now().Unix()
	mem := sigar.Mem{}
	swap := sigar.Swap{}

	mem.Get()
	swap.Get()

	cs <- createMetric(now, hostname, "mem.total", toMegabytesString(mem.Total), "M")
	cs <- createMetric(now, hostname, "mem.used", toMegabytesString(mem.Used), "M")
	cs <- createMetric(now, hostname, "mem.free", toMegabytesString(mem.Free), "M")
	cs <- createMetric(now, hostname, "mem.actualused", toMegabytesString(mem.ActualUsed), "M")
	cs <- createMetric(now, hostname, "mem.actualfree", toMegabytesString(mem.ActualFree), "M")
	cs <- createMetric(now, hostname, "swap.total", toMegabytesString(swap.Total), "M")
	cs <- createMetric(now, hostname, "swap.used", toMegabytesString(swap.Used), "M")
	cs <- createMetric(now, hostname, "swap.free", toMegabytesString(swap.Free), "M")
}

func toMegabytesString(bytes uint64) string {
	return fmt.Sprintf("%.2f", (float64(bytes) / 1024 / 1024))
}

func cpuStatsLoop(ch chan Metric, hostname string, periode time.Duration) {
	concreteSigar := sigar.ConcreteSigar{}
	cpuCh, _ := concreteSigar.CollectCpuStats(periode)
	for cpuStat := range cpuCh {
		now := time.Now().Unix()
		total := float64(cpuStat.Total())

		user := (float64(cpuStat.User) / total) * 100
		sys := (float64(cpuStat.Sys) / total) * 100
		idle := (float64(cpuStat.Idle) / total) * 100
		wait := (float64(cpuStat.Wait) / total) * 100
		stolen := (float64(cpuStat.Stolen) / total) * 100

		ch <- createMetric(now, hostname, "cpu.user", fmt.Sprintf("%.2f", user), "%")
		ch <- createMetric(now, hostname, "cpu.sys", fmt.Sprintf("%.2f", sys), "%")
		ch <- createMetric(now, hostname, "cpu.idle", fmt.Sprintf("%.2f", idle), "%")
		ch <- createMetric(now, hostname, "cpu.wait", fmt.Sprintf("%.2f", wait), "%")
		ch <- createMetric(now, hostname, "cpu.stolen", fmt.Sprintf("%.2f", stolen), "%")
	}
}

func createMetric(timestamp int64, hostname, name, value string, unit string) Metric {
	return Metric{
		Timestamp: timestamp,
		Name:      hostname + "." + name,
		Value:     value,
		Unit:      unit,
	}
}

type ServerStatsPeriodes struct {
	Mem     time.Duration
	LoadAvg time.Duration
	Cpu     time.Duration
}

var DefaultPeriodes = &ServerStatsPeriodes{
	Mem:     1 * time.Second,
	LoadAvg: 1 * time.Second,
	Cpu:     1 * time.Second,
}

type ServerStats struct {
	Metrics chan Metric
}

func NewServerStats(periodes *ServerStatsPeriodes) *ServerStats {
	serverStats := ServerStats{
		Metrics: make(chan Metric),
	}
	hostname, _ := os.Hostname()
	scheduledtask.NewScheduledTask(func() { memStats(serverStats.Metrics, hostname) }, periodes.Mem, 0)
	scheduledtask.NewScheduledTask(func() { loadavg(serverStats.Metrics, hostname) }, periodes.LoadAvg, 0)
	go cpuStatsLoop(serverStats.Metrics, hostname, periodes.Cpu)
	return &serverStats
}
<commit_msg>Added mem used (without buffers and cache) and swap used percent)<commit_after>package serverstats

import (
	"fmt"
	"os"
	"time"

	"github.com/aleasoluciones/goaleasoluciones/scheduledtask"
	"github.com/cloudfoundry/gosigar"
)

const (
	MEMSTATS_PERIODE           = 10 * time.Second
	LOADAVG_AND_UPTIME_PERIODE = 10 * time.Second
	CPU_PERIODE                = 5 * time.Second
)

type Metric struct {
	Timestamp int64  `json:"timestamp"`
	Name      string `json:"name"`
	Value     string `json:"value"`
	Unit      string `json:"unit,omitempty"`
}

func loadavg(cs chan Metric, hostname string) {
	now := time.Now().Unix()
	concreteSigar := sigar.ConcreteSigar{}

	avg, err := concreteSigar.GetLoadAverage()
	if err != nil {
		fmt.Printf("Failed to get load average")
		return
	}

	cs <- createMetric(now, hostname, "loadavg.one", fmt.Sprintf("%.2f", avg.One), "")
	cs <- createMetric(now, hostname, "loadavg.five", fmt.Sprintf("%.2f", avg.Five), "")
	cs <- createMetric(now, hostname, "loadavg.Fifteen", fmt.Sprintf("%.2f", avg.Five), "")
}

func memStats(cs chan Metric, hostname string) {
	now := time.Now().Unix()
	mem := sigar.Mem{}
	swap := sigar.Swap{}

	mem.Get()
	swap.Get()

	cs <- createMetric(now, hostname, "mem.total", toMegabytesString(mem.Total), "M")
	cs <- createMetric(now, hostname, "mem.used", toMegabytesString(mem.Used), "M")
	cs <- createMetric(now, hostname, "mem.free", toMegabytesString(mem.Free), "M")
	cs <- createMetric(now, hostname, "mem.actualused", toMegabytesString(mem.ActualUsed), "M")
	cs <- createMetric(now, hostname, "mem.actualfree", toMegabytesString(mem.ActualFree), "M")
	cs <- createMetric(now, hostname, "swap.total", toMegabytesString(swap.Total), "M")
	cs <- createMetric(now, hostname, "swap.used", toMegabytesString(swap.Used), "M")
	cs <- createMetric(now, hostname, "swap.free", toMegabytesString(swap.Free), "M")

	memActualUsedPercent := (float64(mem.ActualUsed) / float64(mem.Total)) * 100
	swapUsedPercent := (float64(swap.Used) / float64(swap.Total)) * 100
	cs <- createMetric(now, hostname, "mem.actualusedpercent", fmt.Sprintf("%.2f", memActualUsedPercent), "%")
	cs <- createMetric(now, hostname, "swap.usedpercent", fmt.Sprintf("%.2f", swapUsedPercent), "%")
}

func toMegabytesString(bytes uint64) string {
	return fmt.Sprintf("%.2f", (float64(bytes) / 1024 / 1024))
}

func cpuStatsLoop(ch chan Metric, hostname string, periode time.Duration) {
	concreteSigar := sigar.ConcreteSigar{}
	cpuCh, _ := concreteSigar.CollectCpuStats(periode)
	for cpuStat := range cpuCh {
		now := time.Now().Unix()
		total := float64(cpuStat.Total())

		user := (float64(cpuStat.User) / total) * 100
		sys := (float64(cpuStat.Sys) / total) * 100
		idle := (float64(cpuStat.Idle) / total) * 100
		wait := (float64(cpuStat.Wait) / total) * 100
		stolen := (float64(cpuStat.Stolen) / total) * 100

		ch <- createMetric(now, hostname, "cpu.user", fmt.Sprintf("%.2f", user), "%")
		ch <- createMetric(now, hostname, "cpu.sys", fmt.Sprintf("%.2f", sys), "%")
		ch <- createMetric(now, hostname, "cpu.idle", fmt.Sprintf("%.2f", idle), "%")
		ch <- createMetric(now, hostname, "cpu.wait", fmt.Sprintf("%.2f", wait), "%")
		ch <- createMetric(now, hostname, "cpu.stolen", fmt.Sprintf("%.2f", stolen), "%")
	}
}

func createMetric(timestamp int64, hostname, name, value string, unit string) Metric {
	return Metric{
		Timestamp: timestamp,
		Name:      hostname + "." + name,
		Value:     value,
		Unit:      unit,
	}
}

type ServerStatsPeriodes struct {
	Mem     time.Duration
	LoadAvg time.Duration
	Cpu     time.Duration
}

var DefaultPeriodes = &ServerStatsPeriodes{
	Mem:     1 * time.Second,
	LoadAvg: 1 * time.Second,
	Cpu:     1 * time.Second,
}

type ServerStats struct {
	Metrics chan Metric
}

func NewServerStats(periodes *ServerStatsPeriodes) *ServerStats {
	serverStats := ServerStats{
		Metrics: make(chan Metric),
	}
	hostname, _ := os.Hostname()
	scheduledtask.NewScheduledTask(func() { memStats(serverStats.Metrics, hostname) }, periodes.Mem, 0)
	scheduledtask.NewScheduledTask(func() { loadavg(serverStats.Metrics, hostname) }, periodes.LoadAvg, 0)
	go cpuStatsLoop(serverStats.Metrics, hostname, periodes.Cpu)
	return &serverStats
}
<|endoftext|>