<commit_before>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	Aws struct {
		Key    string
		Secret string
	}
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     RuntimeOptions
	}
	Mongo          string
	MongoKontrol   string
	MongoMinWrites int
	Mq             struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Neo4j struct {
		Read    string
		Write   string
		Port    int
		Enabled bool
	}
	GoLogLevel        string
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	Loggr             struct {
		Push   bool
		Url    string
		ApiKey string
	}
	Librato struct {
		Push     bool
		Email    string
		Token    string
		Interval int
	}
	Opsview struct {
		Push bool
		Host string
	}
	ElasticSearch struct {
		Host  string
		Port  int
		Queue string
	}
	NewKites struct {
		UseTLS   bool
		CertFile string
		KeyFile  string
	}
	NewKontrol struct {
		Port           int
		UseTLS         bool
		CertFile       string
		KeyFile        string
		PublicKeyFile  string
		PrivateKeyFile string
	}
	ProxyKite struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Etcd []struct {
		Host string
		Port int
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	FollowFeed struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
	Statsd struct {
		Use  bool
		Ip   string
		Port int
	}
	TopicModifier struct {
		CronSchedule string
	}
	Slack struct {
		Token   string
		Channel string
	}
	Graphite struct {
		Use  bool
		Host string
		Port int
	}
	LogLevel             map[string]string
	Redis                string
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
	SocialApi struct {
		ProxyUrl     string
		CustomDomain struct {
			Public string
			Local  string
		}
	}
	Vmwatcher struct {
		Port           string
		AwsKey         string
		AwsSecret      string
		KloudSecretKey string
		KloudAddr      string
	}
	Segment string
}

type RuntimeOptions struct {
	Kites struct {
		DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
		Stack                     struct {
			Force    bool `json:"force"`
			NewKites bool `json:"newKites"`
		} `json:"stack"`
		Kontrol struct {
			Username string `json:"username"`
		} `json:"kontrol"`
		Os struct {
			Version string `json:"version"`
		} `json:"os"`
		Terminal struct {
			Version string `json:"version"`
		} `json:"terminal"`
		Klient struct {
			Version string `json:"version"`
		} `json:"klient"`
		Kloud struct {
			Version string `json:"version"`
		} `json:"kloud"`
	} `json:"kites"`
	Algolia struct {
		AppId       string `json:"appId"`
		ApiKey      string `json:"apiKey"`
		IndexSuffix string `json:"indexSuffix"`
	} `json:"algolia"`
	LogToExternal   bool   `json:"logToExternal"`
	SuppressLogs    bool   `json:"suppressLogs"`
	LogToInternal   bool   `json:"logToInternal"`
	AuthExchange    string `json:"authExchange"`
	Environment     string `json:"environment"`
	Version         string `json:"version"`
	ResourceName    string `json:"resourceName"`
	UserSitesDomain string `json:"userSitesDomain"`
	LogResourceName string `json:"logResourceName"`
	SocialApiUri    string `json:"socialApiUri"`
	ApiUri          string `json:"apiUri"`
	MainUri         string `json:"mainUri"`
	SourceMapsUri   string `json:"sourceMapsUri"`
	Broker          struct {
		Uri string `json:"uri"`
	} `json:"broker"`
	AppsUri            string `json:"appsUri"`
	UploadsUri         string `json:"uploadsUri"`
	UploadsUriForGroup string `json:"uploadsUriForGroup"`
	FileFetchTimeout   int    `json:"fileFetchTimeout"`
	UserIdleMs         int    `json:"userIdleMs"`
	Embedly            struct {
		ApiKey string `json:"apiKey"`
	} `json:"embedly"`
	Github struct {
		ClientId string `json:"clientId"`
	} `json:"github"`
	Newkontrol struct {
		Url string `json:"url"`
	} `json:"newkontrol"`
	SessionCookie struct {
		MaxAge int  `json:"maxAge"`
		Secure bool `json:"secure"`
	} `json:"sessionCookie"`
	Troubleshoot struct {
		IdleTime    int    `json:"idleTime"`
		ExternalUrl string `json:"externalUrl"`
	} `json:"troubleshoot"`
	Stripe struct {
		Token string `json:"token"`
	} `json:"stripe"`
	ExternalProfiles struct {
		Google struct {
			Nicename string `json:"nicename"`
		} `json:"google"`
		Linkedin struct {
			Nicename string `json:"nicename"`
		} `json:"linkedin"`
		Twitter struct {
			Nicename string `json:"nicename"`
		} `json:"twitter"`
		Odesk struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"odesk"`
		Facebook struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"facebook"`
		Github struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"github"`
	} `json:"externalProfiles"`
	EntryPoint struct {
		Slug string `json:"slug"`
		Type string `json:"type"`
	} `json:"entryPoint"`
	Roles       []string      `json:"roles"`
	Permissions []interface{} `json:"permissions"`
	SiftScience string        `json:"siftScience"`
	Paypal      struct {
		FormUrl string `json:"formUrl"`
	} `json:"paypal"`
	Pubnub struct {
		SubscribeKey string `json:"subscribekey"`
		Enabled      bool   `json:"enabled"`
		SSL          bool   `json:"ssl"`
	} `json:"pubnub"`
	Collaboration struct {
		Timeout int `json:"timeout"`
	} `json:"collaboration"`
	PaymentBlockDuration float64 `json:"paymentBlockDuration"`
	Tokbox               struct {
		ApiKey string `json:"apiKey"`
	} `json:"tokbox"`
	DisabledFeatures struct {
		Moderation bool `json:"moderation"`
		Teams      bool `json:"teams"`
		BotChannel bool `json:"botchannel"`
	} `json:"disabledFeatures"`
	ContentRotatorUrl string `json:"contentRotatorUrl"`
	Google            struct {
		ApiKey string `json:"apiKey"`
	} `json:"google"`
	Recaptcha struct {
		Key     string `json:"key"`
		Enabled bool   `json:"enabled"`
	} `json:"recaptcha"`
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

// TODO: Fix this shit below where dir and profile is not even used ...
func MustConfigDir(dir, profile string) *Config {
	conf, err := readConfig(dir, profile)
	if err != nil {
		panic(err)
	}

	return conf
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<commit_msg>config: add integration to config.go<commit_after>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	Aws struct {
		Key    string
		Secret string
	}
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     RuntimeOptions
	}
	Mongo          string
	MongoKontrol   string
	MongoMinWrites int
	Mq             struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Neo4j struct {
		Read    string
		Write   string
		Port    int
		Enabled bool
	}
	GoLogLevel        string
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	Loggr             struct {
		Push   bool
		Url    string
		ApiKey string
	}
	Librato struct {
		Push     bool
		Email    string
		Token    string
		Interval int
	}
	Opsview struct {
		Push bool
		Host string
	}
	ElasticSearch struct {
		Host  string
		Port  int
		Queue string
	}
	NewKites struct {
		UseTLS   bool
		CertFile string
		KeyFile  string
	}
	NewKontrol struct {
		Port           int
		UseTLS         bool
		CertFile       string
		KeyFile        string
		PublicKeyFile  string
		PrivateKeyFile string
	}
	ProxyKite struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Etcd []struct {
		Host string
		Port int
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	FollowFeed struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
	Statsd struct {
		Use  bool
		Ip   string
		Port int
	}
	TopicModifier struct {
		CronSchedule string
	}
	Slack struct {
		Token   string
		Channel string
	}
	Graphite struct {
		Use  bool
		Host string
		Port int
	}
	LogLevel             map[string]string
	Redis                string
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
	SocialApi struct {
		ProxyUrl     string
		CustomDomain struct {
			Public string
			Local  string
		}
	}
	Vmwatcher struct {
		Port           string
		AwsKey         string
		AwsSecret      string
		KloudSecretKey string
		KloudAddr      string
	}
	Segment string
}

type RuntimeOptions struct {
	Kites struct {
		DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
		Stack                     struct {
			Force    bool `json:"force"`
			NewKites bool `json:"newKites"`
		} `json:"stack"`
		Kontrol struct {
			Username string `json:"username"`
		} `json:"kontrol"`
		Os struct {
			Version string `json:"version"`
		} `json:"os"`
		Terminal struct {
			Version string `json:"version"`
		} `json:"terminal"`
		Klient struct {
			Version string `json:"version"`
		} `json:"klient"`
		Kloud struct {
			Version string `json:"version"`
		} `json:"kloud"`
	} `json:"kites"`
	Algolia struct {
		AppId       string `json:"appId"`
		ApiKey      string `json:"apiKey"`
		IndexSuffix string `json:"indexSuffix"`
	} `json:"algolia"`
	LogToExternal   bool   `json:"logToExternal"`
	SuppressLogs    bool   `json:"suppressLogs"`
	LogToInternal   bool   `json:"logToInternal"`
	AuthExchange    string `json:"authExchange"`
	Environment     string `json:"environment"`
	Version         string `json:"version"`
	ResourceName    string `json:"resourceName"`
	UserSitesDomain string `json:"userSitesDomain"`
	LogResourceName string `json:"logResourceName"`
	SocialApiUri    string `json:"socialApiUri"`
	ApiUri          string `json:"apiUri"`
	MainUri         string `json:"mainUri"`
	SourceMapsUri   string `json:"sourceMapsUri"`
	Broker          struct {
		Uri string `json:"uri"`
	} `json:"broker"`
	AppsUri            string `json:"appsUri"`
	UploadsUri         string `json:"uploadsUri"`
	UploadsUriForGroup string `json:"uploadsUriForGroup"`
	FileFetchTimeout   int    `json:"fileFetchTimeout"`
	UserIdleMs         int    `json:"userIdleMs"`
	Embedly            struct {
		ApiKey string `json:"apiKey"`
	} `json:"embedly"`
	Github struct {
		ClientId string `json:"clientId"`
	} `json:"github"`
	Newkontrol struct {
		Url string `json:"url"`
	} `json:"newkontrol"`
	SessionCookie struct {
		MaxAge int  `json:"maxAge"`
		Secure bool `json:"secure"`
	} `json:"sessionCookie"`
	Troubleshoot struct {
		IdleTime    int    `json:"idleTime"`
		ExternalUrl string `json:"externalUrl"`
	} `json:"troubleshoot"`
	Stripe struct {
		Token string `json:"token"`
	} `json:"stripe"`
	ExternalProfiles struct {
		Google struct {
			Nicename string `json:"nicename"`
		} `json:"google"`
		Linkedin struct {
			Nicename string `json:"nicename"`
		} `json:"linkedin"`
		Twitter struct {
			Nicename string `json:"nicename"`
		} `json:"twitter"`
		Odesk struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"odesk"`
		Facebook struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"facebook"`
		Github struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"github"`
	} `json:"externalProfiles"`
	EntryPoint struct {
		Slug string `json:"slug"`
		Type string `json:"type"`
	} `json:"entryPoint"`
	Roles       []string      `json:"roles"`
	Permissions []interface{} `json:"permissions"`
	SiftScience string        `json:"siftScience"`
	Paypal      struct {
		FormUrl string `json:"formUrl"`
	} `json:"paypal"`
	Pubnub struct {
		SubscribeKey string `json:"subscribekey"`
		Enabled      bool   `json:"enabled"`
		SSL          bool   `json:"ssl"`
	} `json:"pubnub"`
	Collaboration struct {
		Timeout int `json:"timeout"`
	} `json:"collaboration"`
	PaymentBlockDuration float64 `json:"paymentBlockDuration"`
	Tokbox               struct {
		ApiKey string `json:"apiKey"`
	} `json:"tokbox"`
	DisabledFeatures struct {
		Moderation bool `json:"moderation"`
		Teams      bool `json:"teams"`
		BotChannel bool `json:"botchannel"`
	} `json:"disabledFeatures"`
	ContentRotatorUrl string `json:"contentRotatorUrl"`
	Integration       struct {
		Url string `json:"url"`
	} `json:"integration"`
	Google struct {
		ApiKey string `json:"apiKey"`
	} `json:"google"`
	Recaptcha struct {
		Key     string `json:"key"`
		Enabled bool   `json:"enabled"`
	} `json:"recaptcha"`
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

// TODO: Fix this shit below where dir and profile is not even used ...
func MustConfigDir(dir, profile string) *Config {
	conf, err := readConfig(dir, profile)
	if err != nil {
		panic(err)
	}

	return conf
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package server

import (
	"net/http"

	"github.com/issue9/middleware/v4/compress"
	"github.com/issue9/middleware/v4/errorhandler"
	"github.com/issue9/middleware/v4/recovery"
	"github.com/issue9/mux/v4"
	"github.com/issue9/source"
)

// Filter 针对 Context 的中间件
//
// Filter 和 github.com/issue9/middleware.Middleware 本质上没有任何区别，
// 都是作用于 http.Handler 上的中间件，只因参数不同，且两者不能交替出现，
// 派生出两套类型。
//
// 保证针对 middleware.Middleware 的 AddMiddlewares 方法，
// 可以最大限度地利用现有的通用中间件，而 AddFilter
// 方便用户编写针对 Context 的中间件，且 Context 提供了
// http.Handler 不存在的功能。
type Filter func(HandlerFunc) HandlerFunc

// FilterHandler 将过滤器应用于处理函数 next
func FilterHandler(next HandlerFunc, filter ...Filter) HandlerFunc {
	if l := len(filter); l > 0 {
		for i := l - 1; i >= 0; i-- {
			next = filter[i](next)
		}
	}
	return next
}

// AddFilters 添加过滤器
func (srv *Server) AddFilters(filter ...Filter) {
	srv.filters = append(srv.filters, filter...)
}

// Mux 返回 mux.Mux 实例
func (srv *Server) Mux() *mux.Mux { return srv.mux }

func (srv *Server) buildMiddlewares() error {
	srv.SetRecovery(func(w http.ResponseWriter, msg interface{}) {
		w.WriteHeader(http.StatusInternalServerError)
		data, err := source.TraceStack(5, msg)
		if err != nil {
			panic(err)
		}
		srv.Logs().Error(data)
	})

	if err := srv.SetCompressAlgorithm("deflate", compress.NewDeflate); err != nil {
		return err
	}

	if err := srv.SetCompressAlgorithm("gzip", compress.NewGzip); err != nil {
		return err
	}

	if err := srv.SetCompressAlgorithm("br", compress.NewBrotli); err != nil {
		return err
	}

	srv.AddMiddlewares(
		srv.recoveryMiddleware,       // 在最外层，防止协程 panic，崩了整个进程。
		srv.debugger.Middleware,      // 在外层添加调试地址，保证调试内容不会被其它 handler 干扰。
		srv.compress.Middleware,      // srv.errorhandlers 可能会输出大段内容。所以放在其之前。
		srv.errorHandlers.Middleware, // errorHandler 依赖 recovery，必须要在 recovery 之后。
	)

	return nil
}

func (srv *Server) recoveryMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		srv.recoverFunc.Middleware(next).ServeHTTP(w, r)
	})
}

