<commit_before>/*
Copyright 2014 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package rest

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"sort"
	"sync"

	"github.com/gorilla/mux"
)

const (
	defaultLogPrefix     = "rest "
	defaultDocsDirectory = "_docs/"
)

// Address is the address and port to bind to (e.g. ":8080").
type Address string

// FilePath represents a file path.
type FilePath string

// Configuration contains settings for configuring an API.
type Configuration struct {
	Debug         bool
	Logger        *log.Logger
	GenerateDocs  bool
	DocsDirectory string
}

// Debugf prints the formatted string to the Configuration Logger if Debug is enabled.
func (c *Configuration) Debugf(format string, v ...interface{}) {
	if c.Debug {
		c.Logger.Printf(format, v)
	}
}

// NewConfiguration returns a default Configuration.
func NewConfiguration() *Configuration {
	logger := log.New(os.Stdout, defaultLogPrefix, log.LstdFlags)
	return &Configuration{
		Debug:         true,
		Logger:        logger,
		GenerateDocs:  true,
		DocsDirectory: defaultDocsDirectory,
	}
}

// Middleware can be passed in to API#Start and API#StartTLS and will be
// invoked on every request to a route handled by the API. Returns true if the
// request should be terminated, false if it should continue.
type Middleware func(w http.ResponseWriter, r *http.Request) bool

// middlewareProxy proxies an http.Handler by invoking middleware before
// passing the request to the Handler. It implements the http.Handler
// interface.
type middlewareProxy struct {
	handler    http.Handler
	middleware []Middleware
}

// ServeHTTP invokes middleware on the request and then delegates to the
// proxied http.Handler.
func (m *middlewareProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	for _, middleware := range m.middleware {
		if middleware(w, r) {
			return
		}
	}
	m.handler.ServeHTTP(w, r)
}

// wrapMiddleware returns an http.Handler with the Middleware applied.
func wrapMiddleware(handler http.Handler, middleware ...Middleware) http.Handler {
	return &middlewareProxy{
		handler:    handler,
		middleware: middleware,
	}
}

// API is the top-level interface encapsulating an HTTP REST server. It's responsible for
// registering ResourceHandlers and routing requests. Use NewAPI to retrieve an instance.
type API interface {
	http.Handler

	// Start begins serving requests. This will block unless it fails, in which case an
	// error will be returned. This will validate any defined Rules. If any Rules are
	// invalid, it will panic. Any provided Middleware will be invoked for every request
	// handled by the API.
	Start(Address, ...Middleware) error

	// StartTLS begins serving requests received over HTTPS connections. This will block
	// unless it fails, in which case an error will be returned. Files containing a
	// certificate and matching private key for the server must be provided. If the
	// certificate is signed by a certificate authority, the certFile should be the
	// concatenation of the server's certificate followed by the CA's certificate. This
	// will validate any defined Rules. If any Rules are invalid, it will panic. Any
	// provided Middleware will be invoked for every request handled by the API.
	StartTLS(Address, FilePath, FilePath, ...Middleware) error

	// RegisterResourceHandler binds the provided ResourceHandler to the appropriate REST
	// endpoints and applies any specified middleware. Endpoints will have the following
	// base URL: /api/:version/resourceName.
	RegisterResourceHandler(ResourceHandler, ...RequestMiddleware)

	// RegisterHandlerFunc binds the http.HandlerFunc to the provided URI and applies any
	// specified middleware.
	RegisterHandlerFunc(string, http.HandlerFunc, ...RequestMiddleware)

	// RegisterHandler binds the http.Handler to the provided URI and applies any specified
	// middleware.
	RegisterHandler(string, http.Handler, ...RequestMiddleware)

	// RegisterPathPrefix binds the http.HandlerFunc to URIs matched by the given path
	// prefix and applies any specified middleware.
	RegisterPathPrefix(string, http.HandlerFunc, ...RequestMiddleware)

	// RegisterResponseSerializer registers the provided ResponseSerializer with the given
	// format. If the format has already been registered, it will be overwritten.
	RegisterResponseSerializer(string, ResponseSerializer)

	// UnregisterResponseSerializer unregisters the ResponseSerializer with the provided
	// format. If the format hasn't been registered, this is a no-op.
	UnregisterResponseSerializer(string)

	// AvailableFormats returns a slice containing all of the available serialization
	// formats currently available.
	AvailableFormats() []string

	// Configuration returns the API Configuration.
	Configuration() *Configuration

	// ResourceHandlers returns a slice containing the registered ResourceHandlers.
	ResourceHandlers() []ResourceHandler

	// Validate will validate the Rules configured for this API. It returns nil
	// if all Rules are valid, otherwise returns the first encountered
	// validation error.
	Validate() error

	// responseSerializer returns a ResponseSerializer for the given format type. If the
	// format is not implemented, the returned serializer will be nil and the error set.
	responseSerializer(string) (ResponseSerializer, error)
}

// RequestMiddleware is a function that returns a HandlerFunc wrapping the provided HandlerFunc.
// This allows injecting custom logic to operate on requests (e.g. performing authentication).
type RequestMiddleware func(http.HandlerFunc) http.HandlerFunc

// newAuthMiddleware returns a RequestMiddleware used to authenticate requests.
func newAuthMiddleware(authenticate func(*http.Request) error) RequestMiddleware {
	return func(wrapped http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			if err := authenticate(r); err != nil {
				w.WriteHeader(http.StatusUnauthorized)
				w.Write([]byte(err.Error()))
				return
			}
			wrapped(w, r)
		}
	}
}

// muxAPI is an implementation of the API interface which relies on the gorilla/mux
// package to handle request dispatching (see http://www.gorillatoolkit.org/pkg/mux).
type muxAPI struct {
	config             *Configuration
	router             *mux.Router
	mu                 sync.RWMutex
	handler            *requestHandler
	serializerRegistry map[string]ResponseSerializer
	resourceHandlers   []ResourceHandler
}

// NewAPI returns a newly allocated API instance.
func NewAPI(config *Configuration) API {
	r := mux.NewRouter()
	restAPI := &muxAPI{
		config:             config,
		router:             r,
		serializerRegistry: map[string]ResponseSerializer{"json": &jsonSerializer{}},
		resourceHandlers:   make([]ResourceHandler, 0),
	}
	restAPI.handler = &requestHandler{restAPI}
	return restAPI
}

// Start begins serving requests. This will block unless it fails, in which case an error will be
// returned.
func (r *muxAPI) Start(addr Address, middleware ...Middleware) error {
	r.preprocess()
	return http.ListenAndServe(string(addr), wrapMiddleware(r.router, middleware...))
}

// StartTLS begins serving requests received over HTTPS connections. This will block unless it
// fails, in which case an error will be returned. Files containing a certificate and matching
// private key for the server must be provided. If the certificate is signed by a certificate
// authority, the certFile should be the concatenation of the server's certificate followed by
// the CA's certificate.
func (r *muxAPI) StartTLS(addr Address, certFile, keyFile FilePath, middleware ...Middleware) error {
	r.preprocess()
	return http.ListenAndServeTLS(string(addr), string(certFile), string(keyFile), wrapMiddleware(r.router, middleware...))
}

// preprocess performs any necessary preprocessing before the server can be started, including
// Rule validation.
func (r *muxAPI) preprocess() {
	r.validateRulesOrPanic()
	if r.config.GenerateDocs {
		if err := newDocGenerator().generateDocs(r); err != nil {
			log.Printf("documentation could not be generated: %v", err)
		}
	}
}

// RegisterResourceHandler binds the provided ResourceHandler to the appropriate REST endpoints and
// applies any specified middleware. Endpoints will have the following base URL:
// /api/:version/resourceName.
func (r *muxAPI) RegisterResourceHandler(h ResourceHandler, middleware ...RequestMiddleware) {
	h = resourceHandlerProxy{h}
	resource := h.ResourceName()
	middleware = append(middleware, newAuthMiddleware(h.Authenticate))

	r.router.HandleFunc(
		h.CreateURI(), applyMiddleware(r.handler.handleCreate(h), middleware),
	).Methods("POST").Name(resource + ":create")
	r.config.Debugf("Registered create handler at POST %s", h.CreateURI())

	r.router.HandleFunc(
		h.ReadListURI(), applyMiddleware(r.handler.handleReadList(h), middleware),
	).Methods("GET").Name(resource + ":readList")
	r.config.Debugf("Registered read list handler at GET %s", h.ReadListURI())

	r.router.HandleFunc(
		h.ReadURI(), applyMiddleware(r.handler.handleRead(h), middleware),
	).Methods("GET").Name(resource + ":read")
	r.config.Debugf("Registered read handler at GET %s", h.ReadURI())

	r.router.HandleFunc(
		h.UpdateListURI(), applyMiddleware(r.handler.handleUpdateList(h), middleware),
	).Methods("PUT").Name(resource + ":updateList")
	r.config.Debugf("Registered update list handler at PUT %s", h.UpdateListURI())

	r.router.HandleFunc(
		h.UpdateURI(), applyMiddleware(r.handler.handleUpdate(h), middleware),
	).Methods("PUT").Name(resource + ":update")
	r.config.Debugf("Registered update handler at PUT %s", h.UpdateURI())

	r.router.HandleFunc(
		h.DeleteURI(), applyMiddleware(r.handler.handleDelete(h), middleware),
	).Methods("DELETE").Name(resource + ":delete")
	r.config.Debugf("Registered delete handler at DELETE %s", h.DeleteURI())

	// Some browsers don't support PUT and DELETE, so allow method overriding.
	// POST requests with X-HTTP-Method-Override=PUT/DELETE will route to the
	// respective handlers.
	r.router.HandleFunc(
		h.UpdateListURI(), applyMiddleware(r.handler.handleUpdateList(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "PUT").Name(resource + ":updateListOverride")

	r.router.HandleFunc(
		h.UpdateURI(), applyMiddleware(r.handler.handleUpdate(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "PUT").Name(resource + ":updateOverride")

	r.router.HandleFunc(
		h.DeleteURI(), applyMiddleware(r.handler.handleDelete(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "DELETE").Name(resource + ":deleteOverride")

	r.resourceHandlers = append(r.resourceHandlers, h)
}

// RegisterHandlerFunc binds the http.HandlerFunc to the provided URI and applies any
// specified middleware.
func (r *muxAPI) RegisterHandlerFunc(uri string, handler http.HandlerFunc,
	middleware ...RequestMiddleware) {
	r.router.HandleFunc(uri, applyMiddleware(handler, middleware))
}

// RegisterHandler binds the http.Handler to the provided URI and applies any specified
// middleware.
func (r *muxAPI) RegisterHandler(uri string, handler http.Handler, middleware ...RequestMiddleware) {
	r.router.HandleFunc(uri, applyMiddleware(handler.ServeHTTP, middleware))
}

// RegisterPathPrefix binds the http.HandlerFunc to URIs matched by the given path
// prefix and applies any specified middleware.
func (r *muxAPI) RegisterPathPrefix(uri string, handler http.HandlerFunc,
	middleware ...RequestMiddleware) {
	r.router.PathPrefix(uri).HandlerFunc(applyMiddleware(handler, middleware))
}

// ServeHTTP handles an HTTP request.
func (r *muxAPI) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	r.router.ServeHTTP(w, req)
}

// RegisterResponseSerializer registers the provided ResponseSerializer with the given format. If the
// format has already been registered, it will be overwritten.
func (r *muxAPI) RegisterResponseSerializer(format string, serializer ResponseSerializer) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.serializerRegistry[format] = serializer
}

// UnregisterResponseSerializer unregisters the ResponseSerializer with the provided format. If the
// format hasn't been registered, this is a no-op.
func (r *muxAPI) UnregisterResponseSerializer(format string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.serializerRegistry, format)
}

// AvailableFormats returns a slice containing all of the available serialization formats
// currently available.
func (r *muxAPI) AvailableFormats() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	formats := make([]string, 0, len(r.serializerRegistry))
	for format := range r.serializerRegistry {
		formats = append(formats, format)
	}
	sort.Strings(formats)
	return formats
}

// ResourceHandlers returns a slice containing the registered ResourceHandlers.
func (r *muxAPI) ResourceHandlers() []ResourceHandler {
	return r.resourceHandlers
}

// Configuration returns the API Configuration.
func (r *muxAPI) Configuration() *Configuration {
	return r.config
}

// Validate will validate the Rules configured for this API. It returns nil if
// all Rules are valid, otherwise returns the first encountered validation
// error.
func (r *muxAPI) Validate() error {
	for _, handler := range r.resourceHandlers {
		rules := handler.Rules()
		if rules == nil || rules.Size() == 0 {
			continue
		}

		if err := rules.Validate(); err != nil {
			return err
		}
	}
	return nil
}

// validateRulesOrPanic verifies that the Rules for each ResourceHandler
// registered with the muxAPI are valid, meaning they specify fields that exist
// and correct types. If a Rule is invalid, this will panic.
func (r *muxAPI) validateRulesOrPanic() {
	if err := r.Validate(); err != nil {
		panic(err)
	}
}

// responseSerializer returns a ResponseSerializer for the given format type. If the format
// is not implemented, the returned serializer will be nil and the error set.
func (r *muxAPI) responseSerializer(format string) (ResponseSerializer, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if serializer, ok := r.serializerRegistry[format]; ok {
		return serializer, nil
	}
	return nil, fmt.Errorf("Format not implemented: %s", format)
}

// applyMiddleware wraps the HandlerFunc with the provided RequestMiddleware and returns the
// function composition.
func applyMiddleware(h http.HandlerFunc, middleware []RequestMiddleware) http.HandlerFunc {
	for _, m := range middleware {
		h = m(h)
	}

	return h
}
<commit_msg>Move overrides up so they are recognized<commit_after>/*
Copyright 2014 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package rest

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"sort"
	"sync"

	"github.com/gorilla/mux"
)

const (
	defaultLogPrefix     = "rest "
	defaultDocsDirectory = "_docs/"
)

// Address is the address and port to bind to (e.g. ":8080").
type Address string

// FilePath represents a file path.
type FilePath string

// Configuration contains settings for configuring an API.
type Configuration struct {
	Debug         bool
	Logger        *log.Logger
	GenerateDocs  bool
	DocsDirectory string
}

// Debugf prints the formatted string to the Configuration Logger if Debug is enabled.
func (c *Configuration) Debugf(format string, v ...interface{}) {
	if c.Debug {
		c.Logger.Printf(format, v)
	}
}

// NewConfiguration returns a default Configuration.
func NewConfiguration() *Configuration {
	logger := log.New(os.Stdout, defaultLogPrefix, log.LstdFlags)
	return &Configuration{
		Debug:         true,
		Logger:        logger,
		GenerateDocs:  true,
		DocsDirectory: defaultDocsDirectory,
	}
}

// Middleware can be passed in to API#Start and API#StartTLS and will be
// invoked on every request to a route handled by the API. Returns true if the
// request should be terminated, false if it should continue.
type Middleware func(w http.ResponseWriter, r *http.Request) bool

// middlewareProxy proxies an http.Handler by invoking middleware before
// passing the request to the Handler. It implements the http.Handler
// interface.
type middlewareProxy struct {
	handler    http.Handler
	middleware []Middleware
}

// ServeHTTP invokes middleware on the request and then delegates to the
// proxied http.Handler.
func (m *middlewareProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	for _, middleware := range m.middleware {
		if middleware(w, r) {
			return
		}
	}
	m.handler.ServeHTTP(w, r)
}

// wrapMiddleware returns an http.Handler with the Middleware applied.
func wrapMiddleware(handler http.Handler, middleware ...Middleware) http.Handler {
	return &middlewareProxy{
		handler:    handler,
		middleware: middleware,
	}
}

// API is the top-level interface encapsulating an HTTP REST server. It's responsible for
// registering ResourceHandlers and routing requests. Use NewAPI to retrieve an instance.
type API interface {
	http.Handler

	// Start begins serving requests. This will block unless it fails, in which case an
	// error will be returned. This will validate any defined Rules. If any Rules are
	// invalid, it will panic. Any provided Middleware will be invoked for every request
	// handled by the API.
	Start(Address, ...Middleware) error

	// StartTLS begins serving requests received over HTTPS connections. This will block
	// unless it fails, in which case an error will be returned. Files containing a
	// certificate and matching private key for the server must be provided. If the
	// certificate is signed by a certificate authority, the certFile should be the
	// concatenation of the server's certificate followed by the CA's certificate. This
	// will validate any defined Rules. If any Rules are invalid, it will panic. Any
	// provided Middleware will be invoked for every request handled by the API.
	StartTLS(Address, FilePath, FilePath, ...Middleware) error

	// RegisterResourceHandler binds the provided ResourceHandler to the appropriate REST
	// endpoints and applies any specified middleware. Endpoints will have the following
	// base URL: /api/:version/resourceName.
	RegisterResourceHandler(ResourceHandler, ...RequestMiddleware)

	// RegisterHandlerFunc binds the http.HandlerFunc to the provided URI and applies any
	// specified middleware.
	RegisterHandlerFunc(string, http.HandlerFunc, ...RequestMiddleware)

	// RegisterHandler binds the http.Handler to the provided URI and applies any specified
	// middleware.
	RegisterHandler(string, http.Handler, ...RequestMiddleware)

	// RegisterPathPrefix binds the http.HandlerFunc to URIs matched by the given path
	// prefix and applies any specified middleware.
	RegisterPathPrefix(string, http.HandlerFunc, ...RequestMiddleware)

	// RegisterResponseSerializer registers the provided ResponseSerializer with the given
	// format. If the format has already been registered, it will be overwritten.
	RegisterResponseSerializer(string, ResponseSerializer)

	// UnregisterResponseSerializer unregisters the ResponseSerializer with the provided
	// format. If the format hasn't been registered, this is a no-op.
	UnregisterResponseSerializer(string)

	// AvailableFormats returns a slice containing all of the available serialization
	// formats currently available.
	AvailableFormats() []string

	// Configuration returns the API Configuration.
	Configuration() *Configuration

	// ResourceHandlers returns a slice containing the registered ResourceHandlers.
	ResourceHandlers() []ResourceHandler

	// Validate will validate the Rules configured for this API. It returns nil
	// if all Rules are valid, otherwise returns the first encountered
	// validation error.
	Validate() error

	// responseSerializer returns a ResponseSerializer for the given format type. If the
	// format is not implemented, the returned serializer will be nil and the error set.
	responseSerializer(string) (ResponseSerializer, error)
}

// RequestMiddleware is a function that returns a HandlerFunc wrapping the provided HandlerFunc.
// This allows injecting custom logic to operate on requests (e.g. performing authentication).
type RequestMiddleware func(http.HandlerFunc) http.HandlerFunc

// newAuthMiddleware returns a RequestMiddleware used to authenticate requests.
func newAuthMiddleware(authenticate func(*http.Request) error) RequestMiddleware {
	return func(wrapped http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			if err := authenticate(r); err != nil {
				w.WriteHeader(http.StatusUnauthorized)
				w.Write([]byte(err.Error()))
				return
			}
			wrapped(w, r)
		}
	}
}

// muxAPI is an implementation of the API interface which relies on the gorilla/mux
// package to handle request dispatching (see http://www.gorillatoolkit.org/pkg/mux).
type muxAPI struct {
	config             *Configuration
	router             *mux.Router
	mu                 sync.RWMutex
	handler            *requestHandler
	serializerRegistry map[string]ResponseSerializer
	resourceHandlers   []ResourceHandler
}

// NewAPI returns a newly allocated API instance.
func NewAPI(config *Configuration) API {
	r := mux.NewRouter()
	restAPI := &muxAPI{
		config:             config,
		router:             r,
		serializerRegistry: map[string]ResponseSerializer{"json": &jsonSerializer{}},
		resourceHandlers:   make([]ResourceHandler, 0),
	}
	restAPI.handler = &requestHandler{restAPI}
	return restAPI
}

// Start begins serving requests. This will block unless it fails, in which case an error will be
// returned.
func (r *muxAPI) Start(addr Address, middleware ...Middleware) error {
	r.preprocess()
	return http.ListenAndServe(string(addr), wrapMiddleware(r.router, middleware...))
}

// StartTLS begins serving requests received over HTTPS connections. This will block unless it
// fails, in which case an error will be returned. Files containing a certificate and matching
// private key for the server must be provided. If the certificate is signed by a certificate
// authority, the certFile should be the concatenation of the server's certificate followed by
// the CA's certificate.
func (r *muxAPI) StartTLS(addr Address, certFile, keyFile FilePath, middleware ...Middleware) error {
	r.preprocess()
	return http.ListenAndServeTLS(string(addr), string(certFile), string(keyFile), wrapMiddleware(r.router, middleware...))
}

// preprocess performs any necessary preprocessing before the server can be started, including
// Rule validation.
func (r *muxAPI) preprocess() {
	r.validateRulesOrPanic()
	if r.config.GenerateDocs {
		if err := newDocGenerator().generateDocs(r); err != nil {
			log.Printf("documentation could not be generated: %v", err)
		}
	}
}

// Check the route for an error and log the error if it exists.
func (r *muxAPI) checkRoute(method, uri string, route *mux.Route) {
	err := route.GetError()

	if err != nil {
		log.Printf("Failed to setup route %s with %v", uri, err)
	} else {
		r.config.Debugf("Registered create handler at %s %s", method, uri)
	}
}

// RegisterResourceHandler binds the provided ResourceHandler to the appropriate REST endpoints and
// applies any specified middleware. Endpoints will have the following base URL:
// /api/:version/resourceName.
func (r *muxAPI) RegisterResourceHandler(h ResourceHandler, middleware ...RequestMiddleware) {
	h = resourceHandlerProxy{h}
	resource := h.ResourceName()
	middleware = append(middleware, newAuthMiddleware(h.Authenticate))

	// Some browsers don't support PUT and DELETE, so allow method overriding.
	// POST requests with X-HTTP-Method-Override=PUT/DELETE will route to the
	// respective handlers.

	route := r.router.HandleFunc(
		h.ReadListURI(), applyMiddleware(r.handler.handleReadList(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "GET").Name(resource + ":readListOverride")
	r.checkRoute(h.ReadListURI(), "OVERRIDE-GET", route)

	route = r.router.HandleFunc(
		h.UpdateListURI(), applyMiddleware(r.handler.handleUpdateList(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "PUT").Name(resource + ":updateListOverride")
	r.checkRoute(h.UpdateListURI(), "OVERRIDE-PUT", route)

	route = r.router.HandleFunc(
		h.UpdateURI(), applyMiddleware(r.handler.handleUpdate(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "PUT").Name(resource + ":updateOverride")
	r.checkRoute(h.UpdateURI(), "OVERRIDE-PUT", route)

	route = r.router.HandleFunc(
		h.DeleteURI(), applyMiddleware(r.handler.handleDelete(h), middleware),
	).Methods("POST").Headers("X-HTTP-Method-Override", "DELETE").Name(resource + ":deleteOverride")
	r.checkRoute(h.DeleteURI(), "OVERRIDE-DELETE", route)

	// These return a Route which has a GetError command. Probably should check
	// that and log it if it fails :)
	r.router.HandleFunc(
		h.CreateURI(), applyMiddleware(r.handler.handleCreate(h), middleware),
	).Methods("POST").Name(resource + ":create")
	r.checkRoute(h.CreateURI(), "POST", route)

	r.router.HandleFunc(
		h.ReadListURI(), applyMiddleware(r.handler.handleReadList(h), middleware),
	).Methods("GET").Name(resource + ":readList")
	r.checkRoute(h.ReadListURI(), "GET", route)

	r.router.HandleFunc(
		h.ReadURI(), applyMiddleware(r.handler.handleRead(h), middleware),
	).Methods("GET").Name(resource + ":read")
	r.checkRoute(h.ReadURI(), "GET", route)

	r.router.HandleFunc(
		h.UpdateListURI(), applyMiddleware(r.handler.handleUpdateList(h), middleware),
	).Methods("PUT").Name(resource + ":updateList")
	r.checkRoute(h.UpdateListURI(), "PUT", route)

	r.router.HandleFunc(
		h.UpdateURI(), applyMiddleware(r.handler.handleUpdate(h), middleware),
	).Methods("PUT").Name(resource + ":update")
	r.checkRoute(h.UpdateURI(), "PUT", route)

	r.router.HandleFunc(
		h.DeleteURI(), applyMiddleware(r.handler.handleDelete(h), middleware),
	).Methods("DELETE").Name(resource + ":delete")
	r.checkRoute(h.DeleteURI(), "DELETE", route)

	r.resourceHandlers = append(r.resourceHandlers, h)
}

// RegisterHandlerFunc binds the http.HandlerFunc to the provided URI and applies any
// specified middleware.
func (r *muxAPI) RegisterHandlerFunc(uri string, handler http.HandlerFunc,
	middleware ...RequestMiddleware) {
	r.router.HandleFunc(uri, applyMiddleware(handler, middleware))
}

// RegisterHandler binds the http.Handler to the provided URI and applies any specified
// middleware.
func (r *muxAPI) RegisterHandler(uri string, handler http.Handler, middleware ...RequestMiddleware) {
	r.router.HandleFunc(uri, applyMiddleware(handler.ServeHTTP, middleware))
}

// RegisterPathPrefix binds the http.HandlerFunc to URIs matched by the given path
// prefix and applies any specified middleware.
func (r *muxAPI) RegisterPathPrefix(uri string, handler http.HandlerFunc,
	middleware ...RequestMiddleware) {
	r.router.PathPrefix(uri).HandlerFunc(applyMiddleware(handler, middleware))
}

// ServeHTTP handles an HTTP request.
func (r *muxAPI) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	r.router.ServeHTTP(w, req)
}

// RegisterResponseSerializer registers the provided ResponseSerializer with the given format. If the
// format has already been registered, it will be overwritten.
func (r *muxAPI) RegisterResponseSerializer(format string, serializer ResponseSerializer) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.serializerRegistry[format] = serializer
}

// UnregisterResponseSerializer unregisters the ResponseSerializer with the provided format. If the
// format hasn't been registered, this is a no-op.
func (r *muxAPI) UnregisterResponseSerializer(format string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.serializerRegistry, format)
}

// AvailableFormats returns a slice containing all of the available serialization formats
// currently available.
func (r *muxAPI) AvailableFormats() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	formats := make([]string, 0, len(r.serializerRegistry))
	for format := range r.serializerRegistry {
		formats = append(formats, format)
	}
	sort.Strings(formats)
	return formats
}

// ResourceHandlers returns a slice containing the registered ResourceHandlers.
func (r *muxAPI) ResourceHandlers() []ResourceHandler {
	return r.resourceHandlers
}

// Configuration returns the API Configuration.
func (r *muxAPI) Configuration() *Configuration {
	return r.config
}

// Validate will validate the Rules configured for this API. It returns nil if
// all Rules are valid, otherwise returns the first encountered validation
// error.
func (r *muxAPI) Validate() error {
	for _, handler := range r.resourceHandlers {
		rules := handler.Rules()
		if rules == nil || rules.Size() == 0 {
			continue
		}

		if err := rules.Validate(); err != nil {
			return err
		}
	}
	return nil
}

// validateRulesOrPanic verifies that the Rules for each ResourceHandler
// registered with the muxAPI are valid, meaning they specify fields that exist
// and correct types. If a Rule is invalid, this will panic.
func (r *muxAPI) validateRulesOrPanic() {
	if err := r.Validate(); err != nil {
		panic(err)
	}
}

// responseSerializer returns a ResponseSerializer for the given format type. If the format
// is not implemented, the returned serializer will be nil and the error set.
func (r *muxAPI) responseSerializer(format string) (ResponseSerializer, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if serializer, ok := r.serializerRegistry[format]; ok {
		return serializer, nil
	}
	return nil, fmt.Errorf("Format not implemented: %s", format)
}

// applyMiddleware wraps the HandlerFunc with the provided RequestMiddleware and returns the
// function composition.
func applyMiddleware(h http.HandlerFunc, middleware []RequestMiddleware) http.HandlerFunc {
	for _, m := range middleware {
		h = m(h)
	}

	return h
}
<|endoftext|>


<commit_before>// Upload object to QingStor

// +build !plan9

package qingstor

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"hash"
	"io"
	"sort"
	"sync"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	qs "github.com/yunify/qingstor-sdk-go/v3/service"
)

const (
	// maxSinglePartSize = 1024 * 1024 * 1024 * 5 // The maximum allowed size when uploading a single object to QingStor
	// maxMultiPartSize = 1024 * 1024 * 1024 * 1 // The maximum allowed part size when uploading a part to QingStor
	minMultiPartSize = 1024 * 1024 * 4 // The minimum allowed part size when uploading a part to QingStor
	maxMultiParts    = 10000           // The maximum allowed number of parts in an multi-part upload
)

const (
	defaultUploadPartSize    = 1024 * 1024 * 64 // The default part size to buffer chunks of a payload into.
	defaultUploadConcurrency = 4                // the default number of goroutines to spin up when using multiPartUpload.
)

func readFillBuf(r io.Reader, b []byte) (offset int, err error) {
	for offset < len(b) && err == nil {
		var n int
		n, err = r.Read(b[offset:])
		offset += n
	}

	return offset, err
}

// uploadInput contains all input for upload requests to QingStor.
type uploadInput struct {
	body           io.Reader
	qsSvc          *qs.Service
	mimeType       string
	zone           string
	bucket         string
	key            string
	partSize       int64
	concurrency    int
	maxUploadParts int
}

// uploader internal structure to manage an upload to QingStor.
type uploader struct {
	cfg        *uploadInput
	totalSize  int64 // set to -1 if the size is not known
	readerPos  int64 // current reader position
	readerSize int64 // current reader content size
}

// newUploader creates a new Uploader instance to upload objects to QingStor.
func newUploader(in *uploadInput) *uploader {
	u := &uploader{
		cfg: in,
	}
	return u
}

// bucketInit initiate as bucket controller
func (u *uploader) bucketInit() (*qs.Bucket, error) {
	bucketInit, err := u.cfg.qsSvc.Bucket(u.cfg.bucket, u.cfg.zone)
	return bucketInit, err
}

// String converts uploader to a string
func (u *uploader) String() string {
	return fmt.Sprintf("QingStor bucket %s key %s", u.cfg.bucket, u.cfg.key)
}

// nextReader returns a seekable reader representing the next packet of data.
// This operation increases the shared u.readerPos counter, but note that it
// does not need to be wrapped in a mutex because nextReader is only called
// from the main thread.
func (u *uploader) nextReader() (io.ReadSeeker, int, error) {
	type readerAtSeeker interface {
		io.ReaderAt
		io.ReadSeeker
	}
	switch r := u.cfg.body.(type) {
	case readerAtSeeker:
		var err error
		n := u.cfg.partSize
		if u.totalSize >= 0 {
			bytesLeft := u.totalSize - u.readerPos

			if bytesLeft <= u.cfg.partSize {
				err = io.EOF
				n = bytesLeft
			}
		}
		reader := io.NewSectionReader(r, u.readerPos, n)
		u.readerPos += n
		u.readerSize = n
		return reader, int(n), err

	default:
		part := make([]byte, u.cfg.partSize)
		n, err := readFillBuf(r, part)
		u.readerPos += int64(n)
		u.readerSize = int64(n)
		return bytes.NewReader(part[0:n]), n, err
	}
}

// init will initialize all default options.
func (u *uploader) init() {
	if u.cfg.concurrency == 0 {
		u.cfg.concurrency = defaultUploadConcurrency
	}
	if u.cfg.partSize == 0 {
		u.cfg.partSize = defaultUploadPartSize
	}
	if u.cfg.maxUploadParts == 0 {
		u.cfg.maxUploadParts = maxMultiParts
	}
	// Try to get the total size for some optimizations
	u.totalSize = -1
	switch r := u.cfg.body.(type) {
	case io.Seeker:
		pos, _ := r.Seek(0, io.SeekCurrent)
		defer func() {
			_, _ = r.Seek(pos, io.SeekStart)
		}()

		n, err := r.Seek(0, io.SeekEnd)
		if err != nil {
			return
		}
		u.totalSize = n

		// Try to adjust partSize if it is too small and account for
		// integer division truncation.
		if u.totalSize/u.cfg.partSize >= u.cfg.partSize {
			// Add one to the part size to account for remainders
			// during the size calculation. e.g odd number of bytes.
			u.cfg.partSize = (u.totalSize / int64(u.cfg.maxUploadParts)) + 1
		}
	}
}

// singlePartUpload upload a single object that contentLength less than "defaultUploadPartSize"
func (u *uploader) singlePartUpload(buf io.Reader, size int64) error {
	bucketInit, _ := u.bucketInit()

	req := qs.PutObjectInput{
		ContentLength: &size,
		ContentType:   &u.cfg.mimeType,
		Body:          buf,
	}

	_, err := bucketInit.PutObject(u.cfg.key, &req)
	if err == nil {
		fs.Debugf(u, "Upload single object finished")
	}
	return err
}

// Upload upload a object into QingStor
func (u *uploader) upload() error {
	u.init()

	if u.cfg.partSize < minMultiPartSize {
		return errors.Errorf("part size must be at least %d bytes", minMultiPartSize)
	}

	// Do one read to determine if we have more than one part
	reader, _, err := u.nextReader()
	if err == io.EOF { // single part
		fs.Debugf(u, "Uploading as single part object to QingStor")
		return u.singlePartUpload(reader, u.readerPos)
	} else if err != nil {
		return errors.Errorf("read upload data failed: %s", err)
	}

	fs.Debugf(u, "Uploading as multi-part object to QingStor")
	mu := multiUploader{uploader: u}
	return mu.multiPartUpload(reader)
}

// internal structure to manage a specific multipart upload to QingStor.
type multiUploader struct {
	*uploader
	wg          sync.WaitGroup
	mtx         sync.Mutex
	err         error
	uploadID    *string
	objectParts completedParts
	hashMd5     hash.Hash
}

// keeps track of a single chunk of data being sent to QingStor.
type chunk struct {
	buffer     io.ReadSeeker
	partNumber int
	size       int64
}

// completedParts is a wrapper to make parts sortable by their part number,
// since QingStor required this list to be sent in sorted order.
type completedParts []*qs.ObjectPartType

func (a completedParts) Len() int           { return len(a) }
func (a completedParts) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a completedParts) Less(i, j int) bool { return *a[i].PartNumber < *a[j].PartNumber }

// String converts multiUploader to a string
func (mu *multiUploader) String() string {
	if uploadID := mu.uploadID; uploadID != nil {
		return fmt.Sprintf("QingStor bucket %s key %s uploadID %s", mu.cfg.bucket, mu.cfg.key, *uploadID)
	}
	return fmt.Sprintf("QingStor bucket %s key %s uploadID <nil>", mu.cfg.bucket, mu.cfg.key)
}

// getErr is a thread-safe getter for the error object
func (mu *multiUploader) getErr() error {
	mu.mtx.Lock()
	defer mu.mtx.Unlock()
	return mu.err
}

// setErr is a thread-safe setter for the error object
func (mu *multiUploader) setErr(e error) {
	mu.mtx.Lock()
	defer mu.mtx.Unlock()
	mu.err = e
}

// readChunk runs in worker goroutines to pull chunks off of the ch channel
// and send() them as UploadPart requests.
func (mu *multiUploader) readChunk(ch chan chunk) {
	defer mu.wg.Done()
	for {
		c, ok := <-ch
		if !ok {
			break
		}
		if mu.getErr() == nil {
			if err := mu.send(c); err != nil {
				mu.setErr(err)
			}
		}
	}
}

// initiate init an Multiple Object and obtain UploadID
func (mu *multiUploader) initiate() error {
	bucketInit, _ := mu.bucketInit()
	req := qs.InitiateMultipartUploadInput{
		ContentType: &mu.cfg.mimeType,
	}
	fs.Debugf(mu, "Initiating a multi-part upload")
	rsp, err := bucketInit.InitiateMultipartUpload(mu.cfg.key, &req)
	if err == nil {
		mu.uploadID = rsp.UploadID
		mu.hashMd5 = md5.New()
	}
	return err
}

// send upload a part into QingStor
func (mu *multiUploader) send(c chunk) error {
	bucketInit, _ := mu.bucketInit()
	req := qs.UploadMultipartInput{
		PartNumber:    &c.partNumber,
		UploadID:      mu.uploadID,
		ContentLength: &c.size,
		Body:          c.buffer,
	}
	fs.Debugf(mu, "Uploading a part to QingStor with partNumber %d and partSize %d", c.partNumber, c.size)
	_, err := bucketInit.UploadMultipart(mu.cfg.key, &req)
	if err != nil {
		return err
	}
	fs.Debugf(mu, "Done uploading part partNumber %d and partSize %d", c.partNumber, c.size)

	mu.mtx.Lock()
	defer mu.mtx.Unlock()

	_, _ = c.buffer.Seek(0, 0)
	_, _ = io.Copy(mu.hashMd5, c.buffer)

	parts := qs.ObjectPartType{PartNumber: &c.partNumber, Size: &c.size}
	mu.objectParts = append(mu.objectParts, &parts)
	return err
}

// list list the ObjectParts of an multipart upload
func (mu *multiUploader) list() error {
	bucketInit, _ := mu.bucketInit()

	req := qs.ListMultipartInput{
		UploadID: mu.uploadID,
	}
	fs.Debugf(mu, "Reading multi-part details")
	rsp, err := bucketInit.ListMultipart(mu.cfg.key, &req)
	if err == nil {
		mu.objectParts = rsp.ObjectParts
	}
	return err
}

// complete complete an multipart upload
func (mu *multiUploader) complete() error {
	var err error
	if err = mu.getErr(); err != nil {
		return err
	}
	bucketInit, _ := mu.bucketInit()
	//if err = mu.list(); err != nil {
	//	return err
	//}
	//md5String := fmt.Sprintf("\"%s\"", hex.EncodeToString(mu.hashMd5.Sum(nil)))

	md5String := fmt.Sprintf("\"%x\"", mu.hashMd5.Sum(nil))
	sort.Sort(mu.objectParts)
	req := qs.CompleteMultipartUploadInput{
		UploadID:    mu.uploadID,
		ObjectParts: mu.objectParts,
		ETag:        &md5String,
	}
	fs.Debugf(mu, "Completing multi-part object")
	_, err = bucketInit.CompleteMultipartUpload(mu.cfg.key, &req)
	if err == nil {
		fs.Debugf(mu, "Complete multi-part finished")
	}
	return err
}

// abort abort an multipart upload
func (mu *multiUploader) abort() error {
	var err error
	bucketInit, _ := mu.bucketInit()

	if uploadID := mu.uploadID; uploadID != nil {
		req := qs.AbortMultipartUploadInput{
			UploadID: uploadID,
		}
		fs.Debugf(mu, "Aborting multi-part object %q", *uploadID)
		_, err = bucketInit.AbortMultipartUpload(mu.cfg.key, &req)
	}

	return err
}

// multiPartUpload upload a multiple object into QingStor
func (mu *multiUploader) multiPartUpload(firstBuf io.ReadSeeker) error {
	var err error
	//Initiate an multi-part upload
	if err = mu.initiate(); err != nil {
		return err
	}

	ch := make(chan chunk, mu.cfg.concurrency)
	for i := 0; i < mu.cfg.concurrency; i++ {
		mu.wg.Add(1)
		go mu.readChunk(ch)
	}

	var partNumber int
	ch <- chunk{partNumber: partNumber, buffer: firstBuf, size: mu.readerSize}

	for mu.getErr() == nil {
		partNumber++
		// This upload exceeded maximum number of supported parts, error now.
		if partNumber > mu.cfg.maxUploadParts || partNumber > maxMultiParts {
			var msg string
			if partNumber > mu.cfg.maxUploadParts {
				msg = fmt.Sprintf("exceeded total allowed configured maxUploadParts (%d). "+
					"Adjust PartSize to fit in this limit", mu.cfg.maxUploadParts)
			} else {
				msg = fmt.Sprintf("exceeded total allowed QingStor limit maxUploadParts (%d). "+
					"Adjust PartSize to fit in this limit", maxMultiParts)
			}
			mu.setErr(errors.New(msg))
			break
		}

		var reader io.ReadSeeker
		var nextChunkLen int
		reader, nextChunkLen, err = mu.nextReader()
		if err != nil && err != io.EOF {
			// empty ch
			go func() {
				for range ch {
				}
			}()
			// Wait for all goroutines finish
			close(ch)
			mu.wg.Wait()
			return err
		}
		if nextChunkLen == 0 && partNumber > 0 {
			// No need to upload empty part, if file was empty to start
			// with empty single part would of been created and never
			// started multipart upload.
			break
		}
		num := partNumber
		ch <- chunk{partNumber: num, buffer: reader, size: mu.readerSize}
	}
	// Wait for all goroutines finish
	close(ch)
	mu.wg.Wait()
	// Complete Multipart Upload
	err = mu.complete()
	if mu.getErr() != nil || err != nil {
		_ = mu.abort()
	}
	return err
}
<commit_msg>backend/qingstor: prune multiUploader.list()<commit_after>// Upload object to QingStor

// +build !plan9

package qingstor

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"hash"
	"io"
	"sort"
	"sync"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	qs "github.com/yunify/qingstor-sdk-go/v3/service"
)

const (
	// maxSinglePartSize = 1024 * 1024 * 1024 * 5 // The maximum allowed size when uploading a single object to QingStor
	// maxMultiPartSize = 1024 * 1024 * 1024 * 1 // The maximum allowed part size when uploading a part to QingStor
	minMultiPartSize = 1024 * 1024 * 4 // The minimum allowed part size when uploading a part to QingStor
	maxMultiParts    = 10000           // The maximum allowed number of parts in an multi-part upload
)

const (
	defaultUploadPartSize    = 1024 * 1024 * 64 // The default part size to buffer chunks of a payload into.
	defaultUploadConcurrency = 4                // the default number of goroutines to spin up when using multiPartUpload.
)

func readFillBuf(r io.Reader, b []byte) (offset int, err error) {
	for offset < len(b) && err == nil {
		var n int
		n, err = r.Read(b[offset:])
		offset += n
	}

	return offset, err
}

// uploadInput contains all input for upload requests to QingStor.
type uploadInput struct {
	body           io.Reader
	qsSvc          *qs.Service
	mimeType       string
	zone           string
	bucket         string
	key            string
	partSize       int64
	concurrency    int
	maxUploadParts int
}

// uploader internal structure to manage an upload to QingStor.
type uploader struct {
	cfg        *uploadInput
	totalSize  int64 // set to -1 if the size is not known
	readerPos  int64 // current reader position
	readerSize int64 // current reader content size
}

// newUploader creates a new Uploader instance to upload objects to QingStor.
func newUploader(in *uploadInput) *uploader {
	u := &uploader{
		cfg: in,
	}
	return u
}

// bucketInit initiate as bucket controller
func (u *uploader) bucketInit() (*qs.Bucket, error) {
	bucketInit, err := u.cfg.qsSvc.Bucket(u.cfg.bucket, u.cfg.zone)
	return bucketInit, err
}

// String converts uploader to a string
func (u *uploader) String() string {
	return fmt.Sprintf("QingStor bucket %s key %s", u.cfg.bucket, u.cfg.key)
}

// nextReader returns a seekable reader representing the next packet of data.
// This operation increases the shared u.readerPos counter, but note that it
// does not need to be wrapped in a mutex because nextReader is only called
// from the main thread.
func (u *uploader) nextReader() (io.ReadSeeker, int, error) {
	type readerAtSeeker interface {
		io.ReaderAt
		io.ReadSeeker
	}
	switch r := u.cfg.body.(type) {
	case readerAtSeeker:
		var err error
		n := u.cfg.partSize
		if u.totalSize >= 0 {
			bytesLeft := u.totalSize - u.readerPos

			if bytesLeft <= u.cfg.partSize {
				err = io.EOF
				n = bytesLeft
			}
		}
		reader := io.NewSectionReader(r, u.readerPos, n)
		u.readerPos += n
		u.readerSize = n
		return reader, int(n), err

	default:
		part := make([]byte, u.cfg.partSize)
		n, err := readFillBuf(r, part)
		u.readerPos += int64(n)
		u.readerSize = int64(n)
		return bytes.NewReader(part[0:n]), n, err
	}
}

// init will initialize all default options.
func (u *uploader) init() {
	if u.cfg.concurrency == 0 {
		u.cfg.concurrency = defaultUploadConcurrency
	}
	if u.cfg.partSize == 0 {
		u.cfg.partSize = defaultUploadPartSize
	}
	if u.cfg.maxUploadParts == 0 {
		u.cfg.maxUploadParts = maxMultiParts
	}
	// Try to get the total size for some optimizations
	u.totalSize = -1
	switch r := u.cfg.body.(type) {
	case io.Seeker:
		pos, _ := r.Seek(0, io.SeekCurrent)
		defer func() {
			_, _ = r.Seek(pos, io.SeekStart)
		}()

		n, err := r.Seek(0, io.SeekEnd)
		if err != nil {
			return
		}
		u.totalSize = n

		// Try to adjust partSize if it is too small and account for
		// integer division truncation.
		if u.totalSize/u.cfg.partSize >= u.cfg.partSize {
			// Add one to the part size to account for remainders
			// during the size calculation. e.g odd number of bytes.
			u.cfg.partSize = (u.totalSize / int64(u.cfg.maxUploadParts)) + 1
		}
	}
}

// singlePartUpload upload a single object that contentLength less than "defaultUploadPartSize"
func (u *uploader) singlePartUpload(buf io.Reader, size int64) error {
	bucketInit, _ := u.bucketInit()

	req := qs.PutObjectInput{
		ContentLength: &size,
		ContentType:   &u.cfg.mimeType,
		Body:          buf,
	}

	_, err := bucketInit.PutObject(u.cfg.key, &req)
	if err == nil {
		fs.Debugf(u, "Upload single object finished")
	}
	return err
}

// Upload upload a object into QingStor
func (u *uploader) upload() error {
	u.init()

	if u.cfg.partSize < minMultiPartSize {
		return errors.Errorf("part size must be at least %d bytes", minMultiPartSize)
	}

	// Do one read to determine if we have more than one part
	reader, _, err := u.nextReader()
	if err == io.EOF { // single part
		fs.Debugf(u, "Uploading as single part object to QingStor")
		return u.singlePartUpload(reader, u.readerPos)
	} else if err != nil {
		return errors.Errorf("read upload data failed: %s", err)
	}

	fs.Debugf(u, "Uploading as multi-part object to QingStor")
	mu := multiUploader{uploader: u}
	return mu.multiPartUpload(reader)
}

// internal structure to manage a specific multipart upload to QingStor.
type multiUploader struct {
	*uploader
	wg          sync.WaitGroup
	mtx         sync.Mutex
	err         error
	uploadID    *string
	objectParts completedParts
	hashMd5     hash.Hash
}

// keeps track of a single chunk of data being sent to QingStor.
type chunk struct {
	buffer     io.ReadSeeker
	partNumber int
	size       int64
}

// completedParts is a wrapper to make parts sortable by their part number,
// since QingStor required this list to be sent in sorted order.
type completedParts []*qs.ObjectPartType

func (a completedParts) Len() int           { return len(a) }
func (a completedParts) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a completedParts) Less(i, j int) bool { return *a[i].PartNumber < *a[j].PartNumber }

// String converts multiUploader to a string
func (mu *multiUploader) String() string {
	if uploadID := mu.uploadID; uploadID != nil {
		return fmt.Sprintf("QingStor bucket %s key %s uploadID %s", mu.cfg.bucket, mu.cfg.key, *uploadID)
	}
	return fmt.Sprintf("QingStor bucket %s key %s uploadID <nil>", mu.cfg.bucket, mu.cfg.key)
}

// getErr is a thread-safe getter for the error object
func (mu *multiUploader) getErr() error {
	mu.mtx.Lock()
	defer mu.mtx.Unlock()
	return mu.err
}

// setErr is a thread-safe setter for the error object
func (mu *multiUploader) setErr(e error) {
	mu.mtx.Lock()
	defer mu.mtx.Unlock()
	mu.err = e
}

// readChunk runs in worker goroutines to pull chunks off of the ch channel
// and send() them as UploadPart requests.
func (mu *multiUploader) readChunk(ch chan chunk) {
	defer mu.wg.Done()
	for {
		c, ok := <-ch
		if !ok {
			break
		}
		if mu.getErr() == nil {
			if err := mu.send(c); err != nil {
				mu.setErr(err)
			}
		}
	}
}

// initiate init an Multiple Object and obtain UploadID
func (mu *multiUploader) initiate() error {
	bucketInit, _ := mu.bucketInit()
	req := qs.InitiateMultipartUploadInput{
		ContentType: &mu.cfg.mimeType,
	}
	fs.Debugf(mu, "Initiating a multi-part upload")
	rsp, err := bucketInit.InitiateMultipartUpload(mu.cfg.key, &req)
	if err == nil {
		mu.uploadID = rsp.UploadID
		mu.hashMd5 = md5.New()
	}
	return err
}

// send upload a part into QingStor
func (mu *multiUploader) send(c chunk) error {
	bucketInit, _ := mu.bucketInit()
	req := qs.UploadMultipartInput{
		PartNumber:    &c.partNumber,
		UploadID:      mu.uploadID,
		ContentLength: &c.size,
		Body:          c.buffer,
	}
	fs.Debugf(mu, "Uploading a part to QingStor with partNumber %d and partSize %d", c.partNumber, c.size)
	_, err := bucketInit.UploadMultipart(mu.cfg.key, &req)
	if err != nil {
		return err
	}
	fs.Debugf(mu, "Done uploading part partNumber %d and partSize %d", c.partNumber, c.size)

	mu.mtx.Lock()
	defer mu.mtx.Unlock()

	_, _ = c.buffer.Seek(0, 0)
	_, _ = io.Copy(mu.hashMd5, c.buffer)

	parts := qs.ObjectPartType{PartNumber: &c.partNumber, Size: &c.size}
	mu.objectParts = append(mu.objectParts, &parts)
	return err
}

// complete complete an multipart upload
func (mu *multiUploader) complete() error {
	var err error
	if err = mu.getErr(); err != nil {
		return err
	}
	bucketInit, _ := mu.bucketInit()
	//if err = mu.list(); err != nil {
	//	return err
	//}
	//md5String := fmt.Sprintf("\"%s\"", hex.EncodeToString(mu.hashMd5.Sum(nil)))

	md5String := fmt.Sprintf("\"%x\"", mu.hashMd5.Sum(nil))
	sort.Sort(mu.objectParts)
	req := qs.CompleteMultipartUploadInput{
		UploadID:    mu.uploadID,
		ObjectParts: mu.objectParts,
		ETag:        &md5String,
	}
	fs.Debugf(mu, "Completing multi-part object")
	_, err = bucketInit.CompleteMultipartUpload(mu.cfg.key, &req)
	if err == nil {
		fs.Debugf(mu, "Complete multi-part finished")
	}
	return err
}

// abort abort an multipart upload
func (mu *multiUploader) abort() error {
	var err error
	bucketInit, _ := mu.bucketInit()

	if uploadID := mu.uploadID; uploadID != nil {
		req := qs.AbortMultipartUploadInput{
			UploadID: uploadID,
		}
		fs.Debugf(mu, "Aborting multi-part object %q", *uploadID)
		_, err = bucketInit.AbortMultipartUpload(mu.cfg.key, &req)
	}

	return err
}

// multiPartUpload upload a multiple object into QingStor
func (mu *multiUploader) multiPartUpload(firstBuf io.ReadSeeker) error {
	var err error
	//Initiate an multi-part upload
	if err = mu.initiate(); err != nil {
		return err
	}

	ch := make(chan chunk, mu.cfg.concurrency)
	for i := 0; i < mu.cfg.concurrency; i++ {
		mu.wg.Add(1)
		go mu.readChunk(ch)
	}

	var partNumber int
	ch <- chunk{partNumber: partNumber, buffer: firstBuf, size: mu.readerSize}

	for mu.getErr() == nil {
		partNumber++
		// This upload exceeded maximum number of supported parts, error now.
		if partNumber > mu.cfg.maxUploadParts || partNumber > maxMultiParts {
			var msg string
			if partNumber > mu.cfg.maxUploadParts {
				msg = fmt.Sprintf("exceeded total allowed configured maxUploadParts (%d). "+
					"Adjust PartSize to fit in this limit", mu.cfg.maxUploadParts)
			} else {
				msg = fmt.Sprintf("exceeded total allowed QingStor limit maxUploadParts (%d). "+
					"Adjust PartSize to fit in this limit", maxMultiParts)
			}
			mu.setErr(errors.New(msg))
			break
		}

		var reader io.ReadSeeker
		var nextChunkLen int
		reader, nextChunkLen, err = mu.nextReader()
		if err != nil && err != io.EOF {
			// empty ch
			go func() {
				for range ch {
				}
			}()
			// Wait for all goroutines finish
			close(ch)
			mu.wg.Wait()
			return err
		}
		if nextChunkLen == 0 && partNumber > 0 {
			// No need to upload empty part, if file was empty to start
			// with empty single part would of been created and never
			// started multipart upload.
			break
		}
		num := partNumber
		ch <- chunk{partNumber: num, buffer: reader, size: mu.readerSize}
	}
	// Wait for all goroutines finish
	close(ch)
	mu.wg.Wait()
	// Complete Multipart Upload
	err = mu.complete()
	if mu.getErr() != nil || err != nil {
		_ = mu.abort()
	}
	return err
}
<|endoftext|>


<commit_before>package downloader

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/axgle/mahonia"
	"github.com/crawlerclub/x/types"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

func Download(requestInfo *types.HttpRequest) *types.HttpResponse {
	var timeout time.Duration
	if requestInfo.Timeout > 0 {
		timeout = time.Duration(requestInfo.Timeout) * time.Second
	} else {
		timeout = 30 * time.Second
	}
	client := &http.Client{
		Timeout: timeout,
	}
	responseInfo := &types.HttpResponse{
		Url: requestInfo.Url,
	}
	transport := http.Transport{
		DisableKeepAlives: true,
		TLSClientConfig:   &tls.Config{InsecureSkipVerify: true},
	}

	//proxy
	if requestInfo.UseProxy {
		var proxy string
		var err error
		if len(requestInfo.Proxy) > 0 {
			proxy = requestInfo.Proxy
		} else {
			proxy, err = GetProxy()
			if err != nil {
				responseInfo.Error = err
				return responseInfo
			}
		}
		responseInfo.Proxy = proxy
		urlProxy, err := url.Parse(proxy)
		if err != nil {
			responseInfo.Error = errors.New(fmt.Sprintf("failed to parse proxy: %s", proxy))
			return responseInfo
		}
		transport.Proxy = http.ProxyURL(urlProxy)
	}

	client.Transport = &transport

	req, err := http.NewRequest(requestInfo.Method, requestInfo.Url, strings.NewReader(requestInfo.PostData))
	if err != nil {
		responseInfo.Error = err
		return responseInfo
	}
	var resp *http.Response
	resp, err = client.Do(req)
	if err != nil {
		responseInfo.Error = err
		return responseInfo
	}

	responseInfo.StatusCode = resp.StatusCode
	defer resp.Body.Close()

	var contentLen int64
	contentLen, err = strconv.ParseInt(resp.Header.Get("content-length"), 10, 64)
	if err != nil {
		//
	} else if requestInfo.MaxLen > 0 && contentLen > requestInfo.MaxLen {
		responseInfo.Error = errors.New("reponse size too large")
		return responseInfo
	}

	var reader io.ReadCloser
	switch resp.Header.Get("Content-Encoding") {
	case "gzip":
		if reader, err = gzip.NewReader(resp.Body); err != nil {
			responseInfo.Error = err
			return responseInfo
		}
		defer reader.Close()
	case "deflate":
		if reader, err = zlib.NewReader(resp.Body); err != nil {
			responseInfo.Error = err
			return responseInfo
		}
		defer reader.Close()
	default:
		reader = resp.Body
	}

	var readLen int64 = 0
	respBuf := bytes.NewBuffer([]byte{})
	for {
		readData := make([]byte, 4096)
		length, err := reader.Read(readData)
		respBuf.Write(readData[:length])
		readLen += int64(length)
		if err != nil {
			if err == io.EOF {
				break
			}
			responseInfo.Error = errors.New("reponse size too large - count")
			return responseInfo
		}
	}
	responseInfo.Content = respBuf.Bytes()
	var encoding string
	encoding, err = GuessEncoding(responseInfo.Content)
	if err != nil {
		//
		responseInfo.Text = string(responseInfo.Content)
		responseInfo.Encoding = ""
		return responseInfo
	}
	encoder := mahonia.NewDecoder(encoding)
	if encoder == nil {
		responseInfo.Text = string(responseInfo.Content)
		responseInfo.Encoding = ""
		return responseInfo
	}
	responseInfo.Text = encoder.ConvertString(string(responseInfo.Content))
	responseInfo.Encoding = encoding
	return responseInfo
}
<commit_msg>fix: add Content-Type header for POST method<commit_after>package downloader

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/axgle/mahonia"
	"github.com/crawlerclub/x/types"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

func Download(requestInfo *types.HttpRequest) *types.HttpResponse {
	var timeout time.Duration
	if requestInfo.Timeout > 0 {
		timeout = time.Duration(requestInfo.Timeout) * time.Second
	} else {
		timeout = 30 * time.Second
	}
	client := &http.Client{
		Timeout: timeout,
	}
	responseInfo := &types.HttpResponse{
		Url: requestInfo.Url,
	}
	transport := http.Transport{
		DisableKeepAlives: true,
		TLSClientConfig:   &tls.Config{InsecureSkipVerify: true},
	}

	//proxy
	if requestInfo.UseProxy {
		var proxy string
		var err error
		if len(requestInfo.Proxy) > 0 {
			proxy = requestInfo.Proxy
		} else {
			proxy, err = GetProxy()
			if err != nil {
				responseInfo.Error = err
				return responseInfo
			}
		}
		responseInfo.Proxy = proxy
		urlProxy, err := url.Parse(proxy)
		if err != nil {
			responseInfo.Error = errors.New(fmt.Sprintf("failed to parse proxy: %s", proxy))
			return responseInfo
		}
		transport.Proxy = http.ProxyURL(urlProxy)
	}

	client.Transport = &transport

	req, err := http.NewRequest(requestInfo.Method, requestInfo.Url, strings.NewReader(requestInfo.PostData))
	if requestInfo.Method == "POST" {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}
	if err != nil {
		responseInfo.Error = err
		return responseInfo
	}
	var resp *http.Response
	resp, err = client.Do(req)
	if err != nil {
		responseInfo.Error = err
		return responseInfo
	}

	responseInfo.StatusCode = resp.StatusCode
	defer resp.Body.Close()

	var contentLen int64
	contentLen, err = strconv.ParseInt(resp.Header.Get("content-length"), 10, 64)
	if err != nil {
		//
	} else if requestInfo.MaxLen > 0 && contentLen > requestInfo.MaxLen {
		responseInfo.Error = errors.New("reponse size too large")
		return responseInfo
	}

	var reader io.ReadCloser
	switch resp.Header.Get("Content-Encoding") {
	case "gzip":
		if reader, err = gzip.NewReader(resp.Body); err != nil {
			responseInfo.Error = err
			return responseInfo
		}
		defer reader.Close()
	case "deflate":
		if reader, err = zlib.NewReader(resp.Body); err != nil {
			responseInfo.Error = err
			return responseInfo
		}
		defer reader.Close()
	default:
		reader = resp.Body
	}

	var readLen int64 = 0
	respBuf := bytes.NewBuffer([]byte{})
	for {
		readData := make([]byte, 4096)
		length, err := reader.Read(readData)
		respBuf.Write(readData[:length])
		readLen += int64(length)
		if err != nil {
			if err == io.EOF {
				break
			}
			responseInfo.Error = errors.New("reponse size too large - count")
			return responseInfo
		}
	}
	responseInfo.Content = respBuf.Bytes()
	var encoding string
	encoding, err = GuessEncoding(responseInfo.Content)
	if err != nil {
		//
		responseInfo.Text = string(responseInfo.Content)
		responseInfo.Encoding = ""
		return responseInfo
	}
	encoder := mahonia.NewDecoder(encoding)
	if encoder == nil {
		responseInfo.Text = string(responseInfo.Content)
		responseInfo.Encoding = ""
		return responseInfo
	}
	responseInfo.Text = encoder.ConvertString(string(responseInfo.Content))
	responseInfo.Encoding = encoding
	return responseInfo
}
<|endoftext|>


<commit_before>package gitbackup

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

type repository struct {
	name     string
	cloneURL string
}

// BackupTarget backs up an entity that holds one or more git repositories and
// has an interface to retrieve that list of repositories.
// Examples of entities include:
//   - A GitHub user.
//   - A BitBucket user.
//   - A GitHub organization.
func BackupTarget(target Target, backupDirectory string) error {
	fmt.Printf("########## Backing up target %s ##########\n\n", target.Name)

	// Retrieve a list of all the git repositories available from the target.
	var repoList []repository
	var err error
	switch target.Source {
	case "github":
		repoList, err = getGitHubRepoList(target, backupDirectory)
	case "bitbucket":
		repoList, err = getBitBucketRepoList(target, backupDirectory)
	default:
		err = fmt.Errorf(`"%s" is not a recognized source type`, target.Source)
	}
	if err != nil {
		return err
	}

	// Back up each repository found.
	for _, repo := range repoList {
		fmt.Println(fmt.Sprintf("#> %s", repo.name))
		if includeRepository(repo.name, target) {
			backupRepository(
				target.Name,
				repo.name,
				repo.cloneURL,
				backupDirectory,
			)
			fmt.Println("")
		} else {
			fmt.Print("Skipped.\n\n")
		}
	}

	fmt.Println("")

	return nil
}

// getGitHubRepoList finds all the repositories belonging to a given user or
// organization on GitHub.
func getGitHubRepoList(target Target, backupDirectory string) ([]repository, error) {
	// Create URL to request list of repos.
	requestURL := fmt.Sprintf(
		"https://api.github.com/%s/%s/repos?access_token=%s&per_page=200",
		target.Type,
		target.Entity,
		target.Token,
	)

	// Retrieve list of repositories.
	response, err := http.Get(requestURL)
	if err != nil {
		return nil, fmt.Errorf("Failed to connect with the source to retrieve the list of repositories: %s", err)
	}
	defer response.Body.Close()
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve the list of repositories: %s", err)
	}

	// Parse JSON response.
	var data []map[string]interface{}
	if err := json.Unmarshal(contents, &data); err != nil {
		return nil, fmt.Errorf("Failed to parse JSON: %s", err)
	}

	// Make a list of repositories.
	repoList := make([]repository, len(data))
	for i, repo := range data {
		repoName, _ := repo["name"].(string)
		cloneURL, _ := repo["clone_url"].(string)
		cloneURL = strings.Replace(
			cloneURL,
			"https://",
			fmt.Sprintf("https://%s:%s@", target.Entity, target.Token),
			1,
		)
		repoList[i] = repository{name: repoName, cloneURL: cloneURL}
	}

	// No errors.
	return repoList, nil
}

// getBitBucketRepoList finds all the repositories belonging to a given user on
// BitBucket.
func getBitBucketRepoList(target Target, backupDirectory string) ([]repository, error) {
	// Create URL to request list of repos.
	// TODO: support pagination.
	requestURL := fmt.Sprintf(
		"https://%s:%s@bitbucket.org/api/2.0/repositories/%s?page=1&pagelen=100",
		target.Entity,
		target.Password,
		target.Entity,
	)

	// Retrieve list of repositories.
	response, err := http.Get(requestURL)
	if err != nil {
		return nil, fmt.Errorf("Failed to connect with the source to retrieve the list of repositories: %s", err)
	}
	defer response.Body.Close()
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve the list of repositories: %s", err)
	}

	// Parse JSON response.
	var metadata map[string]json.RawMessage
	if err := json.Unmarshal(contents, &metadata); err != nil {
		return nil, fmt.Errorf("Failed to parse JSON: %s", err)
	}
	var data []map[string]json.RawMessage
	if err := json.Unmarshal(metadata["values"], &data); err != nil {
		return nil, fmt.Errorf("Failed to parse JSON: %s", err)
	}

	// Make a list of repositories.
	repoList := make([]repository, len(data))
	for i, repo := range data {
		// Parse the remaining JSON message that pertains to this repository.
		var repoName string
		if err := json.Unmarshal(repo["name"], &repoName); err != nil {
			return nil, fmt.Errorf("Failed to parse JSON: %s", err)
		}
		var links map[string]json.RawMessage
		if err := json.Unmarshal(repo["links"], &links); err != nil {
			return nil, fmt.Errorf("Failed to parse JSON: %s", err)
		}
		var cloneLinks []map[string]string
		if err := json.Unmarshal(links["clone"], &cloneLinks); err != nil {
			return nil, fmt.Errorf("Failed to parse JSON: %s", err)
		}

		// Find the https URL to use for cloning.
		var cloneURL string
		for _, link := range cloneLinks {
			if link["name"] == "https" {
				cloneURL = link["href"]
			}
		}
		if cloneURL == "" {
			return nil, fmt.Errorf("Could not determine HTTPS cloning URL: %s", cloneLinks)
		}

		// Determine URL for cloning.
		cloneURL = strings.Replace(
			cloneURL,
			fmt.Sprintf("https://%s@", target.Entity),
			fmt.Sprintf("https://%s:%s@", target.Entity, target.Password),
			1,
		)

		repoList[i] = repository{name: repoName, cloneURL: cloneURL}
	}

	// No errors.
	return repoList, nil
}

// backupRepository takes a remote git repository and backs it up locally.
// Note that this makes a mirror repository - in other words, the backup only
// contains the content of a normal .git repository but no working directory,
// which saves space. You can always get a normal repository from the backup by
// doing a normal git clone of the backup itself.
func backupRepository(targetName string, repoName string, cloneURL string, backupDirectory string) {
	cloneDirectory := filepath.Join(backupDirectory, targetName, repoName)

	if _, err := os.Stat(cloneDirectory); os.IsNotExist(err) {
		// The repo doesn't exist locally, clone it.
		cmd := exec.Command("git", "clone", "--mirror", cloneURL, cloneDirectory)
		cmdOut, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("Error cloning the repository:", err)
		} else {
			if len(cmdOut) > 0 {
				fmt.Printf(string(cmdOut))
			}
			fmt.Println("Cloned repository.")
		}
	} else {
		// The repo already exists, pull updates.
		cmd := exec.Command("git", "fetch", "-p", cloneURL)
		cmd.Dir = cloneDirectory
		cmdOut, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("Error pulling in the repository:", err)
		} else {
			// Display pulled information.
			if len(cmdOut) > 0 {
				fmt.Printf(string(cmdOut))
			}
			fmt.Println("Pulled latest updates in the repository.")
		}
	}
}

// includeRepository takes a repository name and the information about the
// target it is part of, and determines whether the repository should be backed
// up or not.
func includeRepository(repoName string, target Target) bool {
	if target.Skip != "" {
		r, err := regexp.Compile(target.Skip)
		if err != nil {
			log.Fatalf(`"skip" does not specify a valid regular expression: %s`, err)
		}
		return r.MatchString(repoName) == false
	}
	if target.Only != "" {
		r, err := regexp.Compile(target.Only)
		if err != nil {
			log.Fatalf(`"only" does not specify a valid regular expression: %s`, err)
		}
		return r.MatchString(repoName) == true
	}

	return true
}
<commit_msg>Cleaned up slice.<commit_after>package gitbackup

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

type repository struct {
	name     string
	cloneURL string
}

// BackupTarget backs up an entity that holds one or more git repositories and
// has an interface to retrieve that list of repositories.
// Examples of entities include:
//   - A GitHub user.
//   - A BitBucket user.
//   - A GitHub organization.
func BackupTarget(target Target, backupDirectory string) error {
	fmt.Printf("########## Backing up target %s ##########\n\n", target.Name)

	// Retrieve a list of all the git repositories available from the target.
	var repoList []repository
	var err error
	switch target.Source {
	case "github":
		repoList, err = getGitHubRepoList(target, backupDirectory)
	case "bitbucket":
		repoList, err = getBitBucketRepoList(target, backupDirectory)
	default:
		err = fmt.Errorf(`"%s" is not a recognized source type`, target.Source)
	}
	if err != nil {
		return err
	}

	// Back up each repository found.
	for _, repo := range repoList {
		fmt.Println(fmt.Sprintf("#> %s", repo.name))
		if includeRepository(repo.name, target) {
			backupRepository(
				target.Name,
				repo.name,
				repo.cloneURL,
				backupDirectory,
			)
			fmt.Println("")
		} else {
			fmt.Print("Skipped.\n\n")
		}
	}

	fmt.Println("")

	return nil
}

// getGitHubRepoList finds all the repositories belonging to a given user or
// organization on GitHub.
func getGitHubRepoList(target Target, backupDirectory string) ([]repository, error) {
	// Create URL to request list of repos.
	requestURL := fmt.Sprintf(
		"https://api.github.com/%s/%s/repos?access_token=%s&per_page=200",
		target.Type,
		target.Entity,
		target.Token,
	)

	// Retrieve list of repositories.
	response, err := http.Get(requestURL)
	if err != nil {
		return nil, fmt.Errorf("Failed to connect with the source to retrieve the list of repositories: %s", err)
	}
	defer response.Body.Close()
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve the list of repositories: %s", err)
	}

	// Parse JSON response.
	var data []map[string]interface{}
	if err := json.Unmarshal(contents, &data); err != nil {
		return nil, fmt.Errorf("Failed to parse JSON: %s", err)
	}

	// Make a list of repositories.
	var repoList []repository
	for _, repo := range data {
		repoName, _ := repo["name"].(string)
		cloneURL, _ := repo["clone_url"].(string)
		cloneURL = strings.Replace(
			cloneURL,
			"https://",
			fmt.Sprintf("https://%s:%s@", target.Entity, target.Token),
			1,
		)
		repo := repository{name: repoName, cloneURL: cloneURL}
		repoList = append(repoList,repo)
	}

	// No errors.
	return repoList, nil
}

// getBitBucketRepoList finds all the repositories belonging to a given user on
// BitBucket.
func getBitBucketRepoList(target Target, backupDirectory string) ([]repository, error) {
	// Create URL to request list of repos.
	// TODO: support pagination.
	requestURL := fmt.Sprintf(
		"https://%s:%s@bitbucket.org/api/2.0/repositories/%s?page=1&pagelen=100",
		target.Entity,
		target.Password,
		target.Entity,
	)

	// Retrieve list of repositories.
	response, err := http.Get(requestURL)
	if err != nil {
		return nil, fmt.Errorf("Failed to connect with the source to retrieve the list of repositories: %s", err)
	}
	defer response.Body.Close()
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, fmt.Errorf("Failed to retrieve the list of repositories: %s", err)
	}

	// Parse JSON response.
	var metadata map[string]json.RawMessage
	if err := json.Unmarshal(contents, &metadata); err != nil {
		return nil, fmt.Errorf("Failed to parse JSON: %s", err)
	}
	var data []map[string]json.RawMessage
	if err := json.Unmarshal(metadata["values"], &data); err != nil {
		return nil, fmt.Errorf("Failed to parse JSON: %s", err)
	}

	// Make a list of repositories.
	repoList := make([]repository, len(data))
	for i, repo := range data {
		// Parse the remaining JSON message that pertains to this repository.
		var repoName string
		if err := json.Unmarshal(repo["name"], &repoName); err != nil {
			return nil, fmt.Errorf("Failed to parse JSON: %s", err)
		}
		var links map[string]json.RawMessage
		if err := json.Unmarshal(repo["links"], &links); err != nil {
			return nil, fmt.Errorf("Failed to parse JSON: %s", err)
		}
		var cloneLinks []map[string]string
		if err := json.Unmarshal(links["clone"], &cloneLinks); err != nil {
			return nil, fmt.Errorf("Failed to parse JSON: %s", err)
		}

		// Find the https URL to use for cloning.
		var cloneURL string
		for _, link := range cloneLinks {
			if link["name"] == "https" {
				cloneURL = link["href"]
			}
		}
		if cloneURL == "" {
			return nil, fmt.Errorf("Could not determine HTTPS cloning URL: %s", cloneLinks)
		}

		// Determine URL for cloning.
		cloneURL = strings.Replace(
			cloneURL,
			fmt.Sprintf("https://%s@", target.Entity),
			fmt.Sprintf("https://%s:%s@", target.Entity, target.Password),
			1,
		)

		repoList[i] = repository{name: repoName, cloneURL: cloneURL}
	}

	// No errors.
	return repoList, nil
}

// backupRepository takes a remote git repository and backs it up locally.
// Note that this makes a mirror repository - in other words, the backup only
// contains the content of a normal .git repository but no working directory,
// which saves space. You can always get a normal repository from the backup by
// doing a normal git clone of the backup itself.
func backupRepository(targetName string, repoName string, cloneURL string, backupDirectory string) {
	cloneDirectory := filepath.Join(backupDirectory, targetName, repoName)

	if _, err := os.Stat(cloneDirectory); os.IsNotExist(err) {
		// The repo doesn't exist locally, clone it.
		cmd := exec.Command("git", "clone", "--mirror", cloneURL, cloneDirectory)
		cmdOut, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("Error cloning the repository:", err)
		} else {
			if len(cmdOut) > 0 {
				fmt.Printf(string(cmdOut))
			}
			fmt.Println("Cloned repository.")
		}
	} else {
		// The repo already exists, pull updates.
		cmd := exec.Command("git", "fetch", "-p", cloneURL)
		cmd.Dir = cloneDirectory
		cmdOut, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("Error pulling in the repository:", err)
		} else {
			// Display pulled information.
			if len(cmdOut) > 0 {
				fmt.Printf(string(cmdOut))
			}
			fmt.Println("Pulled latest updates in the repository.")
		}
	}
}

// includeRepository takes a repository name and the information about the
// target it is part of, and determines whether the repository should be backed
// up or not.
func includeRepository(repoName string, target Target) bool {
	if target.Skip != "" {
		r, err := regexp.Compile(target.Skip)
		if err != nil {
			log.Fatalf(`"skip" does not specify a valid regular expression: %s`, err)
		}
		return r.MatchString(repoName) == false
	}
	if target.Only != "" {
		r, err := regexp.Compile(target.Only)
		if err != nil {
			log.Fatalf(`"only" does not specify a valid regular expression: %s`, err)
		}
		return r.MatchString(repoName) == true
	}

	return true
}
<|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/ngaut/systimemon"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/perfschema"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/server"
	"github.com/pingcap/tidb/sessionctx/binloginfo"
	"github.com/pingcap/tidb/store/localstore/boltdb"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/printer"
	"github.com/pingcap/tipb/go-binlog"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
	"google.golang.org/grpc"
)

var (
	version         = flag.Bool("v", false, "print version information and exit")
	store           = flag.String("store", "goleveldb", "registered store name, [memory, goleveldb, boltdb, tikv]")
	storePath       = flag.String("path", "/tmp/tidb", "tidb storage path")
	logLevel        = flag.String("L", "info", "log level: info, debug, warn, error, fatal")
	host            = flag.String("host", "0.0.0.0", "tidb server host")
	port            = flag.String("P", "4000", "tidb server port")
	statusPort      = flag.String("status", "10080", "tidb server status port")
	lease           = flag.String("lease", "1s", "schema lease duration, very dangerous to change only if you know what you do")
	socket          = flag.String("socket", "", "The socket file to use for connection.")
	enablePS        = flag.Bool("perfschema", false, "If enable performance schema.")
	reportStatus    = flag.Bool("report-status", true, "If enable status report HTTP service.")
	logFile         = flag.String("log-file", "", "log file path")
	joinCon         = flag.Int("join-concurrency", 5, "the number of goroutines that participate joining.")
	crossJoin       = flag.Bool("cross-join", true, "whether support cartesian product or not.")
	metricsAddr     = flag.String("metrics-addr", "", "prometheus pushgateway address, leaves it empty will disable prometheus push.")
	metricsInterval = flag.Int("metrics-interval", 15, "prometheus client push interval in second, set \"0\" to disable prometheus push.")
	binlogSocket    = flag.String("binlog-socket", "", "socket file to write binlog")
)

func main() {
	tidb.RegisterLocalStore("boltdb", boltdb.Driver{})
	tidb.RegisterStore("tikv", tikv.Driver{})

	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()
	if *version {
		printer.PrintRawTiDBInfo()
		os.Exit(0)
	}

	leaseDuration := parseLease()
	tidb.SetSchemaLease(leaseDuration)

	cfg := &server.Config{
		Addr:         fmt.Sprintf("%s:%s", *host, *port),
		LogLevel:     *logLevel,
		StatusAddr:   fmt.Sprintf(":%s", *statusPort),
		Socket:       *socket,
		ReportStatus: *reportStatus,
	}

	// set log options
	if len(*logFile) > 0 {
		err := log.SetOutputByName(*logFile)
		if err != nil {
			log.Fatal(errors.ErrorStack(err))
		}
		log.SetRotateByDay()
	}

	if joinCon != nil && *joinCon > 0 {
		plan.JoinConcurrency = *joinCon
	}
	plan.AllowCartesianProduct = *crossJoin
	// Call this before setting log level to make sure that TiDB info could be printed.
	printer.PrintTiDBInfo()
	log.SetLevelByString(cfg.LogLevel)

	store := createStore()

	if *enablePS {
		perfschema.EnablePerfSchema()
	}
	if *binlogSocket != "" {
		createBinlogClient()
	}

	// Create a session to load information schema.
	se, err := tidb.CreateSession(store)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	se.Close()

	var driver server.IDriver
	driver = server.NewTiDBDriver(store)
	var svr *server.Server
	svr, err = server.NewServer(cfg, driver)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go func() {
		sig := <-sc
		log.Infof("Got signal [%d] to exit.", sig)
		svr.Close()
		os.Exit(0)
	}()

	go systimemon.StartMonitor(time.Now, func() {
		log.Error("error: system time jump backward")
	})

	pushMetric(*metricsAddr, time.Duration(*metricsInterval)*time.Second)

	log.Error(svr.Run())
}

func createStore() kv.Storage {
	fullPath := fmt.Sprintf("%s://%s", *store, *storePath)
	store, err := tidb.NewStore(fullPath)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	return store
}

func createBinlogClient() {
	dialerOpt := grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
		return net.DialTimeout("unix", addr, timeout)
	})
	clientCon, err := grpc.Dial(*binlogSocket, dialerOpt, grpc.WithInsecure())
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	binloginfo.PumpClient = binlog.NewPumpClient(clientCon)
}

// Prometheus push.
const zeroDuration = time.Duration(0)

// PushMetric pushs metircs in background.
func pushMetric(addr string, interval time.Duration) {
	if interval == zeroDuration || len(addr) == 0 {
		log.Info("disable Prometheus push client")
		return
	}
	log.Infof("start Prometheus push client with server addr %s and interval %d", addr, interval)
	go prometheusPushClient(addr, interval)
}

// PrometheusPushClient pushs metrics to Prometheus Pushgateway.
func prometheusPushClient(addr string, interval time.Duration) {
	// TODO: TiDB do not have uniq name, so we use host+port to compose a name.
	job := "tidb"
	for {
		err := push.FromGatherer(
			job, push.HostnameGroupingKey(),
			addr,
			prometheus.DefaultGatherer,
		)
		if err != nil {
			log.Errorf("could not push metrics to Prometheus Pushgateway: %v", err)
		}
		time.Sleep(interval)
	}
}

// parseLease parses lease argument string.
func parseLease() time.Duration {
	dur, err := time.ParseDuration(*lease)
	if err != nil {
		dur, err = time.ParseDuration(*lease + "s")
	}
	if err != nil || dur < 0 {
		log.Fatalf("invalid lease duration %s", *lease)
	}
	return dur
}
<commit_msg>tidb-server:fix metrics push way<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/ngaut/systimemon"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/perfschema"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/server"
	"github.com/pingcap/tidb/sessionctx/binloginfo"
	"github.com/pingcap/tidb/store/localstore/boltdb"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/printer"
	"github.com/pingcap/tipb/go-binlog"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
	"google.golang.org/grpc"
)

var (
	version         = flag.Bool("v", false, "print version information and exit")
	store           = flag.String("store", "goleveldb", "registered store name, [memory, goleveldb, boltdb, tikv]")
	storePath       = flag.String("path", "/tmp/tidb", "tidb storage path")
	logLevel        = flag.String("L", "info", "log level: info, debug, warn, error, fatal")
	host            = flag.String("host", "0.0.0.0", "tidb server host")
	port            = flag.String("P", "4000", "tidb server port")
	statusPort      = flag.String("status", "10080", "tidb server status port")
	lease           = flag.String("lease", "1s", "schema lease duration, very dangerous to change only if you know what you do")
	socket          = flag.String("socket", "", "The socket file to use for connection.")
	enablePS        = flag.Bool("perfschema", false, "If enable performance schema.")
	reportStatus    = flag.Bool("report-status", true, "If enable status report HTTP service.")
	logFile         = flag.String("log-file", "", "log file path")
	joinCon         = flag.Int("join-concurrency", 5, "the number of goroutines that participate joining.")
	crossJoin       = flag.Bool("cross-join", true, "whether support cartesian product or not.")
	metricsAddr     = flag.String("metrics-addr", "", "prometheus pushgateway address, leaves it empty will disable prometheus push.")
	metricsInterval = flag.Int("metrics-interval", 15, "prometheus client push interval in second, set \"0\" to disable prometheus push.")
	binlogSocket    = flag.String("binlog-socket", "", "socket file to write binlog")
)

func main() {
	tidb.RegisterLocalStore("boltdb", boltdb.Driver{})
	tidb.RegisterStore("tikv", tikv.Driver{})

	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()
	if *version {
		printer.PrintRawTiDBInfo()
		os.Exit(0)
	}

	leaseDuration := parseLease()
	tidb.SetSchemaLease(leaseDuration)

	cfg := &server.Config{
		Addr:         fmt.Sprintf("%s:%s", *host, *port),
		LogLevel:     *logLevel,
		StatusAddr:   fmt.Sprintf(":%s", *statusPort),
		Socket:       *socket,
		ReportStatus: *reportStatus,
	}

	// set log options
	if len(*logFile) > 0 {
		err := log.SetOutputByName(*logFile)
		if err != nil {
			log.Fatal(errors.ErrorStack(err))
		}
		log.SetRotateByDay()
	}

	if joinCon != nil && *joinCon > 0 {
		plan.JoinConcurrency = *joinCon
	}
	plan.AllowCartesianProduct = *crossJoin
	// Call this before setting log level to make sure that TiDB info could be printed.
	printer.PrintTiDBInfo()
	log.SetLevelByString(cfg.LogLevel)

	store := createStore()

	if *enablePS {
		perfschema.EnablePerfSchema()
	}
	if *binlogSocket != "" {
		createBinlogClient()
	}

	// Create a session to load information schema.
	se, err := tidb.CreateSession(store)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	se.Close()

	var driver server.IDriver
	driver = server.NewTiDBDriver(store)
	var svr *server.Server
	svr, err = server.NewServer(cfg, driver)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go func() {
		sig := <-sc
		log.Infof("Got signal [%d] to exit.", sig)
		svr.Close()
		os.Exit(0)
	}()

	go systimemon.StartMonitor(time.Now, func() {
		log.Error("error: system time jump backward")
	})

	pushMetric(*metricsAddr, time.Duration(*metricsInterval)*time.Second)

	log.Error(svr.Run())
}

func createStore() kv.Storage {
	fullPath := fmt.Sprintf("%s://%s", *store, *storePath)
	store, err := tidb.NewStore(fullPath)
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	return store
}

func createBinlogClient() {
	dialerOpt := grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
		return net.DialTimeout("unix", addr, timeout)
	})
	clientCon, err := grpc.Dial(*binlogSocket, dialerOpt, grpc.WithInsecure())
	if err != nil {
		log.Fatal(errors.ErrorStack(err))
	}
	binloginfo.PumpClient = binlog.NewPumpClient(clientCon)
}

// Prometheus push.
const zeroDuration = time.Duration(0)

// PushMetric pushs metircs in background.
func pushMetric(addr string, interval time.Duration) {
	if interval == zeroDuration || len(addr) == 0 {
		log.Info("disable Prometheus push client")
		return
	}
	log.Infof("start Prometheus push client with server addr %s and interval %d", addr, interval)
	go prometheusPushClient(addr, interval)
}

// PrometheusPushClient pushs metrics to Prometheus Pushgateway.
func prometheusPushClient(addr string, interval time.Duration) {
	// TODO: TiDB do not have uniq name, so we use host+port to compose a name.
	job := "tidb"
	for {
		err := push.AddFromGatherer(
			job, push.HostnameGroupingKey(),
			addr,
			prometheus.DefaultGatherer,
		)
		if err != nil {
			log.Errorf("could not push metrics to Prometheus Pushgateway: %v", err)
		}
		time.Sleep(interval)
	}
}

// parseLease parses lease argument string.
func parseLease() time.Duration {
	dur, err := time.ParseDuration(*lease)
	if err != nil {
		dur, err = time.ParseDuration(*lease + "s")
	}
	if err != nil || dur < 0 {
		log.Fatalf("invalid lease duration %s", *lease)
	}
	return dur
}
<|endoftext|>


<commit_before>// Copyright 2013-2017 Aerospike, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aerospike

import (
	"math/rand"
	"runtime"
	"strings"
	"testing"
	// "time"

	_ "net/http/pprof"
	// . "github.com/aerospike/aerospike-client-go"
)

func doGet(set string, value interface{}, b *testing.B) {
	var err error
	policy := NewPolicy()

	dataBuffer := make([]byte, 1024*1024)

	binNames := []string{}
	key, _ := NewKey("test", set, 1000)

	for i := 0; i < b.N; i++ {
		command := newReadCommand(nil, policy, key, binNames)
		command.baseCommand.dataBuffer = dataBuffer
		err = command.writeBuffer(command)
		if err != nil {
			panic(err)
		}
	}
}

func Benchmark_ReadCommand_________Int64(b *testing.B) {
	set := "put_bench_integer"
	value := rand.Int63()
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_________Int32(b *testing.B) {
	set := "put_bench_integer"
	value := rand.Int31()
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String______1(b *testing.B) {
	set := "put_bench_str_1"
	value := strings.Repeat("s", 1)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String_____10(b *testing.B) {
	set := "put_bench_str_10"
	value := strings.Repeat("s", 10)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String____100(b *testing.B) {
	set := "put_bench_str_100"
	value := strings.Repeat("s", 100)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String___1000(b *testing.B) {
	set := "put_bench_str_1000"
	value := strings.Repeat("s", 1000)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String__10000(b *testing.B) {
	set := "put_bench_str_10000"
	value := strings.Repeat("s", 10000)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String_100000(b *testing.B) {
	set := "put_bench_str_10000"
	value := strings.Repeat("s", 100000)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_Complex_Array(b *testing.B) {
	set := "put_bench_str_10000"
	value := []interface{}{1, 1, 1, "a simple string", nil, rand.Int63(), []byte{12, 198, 211}}
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_Complex_Map(b *testing.B) {
	set := "put_bench_str_10000"
	value := map[interface{}]interface{}{
		rand.Int63(): rand.Int63(),
		nil:          1,
		"s":          491871,
		15892987:     strings.Repeat("s", 100),
		"s2":         []interface{}{"a simple string", nil, rand.Int63(), []byte{12, 198, 211}},
	}
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_JSON_Map(b *testing.B) {
	set := "put_bench_str_10000"
	value := map[string]interface{}{
		strings.Repeat("a", 16): rand.Int63(),
		strings.Repeat("b", 16): strings.Repeat("s", 100),
		strings.Repeat("c", 16): []interface{}{"a simple string", nil, rand.Int63(), []byte{12, 198, 211}},
		strings.Repeat("d", 16): map[interface{}]interface{}{
			rand.Int63(): rand.Int63(),
			nil:          1,
			"s":          491871,
			15892987:     strings.Repeat("s", 100),
			"s2":         []interface{}{"a simple string", nil, rand.Int63(), []byte{12, 198, 211}},
		},
	}
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}
<commit_msg>Update API in the benchmark Code<commit_after>// Copyright 2013-2017 Aerospike, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aerospike

import (
	"math/rand"
	"runtime"
	"strings"
	"testing"
	// "time"

	_ "net/http/pprof"
	// . "github.com/aerospike/aerospike-client-go"
)

func doGet(set string, value interface{}, b *testing.B) {
	var err error
	policy := NewPolicy()

	dataBuffer := make([]byte, 1024*1024)

	binNames := []string{}
	key, _ := NewKey("test", set, 1000)

	for i := 0; i < b.N; i++ {
		command := newReadCommand(nil, policy, key, binNames)
		command.baseCommand.dataBuffer = dataBuffer
		err = command.writeBuffer(&command)
		if err != nil {
			panic(err)
		}
	}
}

func Benchmark_ReadCommand_________Int64(b *testing.B) {
	set := "put_bench_integer"
	value := rand.Int63()
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_________Int32(b *testing.B) {
	set := "put_bench_integer"
	value := rand.Int31()
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String______1(b *testing.B) {
	set := "put_bench_str_1"
	value := strings.Repeat("s", 1)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String_____10(b *testing.B) {
	set := "put_bench_str_10"
	value := strings.Repeat("s", 10)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String____100(b *testing.B) {
	set := "put_bench_str_100"
	value := strings.Repeat("s", 100)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String___1000(b *testing.B) {
	set := "put_bench_str_1000"
	value := strings.Repeat("s", 1000)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String__10000(b *testing.B) {
	set := "put_bench_str_10000"
	value := strings.Repeat("s", 10000)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_String_100000(b *testing.B) {
	set := "put_bench_str_10000"
	value := strings.Repeat("s", 100000)
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_Complex_Array(b *testing.B) {
	set := "put_bench_str_10000"
	value := []interface{}{1, 1, 1, "a simple string", nil, rand.Int63(), []byte{12, 198, 211}}
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_Complex_Map(b *testing.B) {
	set := "put_bench_str_10000"
	value := map[interface{}]interface{}{
		rand.Int63(): rand.Int63(),
		nil:          1,
		"s":          491871,
		15892987:     strings.Repeat("s", 100),
		"s2":         []interface{}{"a simple string", nil, rand.Int63(), []byte{12, 198, 211}},
	}
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}

func Benchmark_ReadCommand_JSON_Map(b *testing.B) {
	set := "put_bench_str_10000"
	value := map[string]interface{}{
		strings.Repeat("a", 16): rand.Int63(),
		strings.Repeat("b", 16): strings.Repeat("s", 100),
		strings.Repeat("c", 16): []interface{}{"a simple string", nil, rand.Int63(), []byte{12, 198, 211}},
		strings.Repeat("d", 16): map[interface{}]interface{}{
			rand.Int63(): rand.Int63(),
			nil:          1,
			"s":          491871,
			15892987:     strings.Repeat("s", 100),
			"s2":         []interface{}{"a simple string", nil, rand.Int63(), []byte{12, 198, 211}},
		},
	}
	b.N = 1000
	runtime.GC()
	b.ResetTimer()
	doGet(set, value, b)
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"strings"

	"gigawatt-common/pkg/gentle"

	"github.com/cenkalti/backoff"
	"github.com/samuel/go-zookeeper/zk"
)

// CreateP functions similarly to `mkdir -p`.
func CreateP(conn *zk.Conn, path string, data []byte, flags int32, acl []zk.ACL) (zxIds []string, err error) {
	zxIds = []string{}
	pieces := strings.Split(strings.Trim(path, "/"), "/")
	var (
		zxId  string
		soFar string
	)
	for _, piece := range pieces {
		soFar += "/" + piece
		if zxId, err = conn.Create(soFar, data, flags, acl); err != nil && err != zk.ErrNodeExists {
			return
		}
		zxIds = append(zxIds, zxId)
	}
	err = nil // Clear out any potential error state, since if we made it this far we're OK.
	return
}

// MustCreateP will keep trying to create the path until it succeeds.
func MustCreateP(conn *zk.Conn, path string, data []byte, flags int32, acl []zk.ACL, strategy backoff.BackOff) (zxIds []string) {
	var err error
	operation := func() error {
		if zxIds, err = CreateP(conn, path, []byte{}, 0, acl); err != nil {
			return err
		}
		return nil
	}
	gentle.RetryUntilSuccess(fmt.Sprintf("conn=%p MustCreateP", conn), operation, strategy)
	return
}

func MustCreateProtectedEphemeralSequential(conn *zk.Conn, path string, data []byte, acl []zk.ACL, strategy backoff.BackOff) (zxId string) {
	var err error
	operation := func() error {
		if zxId, err = conn.CreateProtectedEphemeralSequential(path, data, acl); err != nil {
			return err
		}
		return nil
	}
	gentle.RetryUntilSuccess(fmt.Sprintf("conn=%p MustCreateProtectedEphemeralSequential", conn), operation, strategy)
	return
}
<commit_msg>Protect against infinite failure due to non-existent ZK path.<commit_after>package util

import (
	"fmt"
	"strings"

	"gigawatt-common/pkg/gentle"

	"github.com/cenkalti/backoff"
	"github.com/samuel/go-zookeeper/zk"
)

// CreateP functions similarly to `mkdir -p`.
func CreateP(conn *zk.Conn, path string, data []byte, flags int32, acl []zk.ACL) (zxIds []string, err error) {
	zxIds = []string{}
	pieces := strings.Split(strings.Trim(path, "/"), "/")
	var (
		zxId  string
		soFar string
	)
	for _, piece := range pieces {
		soFar += "/" + piece
		if zxId, err = conn.Create(soFar, data, flags, acl); err != nil && err != zk.ErrNodeExists {
			return
		}
		zxIds = append(zxIds, zxId)
	}
	err = nil // Clear out any potential error state, since if we made it this far we're OK.
	return
}

// MustCreateP will keep trying to create the path until it succeeds.
func MustCreateP(conn *zk.Conn, path string, data []byte, flags int32, acl []zk.ACL, strategy backoff.BackOff) (zxIds []string) {
	var err error
	operation := func() error {
		if zxIds, err = CreateP(conn, path, []byte{}, 0, acl); err != nil {
			return err
		}
		return nil
	}
	gentle.RetryUntilSuccess(fmt.Sprintf("conn=%p MustCreateP", conn), operation, strategy)
	return
}

func MustCreateProtectedEphemeralSequential(conn *zk.Conn, path string, data []byte, acl []zk.ACL, strategy backoff.BackOff) (zxId string) {
	var err error
	operation := func() error {
		if pieces := strings.Split(path, "/"); len(pieces) > 2 {
			if _, err = CreateP(conn, strings.Join(pieces[0:len(pieces)-1], "/"), []byte{}, 0, acl); err != nil {
				return err
			}
		}
		if zxId, err = conn.CreateProtectedEphemeralSequential(path, data, acl); err != nil {
			return err
		}
		return nil
	}
	gentle.RetryUntilSuccess(fmt.Sprintf("conn=%p MustCreateProtectedEphemeralSequential", conn), operation, strategy)
	return
}
<|endoftext|>


<commit_before><commit_msg>Remove .blanco GTLD (#253)<commit_after><|endoftext|>