// SetRecovery 设置在 panic 时的处理函数
func (srv *Server) SetRecovery(f recovery.RecoverFunc) { srv.recoverFunc = f }

// SetErrorHandle 设置指定状态码页面的处理函数
//
// 如果状态码已经存在处理函数，则修改，否则就添加。
func (srv *Server) SetErrorHandle(h errorhandler.HandleFunc, status ...int) {
	srv.errorHandlers.Set(h, status...)
}

// AddMiddlewares 添加中间件
func (srv *Server) AddMiddlewares(middleware ...mux.MiddlewareFunc) {
	for _, m := range middleware {
		srv.mux.AddMiddleware(true, m)
	}
}

// SetDebugger 设置调试地址
func (srv *Server) SetDebugger(pprof, vars string) (err error) {
	if pprof != "" {
		if pprof, err = srv.DefaultRouter().Path(pprof, nil); err != nil {
			return err
		}
	}

	if vars != "" {
		if vars, err = srv.DefaultRouter().Path(vars, nil); err != nil {
			return err
		}
	}

	srv.debugger.Pprof = pprof
	srv.debugger.Vars = vars

	return nil
}

// SetCompressAlgorithm 设置压缩的算法
//
// 默认情况下，支持 gzip、deflate 和 br 三种。
// 如果 w 为 nil，表示删除对该算法的支持。
func (srv *Server) SetCompressAlgorithm(name string, w compress.WriterFunc) error {
	return srv.compress.SetAlgorithm(name, w)
}

// AddCompressTypes 指定哪些内容可以进行压缩传输
//
// 默认情况下是所有内容都将进行压缩传输，
// * 表示所有；
// text/* 表示以 text/ 开头的类型；
// text/plain 表示具体类型 text/plain；
func (srv *Server) AddCompressTypes(types ...string) {
	srv.compress.AddType(types...)
}

// DeleteCompressTypes 删除指定内容的压缩支持
//
// 仅用于删除通过 AddType 添加的内容。
//
// NOTE: 如果指定 * 之后的所有内容都将不支持。
func (srv *Server) DeleteCompressTypes(types ...string) {
	srv.compress.DeleteType(types...)
}
<commit_msg>refactor(server): 删除 Server.AddMiddlewares<commit_after>// SPDX-License-Identifier: MIT

package server

import (
	"net/http"

	"github.com/issue9/middleware/v4/compress"
	"github.com/issue9/middleware/v4/errorhandler"
	"github.com/issue9/middleware/v4/recovery"
	"github.com/issue9/mux/v4"
	"github.com/issue9/source"
)

// Filter 针对 Context 的中间件
//
// Filter 和 github.com/issue9/mux.MiddlewareFunc 本质上没有任何区别，
// mux.MiddlewareFunc 更加的通用，可以复用市面上的大部分中间件，
// Filter 则更加灵活一些，可以针对模块或是某一个路由。
//
// 如果想要使用 mux.MiddlewareFunc，可以调用 Server.Mux().AddMiddleware 方法。
type Filter func(HandlerFunc) HandlerFunc

// FilterHandler 将过滤器应用于处理函数 next
func FilterHandler(next HandlerFunc, filter ...Filter) HandlerFunc {
	if l := len(filter); l > 0 {
		for i := l - 1; i >= 0; i-- {
			next = filter[i](next)
		}
	}
	return next
}

// AddFilters 添加过滤器
func (srv *Server) AddFilters(filter ...Filter) {
	srv.filters = append(srv.filters, filter...)
}

// Mux 返回 mux.Mux 实例
func (srv *Server) Mux() *mux.Mux { return srv.mux }

func (srv *Server) buildMiddlewares() error {
	srv.SetRecovery(func(w http.ResponseWriter, msg interface{}) {
		w.WriteHeader(http.StatusInternalServerError)
		data, err := source.TraceStack(5, msg)
		if err != nil {
			panic(err)
		}
		srv.Logs().Error(data)
	})

	if err := srv.SetCompressAlgorithm("deflate", compress.NewDeflate); err != nil {
		return err
	}

	if err := srv.SetCompressAlgorithm("gzip", compress.NewGzip); err != nil {
		return err
	}

	if err := srv.SetCompressAlgorithm("br", compress.NewBrotli); err != nil {
		return err
	}

	srv.Mux().AddMiddleware(true, srv.recoveryMiddleware)       // 在最外层，防止协程 panic，崩了整个进程。
	srv.Mux().AddMiddleware(true, srv.debugger.Middleware)      // 在外层添加调试地址，保证调试内容不会被其它 handler 干扰。
	srv.Mux().AddMiddleware(true, srv.compress.Middleware)      // srv.errorhandlers 可能会输出大段内容。所以放在其之前。
	srv.Mux().AddMiddleware(true, srv.errorHandlers.Middleware) // errorHandler 依赖 recovery，必须要在 recovery 之后。

	return nil
}

func (srv *Server) recoveryMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		srv.recoverFunc.Middleware(next).ServeHTTP(w, r)
	})
}

// SetRecovery 设置在 panic 时的处理函数
func (srv *Server) SetRecovery(f recovery.RecoverFunc) { srv.recoverFunc = f }

// SetErrorHandle 设置指定状态码页面的处理函数
//
// 如果状态码已经存在处理函数，则修改，否则就添加。
func (srv *Server) SetErrorHandle(h errorhandler.HandleFunc, status ...int) {
	srv.errorHandlers.Set(h, status...)
}

// SetDebugger 设置调试地址
func (srv *Server) SetDebugger(pprof, vars string) (err error) {
	if pprof != "" {
		if pprof, err = srv.DefaultRouter().Path(pprof, nil); err != nil {
			return err
		}
	}

	if vars != "" {
		if vars, err = srv.DefaultRouter().Path(vars, nil); err != nil {
			return err
		}
	}

	srv.debugger.Pprof = pprof
	srv.debugger.Vars = vars

	return nil
}

// SetCompressAlgorithm 设置压缩的算法
//
// 默认情况下，支持 gzip、deflate 和 br 三种。
// 如果 w 为 nil，表示删除对该算法的支持。
func (srv *Server) SetCompressAlgorithm(name string, w compress.WriterFunc) error {
	return srv.compress.SetAlgorithm(name, w)
}

// AddCompressTypes 指定哪些内容可以进行压缩传输
//
// 默认情况下是所有内容都将进行压缩传输，
// * 表示所有；
// text/* 表示以 text/ 开头的类型；
// text/plain 表示具体类型 text/plain；
func (srv *Server) AddCompressTypes(types ...string) {
	srv.compress.AddType(types...)
}

// DeleteCompressTypes 删除指定内容的压缩支持
//
// 仅用于删除通过 AddType 添加的内容。
//
// NOTE: 如果指定 * 之后的所有内容都将不支持。
func (srv *Server) DeleteCompressTypes(types ...string) {
	srv.compress.DeleteType(types...)
}
<|endoftext|>


<commit_before>package proxystorage

import (
	"context"
	"fmt"
	"reflect"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/sirupsen/logrus"

	proxyconfig "github.com/jacksontj/promxy/config"
	"github.com/jacksontj/promxy/promclient"
	"github.com/jacksontj/promxy/proxyquerier"
	"github.com/jacksontj/promxy/servergroup"
)

type proxyStorageState struct {
	sgs            []*servergroup.ServerGroup
	client         promclient.API
	cfg            *proxyconfig.PromxyConfig
	appender       storage.Appender
	appenderCloser func() error
}

func (p *proxyStorageState) Ready() {
	for _, sg := range p.sgs {
		<-sg.Ready
	}
}

func (p *proxyStorageState) Cancel(n *proxyStorageState) {
	if p.sgs != nil {
		for _, sg := range p.sgs {
			sg.Cancel()
		}
	}
	// We call close if the new one is nil, or if the appanders don't match
	if n == nil || p.appender != n.appender {
		if p.appenderCloser != nil {
			p.appenderCloser()
		}
	}
}

func NewProxyStorage() (*ProxyStorage, error) {
	return &ProxyStorage{}, nil
}

// TODO: rename?
type ProxyStorage struct {
	state atomic.Value
}

func (p *ProxyStorage) GetState() *proxyStorageState {
	tmp := p.state.Load()
	if sg, ok := tmp.(*proxyStorageState); ok {
		return sg
	} else {
		return &proxyStorageState{}
	}
}

func (p *ProxyStorage) ApplyConfig(c *proxyconfig.Config) error {
	oldState := p.GetState() // Fetch the old state

	failed := false

	apis := make([]promclient.API, len(c.ServerGroups))
	newState := &proxyStorageState{
		sgs: make([]*servergroup.ServerGroup, len(c.ServerGroups)),
		cfg: &c.PromxyConfig,
	}
	for i, sgCfg := range c.ServerGroups {
		tmp := servergroup.New()
		if err := tmp.ApplyConfig(sgCfg); err != nil {
			failed = true
			logrus.Errorf("Error applying config to server group: %s", err)
		}
		newState.sgs[i] = tmp
		apis[i] = tmp
	}
	newState.client = promclient.NewMultiAPI(apis, model.TimeFromUnix(0), nil)

	if failed {
		newState.Cancel(nil)
		return fmt.Errorf("Error Applying Config to one or more server group(s)")
	}

	// Check for remote_write (for appender)
	if c.PromConfig.RemoteWriteConfigs != nil {
		switch oldAppender := oldState.appender.(type) {
		// If the old one was a remote storage, we just need to apply config
		case *remote.Storage:
			if err := oldAppender.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			newState.appender = oldState.appender
			newState.appenderCloser = oldState.appenderCloser
		// if it was an appenderstub we just need to replace
		default:
			remote := remote.NewStorage(nil, func() (int64, error) { return 0, nil }, 1*time.Second)
			if err := remote.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			newState.appender = remote
			newState.appenderCloser = remote.Close
		}
	} else {
		newState.appender = &appenderStub{}
	}

	newState.Ready()        // Wait for the newstate to be ready
	p.state.Store(newState) // Store the new state
	if oldState != nil {
		oldState.Cancel(newState) // Cancel the old one
	}

	return nil
}

func (p *ProxyStorage) Querier(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
	state := p.GetState()
	return &proxyquerier.ProxyQuerier{
		ctx,
		timestamp.Time(mint),
		timestamp.Time(maxt),
		state.client,

		state.cfg,
	}, nil
}

func (p *ProxyStorage) StartTime() (int64, error) {
	return 0, nil
}

func (p *ProxyStorage) Appender() (storage.Appender, error) {
	state := p.GetState()
	return state.appender, nil
}

// TODO: actually close things?
func (p *ProxyStorage) Close() error { return nil }

// This replaces promql Nodes with more efficient-to-fetch ones. This works by taking lower-layer
// chunks of the query, farming them out to prometheus hosts, then stitching the results back together.
// An example would be a sum, we can sum multiple sums and come up with the same result -- so we do.
// There are a few ground rules for this:
//      - Children cannot be AggregateExpr: aggregates have their own combining logic, so its not safe to send a subquery with additional aggregations
//      - offsets within the subtree must match: if they don't then we'll get mismatched data, so we wait until we are far enough down the tree that they converge
//      - Don't reduce accuracy/granularity: the intention of this is to get the correct data faster, meaning correctness overrules speed.
func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *promql.EvalStmt, node promql.Node) (promql.Node, error) {

	isAgg := func(node promql.Node) bool {
		_, ok := node.(*promql.AggregateExpr)
		return ok
	}

	// If there is a child that is an aggregator we cannot do anything (as they have their own
	// rules around combining). We'll skip this node and let a lower layer take this on
	aggFinder := &BooleanFinder{Func: isAgg}
	offsetFinder := &OffsetFinder{}

	visitor := &MultiVisitor{[]promql.Visitor{aggFinder, offsetFinder}}

	if _, err := promql.Walk(ctx, visitor, s, node, nil, nil); err != nil {
		return nil, err
	}

	if aggFinder.Found > 0 {
		// If there was a single agg and that was us, then we're okay
		if !(isAgg(node) && aggFinder.Found == 1) {
			return nil, nil
		}
	}

	// If the tree below us is not all the same offset, then we can't do anything below -- we'll need
	// to wait until further in execution where they all match
	var offset time.Duration

	// If we couldn't find an offset, then something is wrong-- lets skip
	// Also if there was an error, skip
	if !offsetFinder.Found || offsetFinder.Error != nil {
		return nil, nil
	}
	offset = offsetFinder.Offset

	// Function to recursivelt remove offset. This is needed as we're using
	// the node API to String() the query to downstreams. Promql's iterators require
	// that the time be the absolute time, whereas the API returns them based on the
	// range you ask for (with the offset being implicit)
	// TODO: rename
	removeOffset := func() error {
		_, err := promql.Walk(ctx, &OffsetRemover{}, s, node, nil, nil)
		return err
	}

	state := p.GetState()
	switch n := node.(type) {
	// Some AggregateExprs can be composed (meaning they are "reentrant". If the aggregation op
	// is reentrant/composable then we'll do so, otherwise we let it fall through to normal query mechanisms
	case *promql.AggregateExpr:
		logrus.Debugf("AggregateExpr %v", n)

		var result model.Value
		var err error

		// Not all Aggregation functions are composable, so we'll do what we can
		switch n.Op.String() {
		// All "reentrant" cases (meaning they can be done repeatedly and the outcome doesn't change)
		case "sum", "min", "max", "topk", "bottomk":
			removeOffset()

			if s.Interval > 0 {
				result, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}

		// Convert avg into sum() / count()
		case "avg":
			// Replace with sum() / count()
			return &promql.BinaryExpr{
				Op: 24, // Divide  TODO
				LHS: &promql.AggregateExpr{
					Op:       41, // sum() TODO
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},

				RHS: &promql.AggregateExpr{
					Op:       40, // count() TODO
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},
				VectorMatching: &promql.VectorMatching{Card: promql.CardOneToOne},
			}, nil

		// For count we simply need to change this to a sum over the data we get back
		case "count":
			removeOffset()

			if s.Interval > 0 {
				result, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}
			// TODO: have a reverse method in promql/lex.go
			n.Op = 41 // SUM

		case "stddev": // TODO: something?
		case "stdvar": // TODO: something?

			// Do nothing, we want to allow the VectorSelector to fall through to do a query_range.
			// Unless we find another way to decompose the query, this is all we can do
		case "count_values":
			// DO NOTHING

		case "quantile": // TODO: something?

		}

		if result != nil {
			iterators := promclient.IteratorsForValue(result)

			series := make([]storage.Series, len(iterators))
			for i, iterator := range iterators {
				series[i] = &proxyquerier.Series{iterator}
			}

			ret := &promql.VectorSelector{Offset: offset}
			ret.SetSeries(series)
			n.Expr = ret
			return n, nil
		}

	// Call is for things such as rate() etc. This can be sent directly to the
	// prometheus node to answer
	case *promql.Call:
		logrus.Debugf("call %v %v", n, n.Type())
		removeOffset()

		var result model.Value
		var err error
		if s.Interval > 0 {
			result, err = state.client.QueryRange(ctx, n.String(), v1.Range{
				Start: s.Start.Add(-offset - promql.LookbackDelta),
				End:   s.End.Add(-offset),
				Step:  s.Interval,
			})
		} else {
			result, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
		}

		if err != nil {
			return nil, errors.Cause(err)
		}
		iterators := promclient.IteratorsForValue(result)
		series := make([]storage.Series, len(iterators))
		for i, iterator := range iterators {
			series[i] = &proxyquerier.Series{iterator}
		}

		ret := &promql.VectorSelector{Offset: offset}
		ret.SetSeries(series)
		return ret, nil

		// If we are simply fetching a Vector then we can fetch the data using the same step that
		// the query came in as (reducing the amount of data we need to fetch)
	// If we are simply fetching data, we skip here to let it fall through to the normal
	// storage API
	case *promql.VectorSelector:
		// Do Nothing
		return nil, nil

	// If we hit this someone is asking for a matrix directly, if so then we don't
	// have anyway to ask for less-- since this is exactly what they are asking for
	case *promql.MatrixSelector:
		// DO NOTHING

	default:
		logrus.Debugf("default %v %s", n, reflect.TypeOf(n))

	}
	return nil, nil

}
<commit_msg>Fix comment indentation<commit_after>package proxystorage

import (
	"context"
	"fmt"
	"reflect"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/sirupsen/logrus"

	proxyconfig "github.com/jacksontj/promxy/config"
	"github.com/jacksontj/promxy/promclient"
	"github.com/jacksontj/promxy/proxyquerier"
	"github.com/jacksontj/promxy/servergroup"
)

type proxyStorageState struct {
	sgs            []*servergroup.ServerGroup
	client         promclient.API
	cfg            *proxyconfig.PromxyConfig
	appender       storage.Appender
	appenderCloser func() error
}

func (p *proxyStorageState) Ready() {
	for _, sg := range p.sgs {
		<-sg.Ready
	}
}

func (p *proxyStorageState) Cancel(n *proxyStorageState) {
	if p.sgs != nil {
		for _, sg := range p.sgs {
			sg.Cancel()
		}
	}
	// We call close if the new one is nil, or if the appanders don't match
	if n == nil || p.appender != n.appender {
		if p.appenderCloser != nil {
			p.appenderCloser()
		}
	}
}

func NewProxyStorage() (*ProxyStorage, error) {
	return &ProxyStorage{}, nil
}

// TODO: rename?
type ProxyStorage struct {
	state atomic.Value
}

func (p *ProxyStorage) GetState() *proxyStorageState {
	tmp := p.state.Load()
	if sg, ok := tmp.(*proxyStorageState); ok {
		return sg
	} else {
		return &proxyStorageState{}
	}
}

func (p *ProxyStorage) ApplyConfig(c *proxyconfig.Config) error {
	oldState := p.GetState() // Fetch the old state

	failed := false

	apis := make([]promclient.API, len(c.ServerGroups))
	newState := &proxyStorageState{
		sgs: make([]*servergroup.ServerGroup, len(c.ServerGroups)),
		cfg: &c.PromxyConfig,
	}
	for i, sgCfg := range c.ServerGroups {
		tmp := servergroup.New()
		if err := tmp.ApplyConfig(sgCfg); err != nil {
			failed = true
			logrus.Errorf("Error applying config to server group: %s", err)
		}
		newState.sgs[i] = tmp
		apis[i] = tmp
	}
	newState.client = promclient.NewMultiAPI(apis, model.TimeFromUnix(0), nil)

	if failed {
		newState.Cancel(nil)
		return fmt.Errorf("Error Applying Config to one or more server group(s)")
	}

	// Check for remote_write (for appender)
	if c.PromConfig.RemoteWriteConfigs != nil {
		switch oldAppender := oldState.appender.(type) {
		// If the old one was a remote storage, we just need to apply config
		case *remote.Storage:
			if err := oldAppender.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			newState.appender = oldState.appender
			newState.appenderCloser = oldState.appenderCloser
		// if it was an appenderstub we just need to replace
		default:
			remote := remote.NewStorage(nil, func() (int64, error) { return 0, nil }, 1*time.Second)
			if err := remote.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			newState.appender = remote
			newState.appenderCloser = remote.Close
		}
	} else {
		newState.appender = &appenderStub{}
	}

	newState.Ready()        // Wait for the newstate to be ready
	p.state.Store(newState) // Store the new state
	if oldState != nil {
		oldState.Cancel(newState) // Cancel the old one
	}

	return nil
}

func (p *ProxyStorage) Querier(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
	state := p.GetState()
	return &proxyquerier.ProxyQuerier{
		ctx,
		timestamp.Time(mint),
		timestamp.Time(maxt),
		state.client,

		state.cfg,
	}, nil
}

func (p *ProxyStorage) StartTime() (int64, error) {
	return 0, nil
}

func (p *ProxyStorage) Appender() (storage.Appender, error) {
	state := p.GetState()
	return state.appender, nil
}

// TODO: actually close things?
func (p *ProxyStorage) Close() error { return nil }

// This replaces promql Nodes with more efficient-to-fetch ones. This works by taking lower-layer
// chunks of the query, farming them out to prometheus hosts, then stitching the results back together.
// An example would be a sum, we can sum multiple sums and come up with the same result -- so we do.
// There are a few ground rules for this:
//      - Children cannot be AggregateExpr: aggregates have their own combining logic, so its not safe to send a subquery with additional aggregations
//      - offsets within the subtree must match: if they don't then we'll get mismatched data, so we wait until we are far enough down the tree that they converge
//      - Don't reduce accuracy/granularity: the intention of this is to get the correct data faster, meaning correctness overrules speed.
func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *promql.EvalStmt, node promql.Node) (promql.Node, error) {

	isAgg := func(node promql.Node) bool {
		_, ok := node.(*promql.AggregateExpr)
		return ok
	}

	// If there is a child that is an aggregator we cannot do anything (as they have their own
	// rules around combining). We'll skip this node and let a lower layer take this on
	aggFinder := &BooleanFinder{Func: isAgg}
	offsetFinder := &OffsetFinder{}

	visitor := &MultiVisitor{[]promql.Visitor{aggFinder, offsetFinder}}

	if _, err := promql.Walk(ctx, visitor, s, node, nil, nil); err != nil {
		return nil, err
	}

	if aggFinder.Found > 0 {
		// If there was a single agg and that was us, then we're okay
		if !(isAgg(node) && aggFinder.Found == 1) {
			return nil, nil
		}
	}

	// If the tree below us is not all the same offset, then we can't do anything below -- we'll need
	// to wait until further in execution where they all match
	var offset time.Duration

	// If we couldn't find an offset, then something is wrong-- lets skip
	// Also if there was an error, skip
	if !offsetFinder.Found || offsetFinder.Error != nil {
		return nil, nil
	}
	offset = offsetFinder.Offset

	// Function to recursivelt remove offset. This is needed as we're using
	// the node API to String() the query to downstreams. Promql's iterators require
	// that the time be the absolute time, whereas the API returns them based on the
	// range you ask for (with the offset being implicit)
	// TODO: rename
	removeOffset := func() error {
		_, err := promql.Walk(ctx, &OffsetRemover{}, s, node, nil, nil)
		return err
	}

	state := p.GetState()
	switch n := node.(type) {
	// Some AggregateExprs can be composed (meaning they are "reentrant". If the aggregation op
	// is reentrant/composable then we'll do so, otherwise we let it fall through to normal query mechanisms
	case *promql.AggregateExpr:
		logrus.Debugf("AggregateExpr %v", n)

		var result model.Value
		var err error

		// Not all Aggregation functions are composable, so we'll do what we can
		switch n.Op.String() {
		// All "reentrant" cases (meaning they can be done repeatedly and the outcome doesn't change)
		case "sum", "min", "max", "topk", "bottomk":
			removeOffset()

			if s.Interval > 0 {
				result, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}

		// Convert avg into sum() / count()
		case "avg":
			// Replace with sum() / count()
			return &promql.BinaryExpr{
				Op: 24, // Divide  TODO
				LHS: &promql.AggregateExpr{
					Op:       41, // sum() TODO
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},

				RHS: &promql.AggregateExpr{
					Op:       40, // count() TODO
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},
				VectorMatching: &promql.VectorMatching{Card: promql.CardOneToOne},
			}, nil

		// For count we simply need to change this to a sum over the data we get back
		case "count":
			removeOffset()

			if s.Interval > 0 {
				result, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}
			// TODO: have a reverse method in promql/lex.go
			n.Op = 41 // SUM

		case "stddev": // TODO: something?
		case "stdvar": // TODO: something?

			// Do nothing, we want to allow the VectorSelector to fall through to do a query_range.
			// Unless we find another way to decompose the query, this is all we can do
		case "count_values":
			// DO NOTHING

		case "quantile": // TODO: something?

		}

		if result != nil {
			iterators := promclient.IteratorsForValue(result)

			series := make([]storage.Series, len(iterators))
			for i, iterator := range iterators {
				series[i] = &proxyquerier.Series{iterator}
			}

			ret := &promql.VectorSelector{Offset: offset}
			ret.SetSeries(series)
			n.Expr = ret
			return n, nil
		}

	// Call is for things such as rate() etc. This can be sent directly to the
	// prometheus node to answer
	case *promql.Call:
		logrus.Debugf("call %v %v", n, n.Type())
		removeOffset()

		var result model.Value
		var err error
		if s.Interval > 0 {
			result, err = state.client.QueryRange(ctx, n.String(), v1.Range{
				Start: s.Start.Add(-offset - promql.LookbackDelta),
				End:   s.End.Add(-offset),
				Step:  s.Interval,
			})
		} else {
			result, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
		}

		if err != nil {
			return nil, errors.Cause(err)
		}
		iterators := promclient.IteratorsForValue(result)
		series := make([]storage.Series, len(iterators))
		for i, iterator := range iterators {
			series[i] = &proxyquerier.Series{iterator}
		}

		ret := &promql.VectorSelector{Offset: offset}
		ret.SetSeries(series)
		return ret, nil

	// If we are simply fetching a Vector then we can fetch the data using the same step that
	// the query came in as (reducing the amount of data we need to fetch)
	// If we are simply fetching data, we skip here to let it fall through to the normal
	// storage API
	case *promql.VectorSelector:
		// Do Nothing
		return nil, nil

	// If we hit this someone is asking for a matrix directly, if so then we don't
	// have anyway to ask for less-- since this is exactly what they are asking for
	case *promql.MatrixSelector:
		// DO NOTHING

	default:
		logrus.Debugf("default %v %s", n, reflect.TypeOf(n))

	}
	return nil, nil

}
<|endoftext|>


<commit_before>package httpclient

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/gobs/pretty"
	"github.com/gobs/simplejson"
)

var (
	DefaultClient = &http.Client{} // we use our own default client, so we can change the TLS configuration

	NoRedirect       = errors.New("No redirect")
	TooManyRedirects = errors.New("stopped after 10 redirects")
)

//
// Allow connections via HTTPS even if something is wrong with the certificate
// (self-signed or expired)
//
func AllowInsecure(insecure bool) {
	if insecure {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}

		DefaultClient.Transport = tr
	} else {
		DefaultClient.Transport = nil
	}
}

//
// Set connection timeout
//
func SetTimeout(t time.Duration) {
	DefaultClient.Timeout = t
}

//
// HTTP error
//
type HttpError struct {
	Code    int
	Message string
}

func (e HttpError) Error() string {
	return e.Message
}

func (e HttpError) String() string {
	return fmt.Sprintf("ERROR: %v %v", e.Code, e.Message)
}

//
// CloseResponse makes sure we close the response body
//
func CloseResponse(r *http.Response) {
	if r != nil && r.Body != nil {
		io.Copy(ioutil.Discard, r.Body)
		r.Body.Close()
	}
}

//
// A wrapper for http.Response
//
type HttpResponse struct {
	http.Response
}

func (r *HttpResponse) ContentType() string {
	content_type := r.Header.Get("Content-Type")
	if len(content_type) == 0 {
		return content_type
	}

	return strings.TrimSpace(strings.Split(content_type, ";")[0])
}

//
// Close makes sure that all data from the body is read
// before closing the reader.
//
// If that is not the desider behaviour, just call HttpResponse.Body.Close()
//
func (r *HttpResponse) Close() {
	if r != nil {
		CloseResponse(&r.Response)
	}
}

//
// ResponseError checks the StatusCode and return an error if needed.
// The error is of type HttpError
//
func (r *HttpResponse) ResponseError() error {
	class := r.StatusCode / 100
	if class != 2 && class != 3 {
		return HttpError{Code: r.StatusCode, Message: "HTTP " + r.Status}
	}

	return nil
}

//
// Check if the input value is a "primitive" that can be safely stringified
//
func canStringify(v reflect.Value) bool {
	switch v.Kind() {
	case reflect.Bool,
		reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
		reflect.Float32, reflect.Float64,
		reflect.String:
		return true
	default:
		return false
	}
}

//
// Given a base URL and a bag of parameteters returns the URL with the encoded parameters
//
func URLWithPathParams(base string, path string, params map[string]interface{}) (u *url.URL) {

	u, err := url.Parse(base)
	if err != nil {
		log.Fatal(err)
	}

	if len(path) > 0 {
		u, err = u.Parse(path)
		if err != nil {
			log.Fatal(err)
		}
	}

	q := u.Query()

	for k, v := range params {
		val := reflect.ValueOf(v)

		switch val.Kind() {
		case reflect.Slice:
			if val.IsNil() { // TODO: add an option to ignore empty values
				q.Set(k, "")
				continue
			}
			fallthrough

		case reflect.Array:
			for i := 0; i < val.Len(); i++ {
				av := val.Index(i)

				if canStringify(av) {
					q.Add(k, fmt.Sprintf("%v", av))
				}
			}

		default:
			if canStringify(val) {
				q.Set(k, fmt.Sprintf("%v", v))
			} else {
				log.Fatal("Invalid type ", val)
			}
		}
	}

	u.RawQuery = q.Encode()
	return u
}

func URLWithParams(base string, params map[string]interface{}) (u *url.URL) {
	return URLWithPathParams(base, "", params)
}

//
// http.Get with params
//
func Get(urlStr string, params map[string]interface{}) (*HttpResponse, error) {
	resp, err := DefaultClient.Get(URLWithParams(urlStr, params).String())
	if err == nil {
		return &HttpResponse{*resp}, nil
	} else {
		CloseResponse(resp)
		return nil, err
	}
}

//
// http.Post with params
//
func Post(urlStr string, params map[string]interface{}) (*HttpResponse, error) {
	resp, err := DefaultClient.PostForm(urlStr, URLWithParams(urlStr, params).Query())
	if err == nil {
		return &HttpResponse{*resp}, nil
	} else {
		CloseResponse(resp)
		return nil, err
	}
}

//
//  Read the body
//
func (resp *HttpResponse) Content() []byte {
	body, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	return body
}

//
//  Try to parse the response body as JSON
//
func (resp *HttpResponse) Json() (json *simplejson.Json) {
	json, _ = simplejson.LoadBytes(resp.Content())
	return
}

////////////////////////////////////////////////////////////////////////

//
// http.Client with some defaults and stuff
//
type HttpClient struct {
	// the http.Client
	client *http.Client

	// the base URL for this client
	BaseURL *url.URL

	// the client UserAgent string
	UserAgent string

	// Common headers to be passed on each request
	Headers map[string]string

	// if Verbose, log request and response info
	Verbose bool
}

//
// Create a new HttpClient
//
func NewHttpClient(base string) (httpClient *HttpClient) {
	httpClient = new(HttpClient)
	httpClient.client = &http.Client{CheckRedirect: httpClient.checkRedirect}
	httpClient.Headers = make(map[string]string)

	if u, err := url.Parse(base); err != nil {
		log.Fatal(err)
	} else {
		httpClient.BaseURL = u
	}

	return
}

//
// Set Transport
//
func (self *HttpClient) SetTransport(tr http.RoundTripper) {
	self.client.Transport = tr
}

//
// Get current Transport
//
func (self *HttpClient) GetTransport() http.RoundTripper {
	return self.client.Transport
}

//
// Allow connections via HTTPS even if something is wrong with the certificate
// (self-signed or expired)
//
func (self *HttpClient) AllowInsecure(insecure bool) {
	if insecure {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}

		self.client.Transport = tr
	} else {
		self.client.Transport = nil
	}
}

//
// Set connection timeout
//
func (self *HttpClient) SetTimeout(t time.Duration) {
	self.client.Timeout = t
}

//
// Set LocalAddr in Dialer
//
func (self *HttpClient) SetLocalAddr(addr string) {
	transport, ok := self.client.Transport.(*http.Transport)
	if transport == nil {
		if transport, ok = http.DefaultTransport.(*http.Transport); !ok {
			log.Println("SetLocalAddr for http.DefaultTransport != http.Transport")
			return
		}
	} else if !ok {
		log.Println("SetLocalAddr for client.Transport != http.Transport")
		return
	}
	if tcpaddr, err := net.ResolveTCPAddr("tcp", addr); err == nil {
		transport.Dial = (&net.Dialer{
			Timeout:   30 * time.Second, // defaults from net/http DefaultTransport
			KeepAlive: 30 * time.Second, // defaults from net/http DefaultTransport
			LocalAddr: tcpaddr,
		}).Dial
	} else {
		log.Println("Failed to resolve", addr, " to a TCP address")
	}
}

//
// add default headers plus extra headers
//
func (self *HttpClient) addHeaders(req *http.Request, headers map[string]string) {

	if len(self.UserAgent) > 0 {
		req.Header.Set("User-Agent", self.UserAgent)
	}

	for k, v := range self.Headers {
		req.Header.Set(k, v)
	}

	for k, v := range headers {
		if strings.ToLower(k) == "content-length" {
			if len, err := strconv.Atoi(v); err == nil && req.ContentLength <= 0 {
				req.ContentLength = int64(len)
			}
		} else {
			req.Header.Set(k, v)
		}
	}

}

//
// the callback for CheckRedirect, used to pass along the headers in case of redirection
//
func (self *HttpClient) checkRedirect(req *http.Request, via []*http.Request) error {
	if req.Method == "HEAD" {
		// don't follow redirects on a HEAD request
		return NoRedirect
	}

	if self.Verbose {
		log.Println("REDIRECT:", len(via), req.URL)
	}

	if len(via) >= 10 {
		return TooManyRedirects
	}

	// TODO: check for same host before adding headers
	self.addHeaders(req, nil)
	return nil
}

//
// Create a request object given the method, path, body and extra headers
//
func (self *HttpClient) Request(method string, urlpath string, body io.Reader, headers map[string]string) (req *http.Request) {
	if u, err := self.BaseURL.Parse(urlpath); err != nil {
		log.Fatal(err)
	} else {
		urlpath = u.String()
	}

	req, err := http.NewRequest(method, urlpath, body)
	if err != nil {
		log.Fatal(err)
	}

	self.addHeaders(req, headers)
	return
}

//
// Execute request
//
func (self *HttpClient) Do(req *http.Request) (*HttpResponse, error) {
	if self.Verbose {
		log.Println("REQUEST:", req.Method, req.URL, pretty.PrettyFormat(req.Header))
	}

	resp, err := self.client.Do(req)
	if urlerr, ok := err.(*url.Error); ok && urlerr.Err == NoRedirect {
		err = nil // redirect on HEAD is not an error
	}
	if err == nil {
		if self.Verbose {
			log.Println("RESPONSE:", resp.Status, pretty.PrettyFormat(resp.Header))
		}

		return &HttpResponse{*resp}, nil
	} else {
		if self.Verbose {
			log.Println("ERROR:", err, "REQUEST:", req.Method, req.URL, pretty.PrettyFormat(req.Header))
		}

		CloseResponse(resp)
		return nil, err
	}
}

//
// Execute a DELETE request
//
func (self *HttpClient) Delete(path string, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("DELETE", path, nil, headers)
	return self.Do(req)
}

//
// Execute a HEAD request
//
func (self *HttpClient) Head(path string, params map[string]interface{}, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("HEAD", URLWithParams(path, params).String(), nil, headers)
	return self.Do(req)
}

//
// Execute a GET request
//
func (self *HttpClient) Get(path string, params map[string]interface{}, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("GET", URLWithParams(path, params).String(), nil, headers)
	return self.Do(req)
}

//
// Execute a POST request
//
func (self *HttpClient) Post(path string, content io.Reader, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("POST", path, content, headers)
	return self.Do(req)
}

//
// Execute a PUT request
//
func (self *HttpClient) Put(path string, content io.Reader, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("PUT", path, content, headers)
	return self.Do(req)
}
<commit_msg>Added HttpClient.Close boolean. It propagates to all requests to signal that the connection should be closed when the request complete.<commit_after>package httpclient

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/gobs/pretty"
	"github.com/gobs/simplejson"
)

var (
	DefaultClient = &http.Client{} // we use our own default client, so we can change the TLS configuration

	NoRedirect       = errors.New("No redirect")
	TooManyRedirects = errors.New("stopped after 10 redirects")
)

//
// Allow connections via HTTPS even if something is wrong with the certificate
// (self-signed or expired)
//
func AllowInsecure(insecure bool) {
	if insecure {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}

		DefaultClient.Transport = tr
	} else {
		DefaultClient.Transport = nil
	}
}

//
// Set connection timeout
//
func SetTimeout(t time.Duration) {
	DefaultClient.Timeout = t
}

//
// HTTP error
//
type HttpError struct {
	Code    int
	Message string
}

func (e HttpError) Error() string {
	return e.Message
}

func (e HttpError) String() string {
	return fmt.Sprintf("ERROR: %v %v", e.Code, e.Message)
}

//
// CloseResponse makes sure we close the response body
//
func CloseResponse(r *http.Response) {
	if r != nil && r.Body != nil {
		io.Copy(ioutil.Discard, r.Body)
		r.Body.Close()
	}
}

//
// A wrapper for http.Response
//
type HttpResponse struct {
	http.Response
}

func (r *HttpResponse) ContentType() string {
	content_type := r.Header.Get("Content-Type")
	if len(content_type) == 0 {
		return content_type
	}

	return strings.TrimSpace(strings.Split(content_type, ";")[0])
}

//
// Close makes sure that all data from the body is read
// before closing the reader.
//
// If that is not the desider behaviour, just call HttpResponse.Body.Close()
//
func (r *HttpResponse) Close() {
	if r != nil {
		CloseResponse(&r.Response)
	}
}

//
// ResponseError checks the StatusCode and return an error if needed.
// The error is of type HttpError
//
func (r *HttpResponse) ResponseError() error {
	class := r.StatusCode / 100
	if class != 2 && class != 3 {
		return HttpError{Code: r.StatusCode, Message: "HTTP " + r.Status}
	}

	return nil
}

//
// Check if the input value is a "primitive" that can be safely stringified
//
func canStringify(v reflect.Value) bool {
	switch v.Kind() {
	case reflect.Bool,
		reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
		reflect.Float32, reflect.Float64,
		reflect.String:
		return true
	default:
		return false
	}
}

//
// Given a base URL and a bag of parameteters returns the URL with the encoded parameters
//
func URLWithPathParams(base string, path string, params map[string]interface{}) (u *url.URL) {

	u, err := url.Parse(base)
	if err != nil {
		log.Fatal(err)
	}

	if len(path) > 0 {
		u, err = u.Parse(path)
		if err != nil {
			log.Fatal(err)
		}
	}

	q := u.Query()

	for k, v := range params {
		val := reflect.ValueOf(v)

		switch val.Kind() {
		case reflect.Slice:
			if val.IsNil() { // TODO: add an option to ignore empty values
				q.Set(k, "")
				continue
			}
			fallthrough

		case reflect.Array:
			for i := 0; i < val.Len(); i++ {
				av := val.Index(i)

				if canStringify(av) {
					q.Add(k, fmt.Sprintf("%v", av))
				}
			}

		default:
			if canStringify(val) {
				q.Set(k, fmt.Sprintf("%v", v))
			} else {
				log.Fatal("Invalid type ", val)
			}
		}
	}

	u.RawQuery = q.Encode()
	return u
}

func URLWithParams(base string, params map[string]interface{}) (u *url.URL) {
	return URLWithPathParams(base, "", params)
}

//
// http.Get with params
//
func Get(urlStr string, params map[string]interface{}) (*HttpResponse, error) {
	resp, err := DefaultClient.Get(URLWithParams(urlStr, params).String())
	if err == nil {
		return &HttpResponse{*resp}, nil
	} else {
		CloseResponse(resp)
		return nil, err
	}
}

//
// http.Post with params
//
func Post(urlStr string, params map[string]interface{}) (*HttpResponse, error) {
	resp, err := DefaultClient.PostForm(urlStr, URLWithParams(urlStr, params).Query())
	if err == nil {
		return &HttpResponse{*resp}, nil
	} else {
		CloseResponse(resp)
		return nil, err
	}
}

//
//  Read the body
//
func (resp *HttpResponse) Content() []byte {
	body, _ := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	return body
}

//
//  Try to parse the response body as JSON
//
func (resp *HttpResponse) Json() (json *simplejson.Json) {
	json, _ = simplejson.LoadBytes(resp.Content())
	return
}

////////////////////////////////////////////////////////////////////////

//
// http.Client with some defaults and stuff
//
type HttpClient struct {
	// the http.Client
	client *http.Client

	// the base URL for this client
	BaseURL *url.URL

	// the client UserAgent string
	UserAgent string

	// Common headers to be passed on each request
	Headers map[string]string

	// if Verbose, log request and response info
	Verbose bool

	// if Close, all requests will set Connection: close
	// (no keep-alive)
	Close bool
}

//
// Create a new HttpClient
//
func NewHttpClient(base string) (httpClient *HttpClient) {
	httpClient = new(HttpClient)
	httpClient.client = &http.Client{CheckRedirect: httpClient.checkRedirect}
	httpClient.Headers = make(map[string]string)

	if u, err := url.Parse(base); err != nil {
		log.Fatal(err)
	} else {
		httpClient.BaseURL = u
	}

	return
}

//
// Set Transport
//
func (self *HttpClient) SetTransport(tr http.RoundTripper) {
	self.client.Transport = tr
}

//
// Get current Transport
//
func (self *HttpClient) GetTransport() http.RoundTripper {
	return self.client.Transport
}

//
// Allow connections via HTTPS even if something is wrong with the certificate
// (self-signed or expired)
//
func (self *HttpClient) AllowInsecure(insecure bool) {
	if insecure {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}

		self.client.Transport = tr
	} else {
		self.client.Transport = nil
	}
}

//
// Set connection timeout
//
func (self *HttpClient) SetTimeout(t time.Duration) {
	self.client.Timeout = t
}

//
// Set LocalAddr in Dialer
//
func (self *HttpClient) SetLocalAddr(addr string) {
	transport, ok := self.client.Transport.(*http.Transport)
	if transport == nil {
		if transport, ok = http.DefaultTransport.(*http.Transport); !ok {
			log.Println("SetLocalAddr for http.DefaultTransport != http.Transport")
			return
		}
	} else if !ok {
		log.Println("SetLocalAddr for client.Transport != http.Transport")
		return
	}
	if tcpaddr, err := net.ResolveTCPAddr("tcp", addr); err == nil {
		transport.Dial = (&net.Dialer{
			Timeout:   30 * time.Second, // defaults from net/http DefaultTransport
			KeepAlive: 30 * time.Second, // defaults from net/http DefaultTransport
			LocalAddr: tcpaddr,
		}).Dial
	} else {
		log.Println("Failed to resolve", addr, " to a TCP address")
	}
}

//
// add default headers plus extra headers
//
func (self *HttpClient) addHeaders(req *http.Request, headers map[string]string) {

	if len(self.UserAgent) > 0 {
		req.Header.Set("User-Agent", self.UserAgent)
	}

	for k, v := range self.Headers {
		req.Header.Set(k, v)
	}

	for k, v := range headers {
		if strings.ToLower(k) == "content-length" {
			if len, err := strconv.Atoi(v); err == nil && req.ContentLength <= 0 {
				req.ContentLength = int64(len)
			}
		} else {
			req.Header.Set(k, v)
		}
	}

}

//
// the callback for CheckRedirect, used to pass along the headers in case of redirection
//
func (self *HttpClient) checkRedirect(req *http.Request, via []*http.Request) error {
	if req.Method == "HEAD" {
		// don't follow redirects on a HEAD request
		return NoRedirect
	}

	if self.Verbose {
		log.Println("REDIRECT:", len(via), req.URL)
	}

	if len(via) >= 10 {
		return TooManyRedirects
	}

	// TODO: check for same host before adding headers
	self.addHeaders(req, nil)
	return nil
}

//
// Create a request object given the method, path, body and extra headers
//
func (self *HttpClient) Request(method string, urlpath string, body io.Reader, headers map[string]string) (req *http.Request) {
	if u, err := self.BaseURL.Parse(urlpath); err != nil {
		log.Fatal(err)
	} else {
		urlpath = u.String()
	}

	req, err := http.NewRequest(method, urlpath, body)
	if err != nil {
		log.Fatal(err)
	}

	req.Close = self.Close

	self.addHeaders(req, headers)
	return
}

//
// Execute request
//
func (self *HttpClient) Do(req *http.Request) (*HttpResponse, error) {
	if self.Verbose {
		log.Println("REQUEST:", req.Method, req.URL, pretty.PrettyFormat(req.Header))
	}

	resp, err := self.client.Do(req)
	if urlerr, ok := err.(*url.Error); ok && urlerr.Err == NoRedirect {
		err = nil // redirect on HEAD is not an error
	}
	if err == nil {
		if self.Verbose {
			log.Println("RESPONSE:", resp.Status, pretty.PrettyFormat(resp.Header))
		}

		return &HttpResponse{*resp}, nil
	} else {
		if self.Verbose {
			log.Println("ERROR:", err, "REQUEST:", req.Method, req.URL, pretty.PrettyFormat(req.Header))
		}

		CloseResponse(resp)
		return nil, err
	}
}

//
// Execute a DELETE request
//
func (self *HttpClient) Delete(path string, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("DELETE", path, nil, headers)
	return self.Do(req)
}

//
// Execute a HEAD request
//
func (self *HttpClient) Head(path string, params map[string]interface{}, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("HEAD", URLWithParams(path, params).String(), nil, headers)
	return self.Do(req)
}

//
// Execute a GET request
//
func (self *HttpClient) Get(path string, params map[string]interface{}, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("GET", URLWithParams(path, params).String(), nil, headers)
	return self.Do(req)
}

//
// Execute a POST request
//
func (self *HttpClient) Post(path string, content io.Reader, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("POST", path, content, headers)
	return self.Do(req)
}

//
// Execute a PUT request
//
func (self *HttpClient) Put(path string, content io.Reader, headers map[string]string) (*HttpResponse, error) {
	req := self.Request("PUT", path, content, headers)
	return self.Do(req)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	. "gopkg.in/check.v1"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io/ioutil"
	"vip/fetch"
	"vip/test"
)

var (
	sizes = []int{
		250,
		500,
		160,
		720,
		1024,
		683,
		431,
	}
)

var (
	_ = Suite(&ResizeSuite{})
)

type ResizeSuite struct{}

func (s *ResizeSuite) SetUpSuite(c *C) {
	setUpSuite(c)
}

func (s *ResizeSuite) SetUpTest(c *C) {
	setUpTest(c)

	storage = test.NewStore()
}

func (s *ResizeSuite) BenchmarkThumbnailResize(c *C) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	ctx := &fetch.CacheContext{
		Width: 160,
	}

	for i := 0; i < c.N; i++ {
		// Need a new io.Reader on every iteration
		buf := bytes.NewBuffer(file)
		_, err := fetch.Resize(buf, ctx)
		c.Assert(err, IsNil)
	}
}

func (s *ResizeSuite) BenchmarkLargeResize(c *C) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	ctx := &fetch.CacheContext{
		Width: 720,
	}

	for i := 0; i < c.N; i++ {
		// Need a new io.Reader on every iteration
		buf := bytes.NewBuffer(file)
		_, err := fetch.Resize(buf, ctx)
		c.Assert(err, IsNil)
	}
}

func (s *ResizeSuite) TestResizeImage(c *C) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	for _, size := range sizes {
		ctx := &fetch.CacheContext{
			Width: size,
		}

		buf := bytes.NewBuffer(file)
		resized, err := fetch.Resize(buf, ctx)
		c.Assert(err, IsNil)

		image, _, err := image.Decode(resized)
		c.Assert(err, IsNil)
		c.Assert(image.Bounds().Size().X, Equals, size)
	}
}

func (s *ResizeSuite) insertMockImage() (*fetch.CacheContext, error) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	if err != nil {
		return nil, err
	}

	// Push the file data into the mock datastore
	storage.Put("test_bucket", "test_id", file, "image/jpeg")

	return &fetch.CacheContext{
		ImageId: "test_id",
		Bucket:  "test_bucket",
	}, err
}

func (s *ResizeSuite) TestOriginalColdCache(c *C) {
	// Open the file once to get it's size
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	img, _, err := image.Decode(bytes.NewBuffer(file))
	c.Assert(err, IsNil)

	originalSize := img.Bounds().Size().X

	// A single, unresized image is in the database/store
	ctx, err := s.insertMockImage()
	c.Assert(err, IsNil)

	// Run the image resize request
	data, err := fetch.ImageData(storage, ctx)
	c.Assert(err, IsNil)

	// Verify the size of the resulting byte slice
	img, _, err = image.Decode(bytes.NewBuffer(data))
	c.Assert(err, IsNil)
	c.Assert(img.Bounds().Size().X, Equals, originalSize)
}

func (s *ResizeSuite) TestResizeColdCache(c *C) {
	// A single, unresized image is in the database/store
	mockCtx, err := s.insertMockImage()
	c.Assert(err, IsNil)

	for _, size := range sizes {
		ctx := &fetch.CacheContext{
			ImageId: mockCtx.ImageId,
			Bucket:  mockCtx.Bucket,
			Width:   size,
		}

		// Run the image resize request
		data, err := fetch.ImageData(storage, ctx)
		c.Assert(err, IsNil)

		// Verify the size of the resulting byte slice
		img, _, err := image.Decode(bytes.NewBuffer(data))
		c.Assert(err, IsNil)
		c.Assert(img.Bounds().Size().X, Equals, size)
	}
}

func (s *ResizeSuite) TestResizeCropColdCache(c *C) {
	// A single, unresized image is in the database/store
	mockCtx, err := s.insertMockImage()
	c.Assert(err, IsNil)

	for _, size := range sizes {
		ctx := &fetch.CacheContext{
			ImageId: mockCtx.ImageId,
			Bucket:  mockCtx.Bucket,
			Width:   size,
			Crop:    true,
		}

		// Run the image resize request
		data, err := fetch.ImageData(storage, ctx)
		c.Assert(err, IsNil)

		// Verify the size of the resulting byte slice
		img, _, err := image.Decode(bytes.NewBuffer(data))
		c.Assert(err, IsNil)
		c.Assert(img.Bounds().Size().X, Equals, img.Bounds().Size().Y)
		c.Assert(img.Bounds().Size().X > 0, Equals, true)
		c.Assert(img.Bounds().Size().X <= size, Equals, true)
	}
}
<commit_msg>Use Readers instead of Buffers<commit_after>package main

import (
	"bytes"
	. "gopkg.in/check.v1"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io/ioutil"
	"vip/fetch"
	"vip/test"
)

var (
	sizes = []int{
		250,
		500,
		160,
		720,
		1024,
		683,
		431,
	}
)

var (
	_ = Suite(&ResizeSuite{})
)

type ResizeSuite struct{}

func (s *ResizeSuite) SetUpSuite(c *C) {
	setUpSuite(c)
}

func (s *ResizeSuite) SetUpTest(c *C) {
	setUpTest(c)

	storage = test.NewStore()
}

func (s *ResizeSuite) BenchmarkThumbnailResize(c *C) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	ctx := &fetch.CacheContext{
		Width: 160,
	}

	for i := 0; i < c.N; i++ {
		// Need a new io.Reader on every iteration
		buf := bytes.NewReader(file)
		_, err := fetch.Resize(buf, ctx)
		c.Assert(err, IsNil)
	}
}

func (s *ResizeSuite) BenchmarkLargeResize(c *C) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	ctx := &fetch.CacheContext{
		Width: 720,
	}

	for i := 0; i < c.N; i++ {
		// Need a new io.Reader on every iteration
		buf := bytes.NewReader(file)
		_, err := fetch.Resize(buf, ctx)
		c.Assert(err, IsNil)
	}
}

func (s *ResizeSuite) TestResizeImage(c *C) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	for _, size := range sizes {
		ctx := &fetch.CacheContext{
			Width: size,
		}

		buf := bytes.NewReader(file)
		resized, err := fetch.Resize(buf, ctx)
		c.Assert(err, IsNil)

		image, _, err := image.Decode(resized)
		c.Assert(err, IsNil)
		c.Assert(image.Bounds().Size().X, Equals, size)
	}
}

func (s *ResizeSuite) insertMockImage() (*fetch.CacheContext, error) {
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	if err != nil {
		return nil, err
	}

	// Push the file data into the mock datastore
	storage.Put("test_bucket", "test_id", file, "image/jpeg")

	return &fetch.CacheContext{
		ImageId: "test_id",
		Bucket:  "test_bucket",
	}, err
}

func (s *ResizeSuite) TestOriginalColdCache(c *C) {
	// Open the file once to get it's size
	file, err := ioutil.ReadFile("test/AWESOME.jpg")
	c.Assert(err, IsNil)

	img, _, err := image.Decode(bytes.NewReader(file))
	c.Assert(err, IsNil)

	originalSize := img.Bounds().Size().X

	// A single, unresized image is in the database/store
	ctx, err := s.insertMockImage()
	c.Assert(err, IsNil)

	// Run the image resize request
	data, err := fetch.ImageData(storage, ctx)
	c.Assert(err, IsNil)

	// Verify the size of the resulting byte slice
	img, _, err = image.Decode(bytes.NewReader(data))
	c.Assert(err, IsNil)
	c.Assert(img.Bounds().Size().X, Equals, originalSize)
}

func (s *ResizeSuite) TestResizeColdCache(c *C) {
	// A single, unresized image is in the database/store
	mockCtx, err := s.insertMockImage()
	c.Assert(err, IsNil)

	for _, size := range sizes {
		ctx := &fetch.CacheContext{
			ImageId: mockCtx.ImageId,
			Bucket:  mockCtx.Bucket,
			Width:   size,
		}

		// Run the image resize request
		data, err := fetch.ImageData(storage, ctx)
		c.Assert(err, IsNil)

		// Verify the size of the resulting byte slice
		img, _, err := image.Decode(bytes.NewReader(data))
		c.Assert(err, IsNil)
		c.Assert(img.Bounds().Size().X, Equals, size)
	}
}

func (s *ResizeSuite) TestResizeCropColdCache(c *C) {
	// A single, unresized image is in the database/store
	mockCtx, err := s.insertMockImage()
	c.Assert(err, IsNil)

	for _, size := range sizes {
		ctx := &fetch.CacheContext{
			ImageId: mockCtx.ImageId,
			Bucket:  mockCtx.Bucket,
			Width:   size,
			Crop:    true,
		}

		// Run the image resize request
		data, err := fetch.ImageData(storage, ctx)
		c.Assert(err, IsNil)

		// Verify the size of the resulting byte slice
		img, _, err := image.Decode(bytes.NewReader(data))
		c.Assert(err, IsNil)
		c.Assert(img.Bounds().Size().X, Equals, img.Bounds().Size().Y)
		c.Assert(img.Bounds().Size().X > 0, Equals, true)
		c.Assert(img.Bounds().Size().X <= size, Equals, true)
	}
}
<|endoftext|>


<commit_before>package dbscan

import (
	"github.com/lfritz/clustering/geometry"
	"math"
	"math/rand"
	"reflect"
)

// Kmeans implements the k-mean clustering algorithm
// (https://en.m.wikipedia.org/wiki/K-means_clustering).
func Kmeans(points [][2]float64, k int) []int {
	centroids := initialCentroids(points, k)
	var clustering []int = nil
	for {
		nextClustering := clusteringForCentroids(points, centroids)
		if reflect.DeepEqual(nextClustering, clustering) {
			break
		}
		clustering = nextClustering
		centroids = centroidsForClusters(points, k, clustering)

	}
	return clustering
}

// TODO:
// - try using a struct for intermediate results
// - re-name packages
// - try k-means++
// - try running k-means repeatedly

// randK randomly selects k numbers in [0..n), without duplicates.
func randK(k, n int) []int {
	result := make([]int, k)
	for i := range result {
		x := rand.Intn(n - i)
		done := false
		for j, other := range result[:i] {
			if other > x {
				copy(result[j+1:], result[j:])
				result[j] = x
				done = true
				break
			}
			x++
		}
		if !done {
			result[i] = x
		}
	}
	return result
}

// initialCentroids generates an initial set of centroids for the k-mean algorithm using the Forgy
// method.
func initialCentroids(points [][2]float64, k int) [][2]float64 {
	centroids := make([][2]float64, k)
	for i, x := range randK(k, len(points)) {
		centroids[i] = points[x]
	}
	return centroids
}

func closest(ps [][2]float64, q [2]float64) int {
	closest := 0
	minimumDistance := math.Inf(+1)
	for i, p := range ps {
		distance := geometry.Distance(p, q)
		if distance < minimumDistance {
			minimumDistance = distance
			closest = i
		}
	}
	return closest
}

// clusteringForCentroids takes a set of points and a set of centroids and returns a clustering that
// assigns each point to the closest centroid.
func clusteringForCentroids(points [][2]float64, centroids [][2]float64) []int {
	clustering := make([]int, len(points))
	for i, p := range points {
		clustering[i] = closest(centroids, p)
	}
	return clustering
}

// centroidsForClusters calculates a centroid for each cluster as the mean of its points.
func centroidsForClusters(points [][2]float64, k int, clustering []int) [][2]float64 {
	count := make([]int, k)
	centroids := make([][2]float64, k)
	for i, c := range clustering {
		centroids[c][0] += points[i][0]
		centroids[c][1] += points[i][1]
		count[c]++
	}
	for i, c := range count {
		if c != 0 {
			centroids[i][0] /= float64(c)
			centroids[i][1] /= float64(c)
		}
	}
	return centroids
}
<commit_msg>Fix typos<commit_after>package dbscan

import (
	"github.com/lfritz/clustering/geometry"
	"math"
	"math/rand"
	"reflect"
)

// Kmeans implements the k-means clustering algorithm
// (https://en.m.wikipedia.org/wiki/K-means_clustering).
func Kmeans(points [][2]float64, k int) []int {
	centroids := initialCentroids(points, k)
	var clustering []int = nil
	for {
		nextClustering := clusteringForCentroids(points, centroids)
		if reflect.DeepEqual(nextClustering, clustering) {
			break
		}
		clustering = nextClustering
		centroids = centroidsForClusters(points, k, clustering)

	}
	return clustering
}

// TODO:
// - try using a struct for intermediate results
// - re-name packages
// - try k-means++
// - try running k-means repeatedly

// randK randomly selects k numbers in [0..n), without duplicates.
func randK(k, n int) []int {
	result := make([]int, k)
	for i := range result {
		x := rand.Intn(n - i)
		done := false
		for j, other := range result[:i] {
			if other > x {
				copy(result[j+1:], result[j:])
				result[j] = x
				done = true
				break
			}
			x++
		}
		if !done {
			result[i] = x
		}
	}
	return result
}

// initialCentroids generates an initial set of centroids for the k-means algorithm using the Forgy
// method.
func initialCentroids(points [][2]float64, k int) [][2]float64 {
	centroids := make([][2]float64, k)
	for i, x := range randK(k, len(points)) {
		centroids[i] = points[x]
	}
	return centroids
}

func closest(ps [][2]float64, q [2]float64) int {
	closest := 0
	minimumDistance := math.Inf(+1)
	for i, p := range ps {
		distance := geometry.Distance(p, q)
		if distance < minimumDistance {
			minimumDistance = distance
			closest = i
		}
	}
	return closest
}

// clusteringForCentroids takes a set of points and a set of centroids and returns a clustering that
// assigns each point to the closest centroid.
func clusteringForCentroids(points [][2]float64, centroids [][2]float64) []int {
	clustering := make([]int, len(points))
	for i, p := range points {
		clustering[i] = closest(centroids, p)
	}
	return clustering
}

// centroidsForClusters calculates a centroid for each cluster as the mean of its points.
func centroidsForClusters(points [][2]float64, k int, clustering []int) [][2]float64 {
	count := make([]int, k)
	centroids := make([][2]float64, k)
	for i, c := range clustering {
		centroids[c][0] += points[i][0]
		centroids[c][1] += points[i][1]
		count[c]++
	}
	for i, c := range count {
		if c != 0 {
			centroids[i][0] /= float64(c)
			centroids[i][1] /= float64(c)
		}
	}
	return centroids
}
<|endoftext|>


<commit_before>package goat

import (
	"net"
)

// ConnHandler interface method Handle defines how to handle incoming network connections
type ConnHandler interface {
	Handle(c net.Conn) bool
}

// HttpConnHandler handles incoming HTTP (TCP) network connections
type HttpConnHandler struct {
}

func (h HttpConnHandler) Handle(c net.Conn) bool {
	return true
}

// UdpConnHandler handles incoming UDP network connections
type UdpConnHandler struct {
}

func (u UdpConnHandler) Handle(c net.Conn) bool {
	return true
}
<commit_msg>Add basic HTTP read/write over socket<commit_after>package goat

import (
	"fmt"
	"net"
	"strings"
)

// ConnHandler interface method Handle defines how to handle incoming network connections
type ConnHandler interface {
	Handle(c net.Conn) bool
}

// HttpConnHandler handles incoming HTTP (TCP) network connections
type HttpConnHandler struct {
}

// Handle an incoming HTTP request and provide a HTTP response
func (h HttpConnHandler) Handle(c net.Conn) bool {
	// Read in data from socket
	var buf = make([]byte, 1024)
	c.Read(buf)

	// TODO: remove temporary printing and fake response
	fmt.Println("http: ", string(buf))
	res := []string {
		"HTTP/1.1 200 OK\r\n",
		"Content-Type: text/plain\r\n",
		"Content-Length: 4\r\n",
		"Connection: close\r\n\r\n",
		"goat\r\n",
	}

	// Write response
	c.Write(strings.Join(res, ""))
	c.Close()

	return true
}

// UdpConnHandler handles incoming UDP network connections
type UdpConnHandler struct {
}

func (u UdpConnHandler) Handle(c net.Conn) bool {
	return true
}
<|endoftext|>


<commit_before><commit_msg>Add clojure catching. Change the collaspe goroutine strategy.<commit_after><|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner

import (
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/names"

	"github.com/juju/juju/agent"
	apiprovisioner "github.com/juju/juju/api/provisioner"
	"github.com/juju/juju/controller/authentication"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/watcher"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/worker/catacomb"
	"github.com/juju/juju/worker/environ"
)

var logger = loggo.GetLogger("juju.provisioner")

// Ensure our structs implement the required Provisioner interface.
var _ Provisioner = (*environProvisioner)(nil)
var _ Provisioner = (*containerProvisioner)(nil)

var (
	retryStrategyDelay = 10 * time.Second
	retryStrategyCount = 3
)

// Provisioner represents a running provisioner worker.
type Provisioner interface {
	worker.Worker
	getMachineWatcher() (watcher.StringsWatcher, error)
	getRetryWatcher() (watcher.NotifyWatcher, error)
}

// environProvisioner represents a running provisioning worker for machine nodes
// belonging to an environment.
type environProvisioner struct {
	provisioner
	environ environs.Environ
	configObserver
}

// containerProvisioner represents a running provisioning worker for containers
// hosted on a machine.
type containerProvisioner struct {
	provisioner
	containerType instance.ContainerType
	machine       *apiprovisioner.Machine
	configObserver
}

// provisioner providers common behaviour for a running provisioning worker.
type provisioner struct {
	Provisioner
	st          *apiprovisioner.State
	agentConfig agent.Config
	broker      environs.InstanceBroker
	toolsFinder ToolsFinder
	catacomb    catacomb.Catacomb
}

// RetryStrategy defines the retry behavior when encountering a retryable
// error during provisioning.
type RetryStrategy struct {
	retryDelay time.Duration
	retryCount int
}

// NewRetryStrategy returns a new retry strategy with the specified delay and
// count for use with retryable provisioning errors.
func NewRetryStrategy(delay time.Duration, count int) RetryStrategy {
	return RetryStrategy{
		retryDelay: delay,
		retryCount: count,
	}
}

// configObserver is implemented so that tests can see
// when the environment configuration changes.
type configObserver struct {
	sync.Mutex
	observer chan<- *config.Config
}

// notify notifies the observer of a configuration change.
func (o *configObserver) notify(cfg *config.Config) {
	o.Lock()
	if o.observer != nil {
		o.observer <- cfg
	}
	o.Unlock()
}

// Kill implements worker.Worker.Kill.
func (p *provisioner) Kill() {
	p.catacomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (p *provisioner) Wait() error {
	return p.catacomb.Wait()
}

// getToolsFinder returns a ToolsFinder for the provided State.
// This exists for mocking.
var getToolsFinder = func(st *apiprovisioner.State) ToolsFinder {
	return st
}

// getStartTask creates a new worker for the provisioner,
func (p *provisioner) getStartTask(harvestMode config.HarvestMode) (ProvisionerTask, error) {
	auth, err := authentication.NewAPIAuthenticator(p.st)
	if err != nil {
		return nil, err
	}
	// Start responding to changes in machines, and to any further updates
	// to the environment config.
	machineWatcher, err := p.getMachineWatcher()
	if err != nil {
		return nil, err
	}
	retryWatcher, err := p.getRetryWatcher()
	if err != nil && !errors.IsNotImplemented(err) {
		return nil, err
	}
	tag := p.agentConfig.Tag()
	machineTag, ok := tag.(names.MachineTag)
	if !ok {
		errors.Errorf("expacted names.MachineTag, got %T", tag)
	}

	envCfg, err := p.st.ModelConfig()
	if err != nil {
		return nil, errors.Annotate(err, "could not retrieve the model config.")
	}

	secureServerConnection := false
	if info, ok := p.agentConfig.StateServingInfo(); ok {
		secureServerConnection = info.CAPrivateKey != ""
	}
	task, err := NewProvisionerTask(
		machineTag,
		harvestMode,
		p.st,
		p.toolsFinder,
		machineWatcher,
		retryWatcher,
		p.broker,
		auth,
		envCfg.ImageStream(),
		secureServerConnection,
		RetryStrategy{retryDelay: retryStrategyDelay, retryCount: retryStrategyCount},
	)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return task, nil
}

// NewEnvironProvisioner returns a new Provisioner for an environment.
// When new machines are added to the state, it allocates instances
// from the environment and allocates them to the new machines.
func NewEnvironProvisioner(st *apiprovisioner.State, agentConfig agent.Config) (Provisioner, error) {
	p := &environProvisioner{
		provisioner: provisioner{
			st:          st,
			agentConfig: agentConfig,
			toolsFinder: getToolsFinder(st),
		},
	}
	p.Provisioner = p
	logger.Tracef("Starting environ provisioner for %q", p.agentConfig.Tag())

	err := catacomb.Invoke(catacomb.Plan{
		Site: &p.catacomb,
		Work: p.loop,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return p, nil
}

func (p *environProvisioner) loop() error {
	var modelConfigChanges <-chan struct{}
	modelWatcher, err := p.st.WatchForModelConfigChanges()
	if err != nil {
		return loggedErrorStack(errors.Trace(err))
	}
	if err := p.catacomb.Add(modelWatcher); err != nil {
		return errors.Trace(err)
	}
	modelConfigChanges = modelWatcher.Changes()

	p.environ, err = environ.WaitForEnviron(modelWatcher, p.st, p.catacomb.Dying())
	if err != nil {
		if err == environ.ErrWaitAborted {
			return p.catacomb.ErrDying()
		}
		return loggedErrorStack(errors.Trace(err))
	}
	p.broker = p.environ

	harvestMode := p.environ.Config().ProvisionerHarvestMode()
	task, err := p.getStartTask(harvestMode)
	if err != nil {
		return loggedErrorStack(errors.Trace(err))
	}
	if err := p.catacomb.Add(task); err != nil {
		return errors.Trace(err)
	}

	for {
		select {
		case <-p.catacomb.Dying():
			return p.catacomb.ErrDying()
		case _, ok := <-modelConfigChanges:
			if !ok {
				return errors.New("model configuration watcher closed")
			}
			environConfig, err := p.st.ModelConfig()
			if err != nil {
				return errors.Annotate(err, "cannot load model configuration")
			}
			if err := p.setConfig(environConfig); err != nil {
				return errors.Annotate(err, "loaded invalid model configuration")
			}
			task.SetHarvestMode(environConfig.ProvisionerHarvestMode())
		}
	}
}

func (p *environProvisioner) getMachineWatcher() (watcher.StringsWatcher, error) {
	return p.st.WatchModelMachines()
}

func (p *environProvisioner) getRetryWatcher() (watcher.NotifyWatcher, error) {
	return p.st.WatchMachineErrorRetry()
}

// setConfig updates the environment configuration and notifies
// the config observer.
func (p *environProvisioner) setConfig(environConfig *config.Config) error {
	if err := p.environ.SetConfig(environConfig); err != nil {
		return err
	}
	p.configObserver.notify(environConfig)
	return nil
}

// NewContainerProvisioner returns a new Provisioner. When new machines
// are added to the state, it allocates instances from the environment
// and allocates them to the new machines.
func NewContainerProvisioner(
	containerType instance.ContainerType,
	st *apiprovisioner.State,
	agentConfig agent.Config,
	broker environs.InstanceBroker,
	toolsFinder ToolsFinder,
) (Provisioner, error) {

	p := &containerProvisioner{
		provisioner: provisioner{
			st:          st,
			agentConfig: agentConfig,
			broker:      broker,
			toolsFinder: toolsFinder,
		},
		containerType: containerType,
	}
	p.Provisioner = p
	logger.Tracef("Starting %s provisioner for %q", p.containerType, p.agentConfig.Tag())

	err := catacomb.Invoke(catacomb.Plan{
		Site: &p.catacomb,
		Work: p.loop,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return p, nil
}

func (p *containerProvisioner) loop() error {
	modelWatcher, err := p.st.WatchForModelConfigChanges()
	if err != nil {
		return errors.Trace(err)
	}
	if err := p.catacomb.Add(modelWatcher); err != nil {
		return errors.Trace(err)
	}

	config, err := p.st.ModelConfig()
	if err != nil {
		return err
	}
	harvestMode := config.ProvisionerHarvestMode()

	task, err := p.getStartTask(harvestMode)
	if err != nil {
		return err
	}
	if err := p.catacomb.Add(task); err != nil {
		return errors.Trace(err)
	}

	for {
		select {
		case <-p.catacomb.Dying():
			return p.catacomb.ErrDying()
		case _, ok := <-modelWatcher.Changes():
			if !ok {
				return errors.New("model configuratioon watch closed")
			}
			modelConfig, err := p.st.ModelConfig()
			if err != nil {
				return errors.Annotate(err, "cannot load model configuration")
			}
			p.configObserver.notify(modelConfig)
			task.SetHarvestMode(modelConfig.ProvisionerHarvestMode())
		}
	}
}

func (p *containerProvisioner) getMachine() (*apiprovisioner.Machine, error) {
	if p.machine == nil {
		tag := p.agentConfig.Tag()
		machineTag, ok := tag.(names.MachineTag)
		if !ok {
			return nil, errors.Errorf("expected names.MachineTag, got %T", tag)
		}
		var err error
		if p.machine, err = p.st.Machine(machineTag); err != nil {
			logger.Errorf("%s is not in state", machineTag)
			return nil, err
		}
	}
	return p.machine, nil
}

func (p *containerProvisioner) getMachineWatcher() (watcher.StringsWatcher, error) {
	machine, err := p.getMachine()
	if err != nil {
		return nil, err
	}
	return machine.WatchContainers(p.containerType)
}

func (p *containerProvisioner) getRetryWatcher() (watcher.NotifyWatcher, error) {
	return nil, errors.NotImplementedf("getRetryWatcher")
}
<commit_msg>Fix intermittent test bug #1552589.<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner

import (
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/names"

	"github.com/juju/juju/agent"
	apiprovisioner "github.com/juju/juju/api/provisioner"
	"github.com/juju/juju/controller/authentication"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/watcher"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/worker/catacomb"
	"github.com/juju/juju/worker/environ"
)

var logger = loggo.GetLogger("juju.provisioner")

// Ensure our structs implement the required Provisioner interface.
var _ Provisioner = (*environProvisioner)(nil)
var _ Provisioner = (*containerProvisioner)(nil)

var (
	retryStrategyDelay = 10 * time.Second
	retryStrategyCount = 3
)

// Provisioner represents a running provisioner worker.
type Provisioner interface {
	worker.Worker
	getMachineWatcher() (watcher.StringsWatcher, error)
	getRetryWatcher() (watcher.NotifyWatcher, error)
}

// environProvisioner represents a running provisioning worker for machine nodes
// belonging to an environment.
type environProvisioner struct {
	provisioner
	environ environs.Environ
	configObserver
}

// containerProvisioner represents a running provisioning worker for containers
// hosted on a machine.
type containerProvisioner struct {
	provisioner
	containerType instance.ContainerType
	machine       *apiprovisioner.Machine
	configObserver
}

// provisioner providers common behaviour for a running provisioning worker.
type provisioner struct {
	Provisioner
	st          *apiprovisioner.State
	agentConfig agent.Config
	broker      environs.InstanceBroker
	toolsFinder ToolsFinder
	catacomb    catacomb.Catacomb
}

// RetryStrategy defines the retry behavior when encountering a retryable
// error during provisioning.
type RetryStrategy struct {
	retryDelay time.Duration
	retryCount int
}

// NewRetryStrategy returns a new retry strategy with the specified delay and
// count for use with retryable provisioning errors.
func NewRetryStrategy(delay time.Duration, count int) RetryStrategy {
	return RetryStrategy{
		retryDelay: delay,
		retryCount: count,
	}
}

// configObserver is implemented so that tests can see
// when the environment configuration changes.
type configObserver struct {
	sync.Mutex
	observer chan<- *config.Config
}

// notify notifies the observer of a configuration change.
func (o *configObserver) notify(cfg *config.Config) {
	o.Lock()
	if o.observer != nil {
		o.observer <- cfg
	}
	o.Unlock()
}

// Kill implements worker.Worker.Kill.
func (p *provisioner) Kill() {
	p.catacomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (p *provisioner) Wait() error {
	return p.catacomb.Wait()
}

// getToolsFinder returns a ToolsFinder for the provided State.
// This exists for mocking.
var getToolsFinder = func(st *apiprovisioner.State) ToolsFinder {
	return st
}

// getStartTask creates a new worker for the provisioner,
func (p *provisioner) getStartTask(harvestMode config.HarvestMode) (ProvisionerTask, error) {
	auth, err := authentication.NewAPIAuthenticator(p.st)
	if err != nil {
		return nil, err
	}
	// Start responding to changes in machines, and to any further updates
	// to the environment config.
	machineWatcher, err := p.getMachineWatcher()
	if err != nil {
		return nil, err
	}
	retryWatcher, err := p.getRetryWatcher()
	if err != nil && !errors.IsNotImplemented(err) {
		return nil, err
	}
	tag := p.agentConfig.Tag()
	machineTag, ok := tag.(names.MachineTag)
	if !ok {
		errors.Errorf("expected names.MachineTag, got %T", tag)
	}

	envCfg, err := p.st.ModelConfig()
	if err != nil {
		return nil, errors.Annotate(err, "could not retrieve the model config.")
	}

	secureServerConnection := false
	if info, ok := p.agentConfig.StateServingInfo(); ok {
		secureServerConnection = info.CAPrivateKey != ""
	}
	task, err := NewProvisionerTask(
		machineTag,
		harvestMode,
		p.st,
		p.toolsFinder,
		machineWatcher,
		retryWatcher,
		p.broker,
		auth,
		envCfg.ImageStream(),
		secureServerConnection,
		RetryStrategy{retryDelay: retryStrategyDelay, retryCount: retryStrategyCount},
	)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return task, nil
}

// NewEnvironProvisioner returns a new Provisioner for an environment.
// When new machines are added to the state, it allocates instances
// from the environment and allocates them to the new machines.
func NewEnvironProvisioner(st *apiprovisioner.State, agentConfig agent.Config) (Provisioner, error) {
	p := &environProvisioner{
		provisioner: provisioner{
			st:          st,
			agentConfig: agentConfig,
			toolsFinder: getToolsFinder(st),
		},
	}
	p.Provisioner = p
	logger.Tracef("Starting environ provisioner for %q", p.agentConfig.Tag())

	err := catacomb.Invoke(catacomb.Plan{
		Site: &p.catacomb,
		Work: p.loop,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return p, nil
}

func (p *environProvisioner) loop() error {
	var modelConfigChanges <-chan struct{}
	modelWatcher, err := p.st.WatchForModelConfigChanges()
	if err != nil {
		return loggedErrorStack(errors.Trace(err))
	}
	if err := p.catacomb.Add(modelWatcher); err != nil {
		return errors.Trace(err)
	}
	modelConfigChanges = modelWatcher.Changes()

	p.environ, err = environ.WaitForEnviron(modelWatcher, p.st, p.catacomb.Dying())
	if err != nil {
		if err == environ.ErrWaitAborted {
			return p.catacomb.ErrDying()
		}
		return loggedErrorStack(errors.Trace(err))
	}
	p.broker = p.environ

	modelConfig := p.environ.Config()
	p.configObserver.notify(modelConfig)
	harvestMode := modelConfig.ProvisionerHarvestMode()
	task, err := p.getStartTask(harvestMode)
	if err != nil {
		return loggedErrorStack(errors.Trace(err))
	}
	if err := p.catacomb.Add(task); err != nil {
		return errors.Trace(err)
	}

	for {
		select {
		case <-p.catacomb.Dying():
			return p.catacomb.ErrDying()
		case _, ok := <-modelConfigChanges:
			if !ok {
				return errors.New("model configuration watcher closed")
			}
			modelConfig, err := p.st.ModelConfig()
			if err != nil {
				return errors.Annotate(err, "cannot load model configuration")
			}
			if err := p.setConfig(modelConfig); err != nil {
				return errors.Annotate(err, "loaded invalid model configuration")
			}
			task.SetHarvestMode(modelConfig.ProvisionerHarvestMode())
		}
	}
}

func (p *environProvisioner) getMachineWatcher() (watcher.StringsWatcher, error) {
	return p.st.WatchModelMachines()
}

func (p *environProvisioner) getRetryWatcher() (watcher.NotifyWatcher, error) {
	return p.st.WatchMachineErrorRetry()
}

// setConfig updates the environment configuration and notifies
// the config observer.
func (p *environProvisioner) setConfig(modelConfig *config.Config) error {
	if err := p.environ.SetConfig(modelConfig); err != nil {
		return err
	}
	p.configObserver.notify(modelConfig)
	return nil
}

// NewContainerProvisioner returns a new Provisioner. When new machines
// are added to the state, it allocates instances from the environment
// and allocates them to the new machines.
func NewContainerProvisioner(
	containerType instance.ContainerType,
	st *apiprovisioner.State,
	agentConfig agent.Config,
	broker environs.InstanceBroker,
	toolsFinder ToolsFinder,
) (Provisioner, error) {

	p := &containerProvisioner{
		provisioner: provisioner{
			st:          st,
			agentConfig: agentConfig,
			broker:      broker,
			toolsFinder: toolsFinder,
		},
		containerType: containerType,
	}
	p.Provisioner = p
	logger.Tracef("Starting %s provisioner for %q", p.containerType, p.agentConfig.Tag())

	err := catacomb.Invoke(catacomb.Plan{
		Site: &p.catacomb,
		Work: p.loop,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return p, nil
}

func (p *containerProvisioner) loop() error {
	modelWatcher, err := p.st.WatchForModelConfigChanges()
	if err != nil {
		return errors.Trace(err)
	}
	if err := p.catacomb.Add(modelWatcher); err != nil {
		return errors.Trace(err)
	}

	modelConfig, err := p.st.ModelConfig()
	if err != nil {
		return err
	}
	p.configObserver.notify(modelConfig)
	harvestMode := modelConfig.ProvisionerHarvestMode()

	task, err := p.getStartTask(harvestMode)
	if err != nil {
		return err
	}
	if err := p.catacomb.Add(task); err != nil {
		return errors.Trace(err)
	}

	for {
		select {
		case <-p.catacomb.Dying():
			return p.catacomb.ErrDying()
		case _, ok := <-modelWatcher.Changes():
			if !ok {
				return errors.New("model configuratioon watch closed")
			}
			modelConfig, err := p.st.ModelConfig()
			if err != nil {
				return errors.Annotate(err, "cannot load model configuration")
			}
			p.configObserver.notify(modelConfig)
			task.SetHarvestMode(modelConfig.ProvisionerHarvestMode())
		}
	}
}

func (p *containerProvisioner) getMachine() (*apiprovisioner.Machine, error) {
	if p.machine == nil {
		tag := p.agentConfig.Tag()
		machineTag, ok := tag.(names.MachineTag)
		if !ok {
			return nil, errors.Errorf("expected names.MachineTag, got %T", tag)
		}
		var err error
		if p.machine, err = p.st.Machine(machineTag); err != nil {
			logger.Errorf("%s is not in state", machineTag)
			return nil, err
		}
	}
	return p.machine, nil
}

func (p *containerProvisioner) getMachineWatcher() (watcher.StringsWatcher, error) {
	machine, err := p.getMachine()
	if err != nil {
		return nil, err
	}
	return machine.WatchContainers(p.containerType)
}

func (p *containerProvisioner) getRetryWatcher() (watcher.NotifyWatcher, error) {
	return nil, errors.NotImplementedf("getRetryWatcher")
}
<|endoftext|>