<commit_before>// Package simplehstore offers a simple way to use a PostgreSQL database with HSTORE.
// The database back end is interchangeable with Redis (xyproto/simpleredis), BoltDB (xyproto/simplebolt) and
// MariaDB/MySQL (xyproto/simplemaria) by using the interfaces in the xyproto/pinterface package.
package simplehstore

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strings"

	// Using the PostgreSQL database engine
	pq "github.com/lib/pq"
	"github.com/xyproto/env"
)

const (
	VersionString     = "1.6.5"
	defaultStringType = "TEXT"
	defaultPort       = 5432
	encoding          = "UTF8"
)

// Host represents a PostgreSQL database
type Host struct {
	db     *sql.DB
	dbname string

	// If set to true, any UTF-8 string will be let through as it is.
	// Some UTF-8 strings may be unpalatable for PostgreSQL when performing
	// SQL queries. The default is "false".
	rawUTF8 bool
}

// Common for each of the db data structures used here
type dbDatastructure struct {
	host  *Host
	table string
}

var defaultConnectionString = func() string {
	password := env.Str("POSTGRES_PASSWORD")
	s := env.Str("POSTGRES_USER", "postgres")
	if password != "" {
		s += ":" + password
	}
	s += "@127.0.0.1/" // for CI testing
	return s
}()

var (
	// The default "username:password@host:port/database" that the database is running at
	defaultDatabaseName = env.Str("POSTGRES_DB", "postgres")

	// ErrNoAvailableValues is used as an error if an SQL query returns no values
	ErrNoAvailableValues = errors.New("no available values")
	// ErrTooFewResults is used as an error if an SQL query returns too few results
	ErrTooFewResults = errors.New("too few results")

	// Column names
	listCol  = "a_list"
	setCol   = "a_set"
	ownerCol = "owner"
	kvPrefix = "a_kv_"
)

// SetColumnNames can be used to change the column names and prefixes that are used in the PostgreSQL tables.
// The default values are: "a_list", "a_set", "owner" and "a_kv_".
func SetColumnNames(list, set, hashMapOwner, keyValuePrefix string) {
	listCol = list
	setCol = set
	ownerCol = hashMapOwner
	kvPrefix = keyValuePrefix
}

// TestConnection checks if the local database server is up and running
func TestConnection() (err error) {
	return TestConnectionHost(defaultConnectionString)
}

// TestConnectionHost checks if a given database server is up and running.
// connectionString may be on the form "username:password@host:port/database".
// The database name is ignored.
func TestConnectionHost(connectionString string) error {
	newConnectionString, _ := rebuildConnectionString(connectionString, false)
	// Connect to the given host:port
	db, err := sql.Open("postgres", newConnectionString)
	if err != nil {
		return err
	}
	defer db.Close()
	err = db.Ping()
	if Verbose {
		if err != nil {
			log.Println("Ping: failed")
		} else {
			log.Println("Ping: ok")
		}
	}
	return err
}

// TestConnectionHostWithDSN checks if a given database server is up and running.
func TestConnectionHostWithDSN(connectionString string) error {
	// Connect to the given host:port
	db, err := sql.Open("postgres", connectionString)
	if err != nil {
		return err
	}
	defer db.Close()
	err = db.Ping()
	if Verbose {
		if err != nil {
			log.Println("Ping: failed")
		} else {
			log.Println("Ping: ok")
		}
	}
	return err
}

/* --- Host functions --- */

// NewHost sets up a new database connection.
// connectionString may be on the form "username:password@host:port/database".
func NewHost(connectionString string) *Host {
	host, err := NewHost2(connectionString)
	if err != nil {
		log.Fatalln(err)
	}
	return host
}

// NewHost2 sets up a new database connection.
// connectionString may be on the form "username:password@host:port/database".
// An error may be returned.
func NewHost2(connectionString string) (*Host, error) {
	newConnectionString, dbname := rebuildConnectionString(connectionString, true)
	db, err := sql.Open("postgres", newConnectionString)
	if err != nil {
		return nil, fmt.Errorf("could not connect to %s", newConnectionString)
	}
	host := &Host{db, pq.QuoteIdentifier(dbname), false}
	if err := host.Ping(); err != nil {
		return nil, fmt.Errorf("database host does not reply to ping: %s", err)
	}
	if err := host.createDatabase(); err != nil {
		return nil, fmt.Errorf("could not create database %s: %s", host.dbname, err)
	}
	if err := host.useDatabase(); err != nil {
		return nil, fmt.Errorf("could not use database %s: %s", host.dbname, err)
	}
	return host, nil
}

// NewHostWithDSN creates a new database connection with a valid DSN.
func NewHostWithDSN(connectionString string, dbname string) *Host {
	host, err := NewHostWithDSN2(connectionString, dbname)
	if err != nil {
		log.Fatalln(err)
	}
	return host
}

// NewHostWithDSN2 creates a new database connection with a valid DSN.
// An error may be returned.
func NewHostWithDSN2(connectionString string, dbname string) (*Host, error) {
	db, err := sql.Open("postgres", connectionString)
	if err != nil {
		return nil, fmt.Errorf("could not connect to %s", connectionString)
	}
	host := &Host{db, pq.QuoteIdentifier(dbname), false}
	if err := host.Ping(); err != nil {
		return nil, fmt.Errorf("database host does not reply to ping: %s", err)
	}
	if err := host.createDatabase(); err != nil {
		return nil, fmt.Errorf("could not create database %s: %s", host.dbname, err)
	}
	if err := host.useDatabase(); err != nil {
		return nil, fmt.Errorf("could not use database %s: %s", host.dbname, err)
	}
	return host, nil
}

// New sets up a connection to the default (local) database host
func New() *Host {
	connectionString := defaultConnectionString + defaultDatabaseName
	if !strings.HasSuffix(defaultConnectionString, "/") {
		connectionString = defaultConnectionString + "/" + defaultDatabaseName
	}
	return NewHost(connectionString)
}

// SetRawUTF8 can be used to select if the UTF-8 data be unprocessed, and not
// hex encoded and compressed. Unprocessed UTF-8 may be slightly faster,
// but malformed UTF-8 strings can potentially cause problems.
// Encoding the strings before sending them to PostgreSQL is the default.
// Choose the setting that best suits your situation.
func (host *Host) SetRawUTF8(enabled bool) {
	host.rawUTF8 = enabled
}

// SelectDatabase sets a different database name and creates the database if needed.
func (host *Host) SelectDatabase(dbname string) error {
	host.dbname = dbname
	if err := host.createDatabase(); err != nil {
		return err
	}
	return host.useDatabase()
}

// Will create the database if it does not already exist
func (host *Host) createDatabase() error {
	if _, err := host.db.Exec(fmt.Sprintf("CREATE DATABASE %s WITH ENCODING '%s'", host.dbname, encoding)); err != nil {
		if !strings.HasSuffix(err.Error(), "already exists") {
			return err
		}
	}
	return nil
}

// Use the host.dbname database
func (host *Host) useDatabase() error {
	if Verbose {
		log.Println("Using database " + host.dbname)
	}
	return nil
}

// Database returns the underlying *sql.DB database struct
func (host *Host) Database() *sql.DB {
	return host.db
}

// Close the connection
func (host *Host) Close() {
	host.db.Close()
}

// Ping the host
func (host *Host) Ping() error {
	return host.db.Ping()
}
<commit_msg>Add a comment<commit_after>// Package simplehstore offers a simple way to use a PostgreSQL database with HSTORE.
// The database back end is interchangeable with Redis (xyproto/simpleredis), BoltDB (xyproto/simplebolt) and
// MariaDB/MySQL (xyproto/simplemaria) by using the interfaces in the xyproto/pinterface package.
package simplehstore

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strings"

	// Using the PostgreSQL database engine
	pq "github.com/lib/pq"
	"github.com/xyproto/env"
)

const (
	// VersionString is the current version of simplehstore.
	VersionString = "1.6.5"

	defaultStringType = "TEXT"
	defaultPort       = 5432
	encoding          = "UTF8"
)

// Host represents a PostgreSQL database
type Host struct {
	db     *sql.DB
	dbname string

	// If set to true, any UTF-8 string will be let through as it is.
	// Some UTF-8 strings may be unpalatable for PostgreSQL when performing
	// SQL queries. The default is "false".
	rawUTF8 bool
}

// Common for each of the db data structures used here
type dbDatastructure struct {
	host  *Host
	table string
}

var defaultConnectionString = func() string {
	password := env.Str("POSTGRES_PASSWORD")
	s := env.Str("POSTGRES_USER", "postgres")
	if password != "" {
		s += ":" + password
	}
	s += "@127.0.0.1/" // for CI testing
	return s
}()

var (
	// The default "username:password@host:port/database" that the database is running at
	defaultDatabaseName = env.Str("POSTGRES_DB", "postgres")

	// ErrNoAvailableValues is used as an error if an SQL query returns no values
	ErrNoAvailableValues = errors.New("no available values")
	// ErrTooFewResults is used as an error if an SQL query returns too few results
	ErrTooFewResults = errors.New("too few results")

	// Column names
	listCol  = "a_list"
	setCol   = "a_set"
	ownerCol = "owner"
	kvPrefix = "a_kv_"
)

// SetColumnNames can be used to change the column names and prefixes that are used in the PostgreSQL tables.
// The default values are: "a_list", "a_set", "owner" and "a_kv_".
func SetColumnNames(list, set, hashMapOwner, keyValuePrefix string) {
	listCol = list
	setCol = set
	ownerCol = hashMapOwner
	kvPrefix = keyValuePrefix
}

// TestConnection checks if the local database server is up and running
func TestConnection() (err error) {
	return TestConnectionHost(defaultConnectionString)
}

// TestConnectionHost checks if a given database server is up and running.
// connectionString may be on the form "username:password@host:port/database".
// The database name is ignored.
func TestConnectionHost(connectionString string) error {
	newConnectionString, _ := rebuildConnectionString(connectionString, false)
	// Connect to the given host:port
	db, err := sql.Open("postgres", newConnectionString)
	if err != nil {
		return err
	}
	defer db.Close()
	err = db.Ping()
	if Verbose {
		if err != nil {
			log.Println("Ping: failed")
		} else {
			log.Println("Ping: ok")
		}
	}
	return err
}

// TestConnectionHostWithDSN checks if a given database server is up and running.
func TestConnectionHostWithDSN(connectionString string) error {
	// Connect to the given host:port
	db, err := sql.Open("postgres", connectionString)
	if err != nil {
		return err
	}
	defer db.Close()
	err = db.Ping()
	if Verbose {
		if err != nil {
			log.Println("Ping: failed")
		} else {
			log.Println("Ping: ok")
		}
	}
	return err
}

/* --- Host functions --- */

// NewHost sets up a new database connection.
// connectionString may be on the form "username:password@host:port/database".
func NewHost(connectionString string) *Host {
	host, err := NewHost2(connectionString)
	if err != nil {
		log.Fatalln(err)
	}
	return host
}

// NewHost2 sets up a new database connection.
// connectionString may be on the form "username:password@host:port/database".
// An error may be returned.
func NewHost2(connectionString string) (*Host, error) {
	newConnectionString, dbname := rebuildConnectionString(connectionString, true)
	db, err := sql.Open("postgres", newConnectionString)
	if err != nil {
		return nil, fmt.Errorf("could not connect to %s", newConnectionString)
	}
	host := &Host{db, pq.QuoteIdentifier(dbname), false}
	if err := host.Ping(); err != nil {
		return nil, fmt.Errorf("database host does not reply to ping: %s", err)
	}
	if err := host.createDatabase(); err != nil {
		return nil, fmt.Errorf("could not create database %s: %s", host.dbname, err)
	}
	if err := host.useDatabase(); err != nil {
		return nil, fmt.Errorf("could not use database %s: %s", host.dbname, err)
	}
	return host, nil
}

// NewHostWithDSN creates a new database connection with a valid DSN.
func NewHostWithDSN(connectionString string, dbname string) *Host {
	host, err := NewHostWithDSN2(connectionString, dbname)
	if err != nil {
		log.Fatalln(err)
	}
	return host
}

// NewHostWithDSN2 creates a new database connection with a valid DSN.
// An error may be returned.
func NewHostWithDSN2(connectionString string, dbname string) (*Host, error) {
	db, err := sql.Open("postgres", connectionString)
	if err != nil {
		return nil, fmt.Errorf("could not connect to %s", connectionString)
	}
	host := &Host{db, pq.QuoteIdentifier(dbname), false}
	if err := host.Ping(); err != nil {
		return nil, fmt.Errorf("database host does not reply to ping: %s", err)
	}
	if err := host.createDatabase(); err != nil {
		return nil, fmt.Errorf("could not create database %s: %s", host.dbname, err)
	}
	if err := host.useDatabase(); err != nil {
		return nil, fmt.Errorf("could not use database %s: %s", host.dbname, err)
	}
	return host, nil
}

// New sets up a connection to the default (local) database host
func New() *Host {
	connectionString := defaultConnectionString + defaultDatabaseName
	if !strings.HasSuffix(defaultConnectionString, "/") {
		connectionString = defaultConnectionString + "/" + defaultDatabaseName
	}
	return NewHost(connectionString)
}

// SetRawUTF8 can be used to select if the UTF-8 data be unprocessed, and not
// hex encoded and compressed. Unprocessed UTF-8 may be slightly faster,
// but malformed UTF-8 strings can potentially cause problems.
// Encoding the strings before sending them to PostgreSQL is the default.
// Choose the setting that best suits your situation.
func (host *Host) SetRawUTF8(enabled bool) {
	host.rawUTF8 = enabled
}

// SelectDatabase sets a different database name and creates the database if needed.
func (host *Host) SelectDatabase(dbname string) error {
	host.dbname = dbname
	if err := host.createDatabase(); err != nil {
		return err
	}
	return host.useDatabase()
}

// Will create the database if it does not already exist
func (host *Host) createDatabase() error {
	if _, err := host.db.Exec(fmt.Sprintf("CREATE DATABASE %s WITH ENCODING '%s'", host.dbname, encoding)); err != nil {
		if !strings.HasSuffix(err.Error(), "already exists") {
			return err
		}
	}
	return nil
}

// Use the host.dbname database
func (host *Host) useDatabase() error {
	if Verbose {
		log.Println("Using database " + host.dbname)
	}
	return nil
}

// Database returns the underlying *sql.DB database struct
func (host *Host) Database() *sql.DB {
	return host.db
}

// Close the connection
func (host *Host) Close() {
	host.db.Close()
}

// Ping the host
func (host *Host) Ping() error {
	return host.db.Ping()
}
<|endoftext|>


<commit_before>package httplog

import (
	"fmt"
	"net/http"
	"strings"
	"time"
)

/* kindof a common log type output */
func LogRequest(r *http.Request, statusCode int) {
	var addr string
	var user_agent string

	user_agent = ""
	addr = RealIP(r)

	for k, v := range r.Header {
		if k == "User-Agent" {
			user_agent = strings.Join(v, " ")
		}
		if k == "X-Forwarded-For" {
			addr = strings.Join(v, " ")
		}
	}

	fmt.Printf("%s - - [%s] \"%s %s\" \"%s\" %d %d\n",
		addr,
		time.Now(),
		r.Method,
		r.URL.String(),
		user_agent,
		statusCode,
		r.ContentLength)
}

func RealIP(r *http.Request) (ip string) {
	ip = r.RemoteAddr

	port_pos := strings.LastIndex(ip, ":")
	if port_pos != -1 {
		ip = ip[0:port_pos]
	}

	for k, v := range r.Header {
		if k == "X-Forwarded-For" {
			ip = strings.Join(v, " ")
		}
	}

	return ip
}
<commit_msg>adding a debugging Header printout<commit_after>package httplog

import (
	"fmt"
	"net/http"
	"strings"
	"time"
)

// for debugging request headers
func LogHeaders(r *http.Request) {
	fmt.Printf("HEADERS:\n")
	for k, v := range r.Header {
		fmt.Printf("\t%s\n", k)
		for i, _ := range v {
			fmt.Printf("\t\t%s\n", v[i])
		}
	}
}

/* kindof a common log type output */
func LogRequest(r *http.Request, statusCode int) {
	var addr string
	var user_agent string

	user_agent = ""
	addr = RealIP(r)

	for k, v := range r.Header {
		if k == "User-Agent" {
			user_agent = strings.Join(v, " ")
		}
		if k == "X-Forwarded-For" {
			addr = strings.Join(v, " ")
		}
	}

	fmt.Printf("%s - - [%s] \"%s %s\" \"%s\" %d %d\n",
		addr,
		time.Now(),
		r.Method,
		r.URL.String(),
		user_agent,
		statusCode,
		r.ContentLength)
}

func RealIP(r *http.Request) (ip string) {
	ip = r.RemoteAddr

	port_pos := strings.LastIndex(ip, ":")
	if port_pos != -1 {
		ip = ip[0:port_pos]
	}

	for k, v := range r.Header {
		if k == "X-Forwarded-For" {
			ip = strings.Join(v, " ")
		}
	}

	return ip
}
<|endoftext|>


<commit_before>package cdn

import (
	"context"
	"fmt"
	"net/http"

	"github.com/go-gorp/gorp"
	"github.com/gorilla/mux"
	"github.com/rockbears/log"

	"github.com/ovh/cds/engine/api"
	"github.com/ovh/cds/engine/cache"
	"github.com/ovh/cds/engine/cdn/item"
	"github.com/ovh/cds/engine/cdn/lru"
	"github.com/ovh/cds/engine/cdn/storage"
	"github.com/ovh/cds/engine/cdn/storage/cds"
	_ "github.com/ovh/cds/engine/cdn/storage/local"
	_ "github.com/ovh/cds/engine/cdn/storage/redis"
	_ "github.com/ovh/cds/engine/cdn/storage/swift"
	"github.com/ovh/cds/engine/database"
	"github.com/ovh/cds/engine/gorpmapper"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/cdsclient"
)

const (
	defaultLruSize            = 128 * 1024 * 1024 // 128Mb
	defaultStepMaxSize        = 15 * 1024 * 1024  // 15Mb
	defaultStepLinesRateLimit = 1800
	defaultGlobalTCPRateLimit = 2 * 1024 * 1024 // 2Mb
)

// New returns a new service
func New() *Service {
	s := new(Service)
	s.GoRoutines = sdk.NewGoRoutines()
	s.Router = &api.Router{
		Mux: mux.NewRouter(),
	}

	return s
}

func (s *Service) Init(config interface{}) (cdsclient.ServiceConfig, error) {
	var cfg cdsclient.ServiceConfig
	sConfig, ok := config.(Configuration)
	if !ok {
		return cfg, sdk.WithStack(fmt.Errorf("invalid CDN service configuration"))
	}

	cfg.Host = sConfig.API.HTTP.URL
	cfg.Token = sConfig.API.Token
	cfg.InsecureSkipVerifyTLS = sConfig.API.HTTP.Insecure
	cfg.RequestSecondsTimeout = sConfig.API.RequestTimeout

	return cfg, nil
}

// ApplyConfiguration apply an object of type CDN.Configuration after checking it
func (s *Service) ApplyConfiguration(config interface{}) error {
	if err := s.CheckConfiguration(config); err != nil {
		return err
	}
	var ok bool
	s.Cfg, ok = config.(Configuration)
	if !ok {
		return fmt.Errorf("invalid configuration")
	}
	s.ServiceName = s.Cfg.Name
	s.ServiceType = sdk.TypeCDN
	s.HTTPURL = s.Cfg.URL
	s.MaxHeartbeatFailures = s.Cfg.API.MaxHeartbeatFailures

	if s.Cfg.Cache.LruSize == 0 {
		s.Cfg.Cache.LruSize = defaultLruSize
	}
	if s.Cfg.Log.StepMaxSize == 0 {
		s.Cfg.Log.StepMaxSize = defaultStepMaxSize
	}
	if s.Cfg.Log.StepLinesRateLimit == 0 {
		s.Cfg.Log.StepLinesRateLimit = defaultStepLinesRateLimit
	}
	if s.Cfg.TCP.GlobalTCPRateLimit == 0 {
		s.Cfg.TCP.GlobalTCPRateLimit = defaultGlobalTCPRateLimit
	}

	return nil
}

// CheckConfiguration checks the validity of the configuration object
func (s *Service) CheckConfiguration(config interface{}) error {
	sConfig, ok := config.(Configuration)
	if !ok {
		return fmt.Errorf("invalid configuration")
	}

	if sConfig.URL == "" {
		return fmt.Errorf("your CDS configuration seems to be empty. Please use environment variables, file or Consul to set your configuration")
	}
	if sConfig.Name == "" {
		return fmt.Errorf("please enter a name in your CDN configuration")
	}

	return nil
}

func (s *Service) Start(c context.Context, cfg cdsclient.ServiceConfig) error {
	if err := s.Common.Start(c, cfg); err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(c)
	defer cancel()

	var err error
	log.Info(ctx, "Initializing redis cache on %s...", s.Cfg.Cache.Redis.Host)
	s.Cache, err = cache.New(s.Cfg.Cache.Redis.Host, s.Cfg.Cache.Redis.Password, s.Cfg.Cache.TTL)
	if err != nil {
		return sdk.WrapError(err, "cannot connect to redis instance")
	}

	if s.Cfg.EnableLogProcessing {
		log.Info(ctx, "Initializing database connection...")
		// Intialize database
		s.DBConnectionFactory, err = database.Init(
			ctx,
			s.Cfg.Database.User,
			s.Cfg.Database.Role,
			s.Cfg.Database.Password,
			s.Cfg.Database.Name,
			s.Cfg.Database.Schema,
			s.Cfg.Database.Host,
			s.Cfg.Database.Port,
			s.Cfg.Database.SSLMode,
			s.Cfg.Database.ConnectTimeout,
			s.Cfg.Database.Timeout,
			s.Cfg.Database.MaxConn)
		if err != nil {
			return sdk.WrapError(err, "cannot connect to database")
		}

		log.Info(ctx, "Setting up database keys...")
		s.Mapper = gorpmapper.New()
		encryptionKeyConfig := s.Cfg.Database.EncryptionKey.GetKeys(gorpmapper.KeyEcnryptionIdentifier)
		signatureKeyConfig := s.Cfg.Database.SignatureKey.GetKeys(gorpmapper.KeySignIdentifier)
		if err := s.Mapper.ConfigureKeys(&signatureKeyConfig, &encryptionKeyConfig); err != nil {
			return sdk.WrapError(err, "cannot setup database keys")
		}

		// Init dao packages
		item.InitDBMapping(s.Mapper)
		storage.InitDBMapping(s.Mapper)
	}

	return nil
}

// Serve will start the http api server
func (s *Service) Serve(c context.Context) error {
	ctx, cancel := context.WithCancel(c)
	defer cancel()

	if s.Cfg.EnableLogProcessing {
		var err error
		s.LogCache, err = lru.NewRedisLRU(s.mustDBWithCtx(ctx), s.Cfg.Cache.LruSize, s.Cfg.Cache.Redis.Host, s.Cfg.Cache.Redis.Password)
		if err != nil {
			return sdk.WrapError(err, "cannot connect to redis instance for lru")
		}

		// Init storage units
		s.Units, err = storage.Init(ctx, s.Mapper, s.Cache, s.mustDBWithCtx(ctx), s.GoRoutines, s.Cfg.Units)
		if err != nil {
			return err
		}

		s.Units.Start(ctx, s.GoRoutines)

		s.GoRoutines.Run(ctx, "service.cdn-gc-items", func(ctx context.Context) {
			s.itemsGC(ctx)
		})
		s.GoRoutines.Run(ctx, "service.cdn-purge-items", func(ctx context.Context) {
			s.itemPurge(ctx)
		})

		// Start CDS Backend migration
		for _, st := range s.Units.Storages {
			cdsStorage, ok := st.(*cds.CDS)
			if !ok {
				continue
			}
			s.GoRoutines.Exec(ctx, "cdn-cds-backend-migration", func(ctx context.Context) {
				if err := s.listenCDSSync(ctx, cdsStorage); err != nil {
					log.Error(ctx, "unable to listen pubsub for cds sync: %v", err)
				}
			})
		}

		s.GoRoutines.Run(ctx, "service.log-cache-eviction", func(ctx context.Context) {
			s.LogCache.Evict(ctx)
		})
	}

	if err := s.initMetrics(ctx); err != nil {
		return err
	}

	if err := s.runTCPLogServer(ctx); err != nil {
		return err
	}

	log.Info(ctx, "Initializing HTTP router")
	s.initRouter(ctx)
	if err := s.initWebsocket(); err != nil {
		return err
	}
	server := &http.Server{
		Addr:           fmt.Sprintf("%s:%d", s.Cfg.HTTP.Addr, s.Cfg.HTTP.Port),
		Handler:        s.Router.Mux,
		MaxHeaderBytes: 1 << 20,
	}

	// Gracefully shutdown the http server
	s.GoRoutines.Exec(ctx, "service.httpserver-shutdown", func(ctx context.Context) {
		<-ctx.Done()
		log.Info(ctx, "CDN> Shutdown HTTP Server")
		_ = server.Shutdown(ctx)
	})

	// Start the http server
	log.Info(ctx, "CDN> Starting HTTP Server on port %d", s.Cfg.HTTP.Port)
	if err := server.ListenAndServe(); err != nil {
		log.Fatal(ctx, "CDN> Cannot start cds-cdn: %v", err)
	}
	return ctx.Err()
}

func (s *Service) mustDBWithCtx(ctx context.Context) *gorp.DbMap {
	db := s.DBConnectionFactory.GetDBMap(s.Mapper)()
	db = db.WithContext(ctx).(*gorp.DbMap)
	if db == nil {
		panic(fmt.Errorf("database unavailable"))
	}
	return db
}
<commit_msg>fix(cdn): start units and lru in service start func (#5663)<commit_after>package cdn

import (
	"context"
	"fmt"
	"net/http"

	"github.com/go-gorp/gorp"
	"github.com/gorilla/mux"
	"github.com/rockbears/log"

	"github.com/ovh/cds/engine/api"
	"github.com/ovh/cds/engine/cache"
	"github.com/ovh/cds/engine/cdn/item"
	"github.com/ovh/cds/engine/cdn/lru"
	"github.com/ovh/cds/engine/cdn/storage"
	"github.com/ovh/cds/engine/cdn/storage/cds"
	_ "github.com/ovh/cds/engine/cdn/storage/local"
	_ "github.com/ovh/cds/engine/cdn/storage/redis"
	_ "github.com/ovh/cds/engine/cdn/storage/swift"
	"github.com/ovh/cds/engine/database"
	"github.com/ovh/cds/engine/gorpmapper"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/cdsclient"
)

const (
	defaultLruSize            = 128 * 1024 * 1024 // 128Mb
	defaultStepMaxSize        = 15 * 1024 * 1024  // 15Mb
	defaultStepLinesRateLimit = 1800
	defaultGlobalTCPRateLimit = 2 * 1024 * 1024 // 2Mb
)

// New returns a new service
func New() *Service {
	s := new(Service)
	s.GoRoutines = sdk.NewGoRoutines()
	s.Router = &api.Router{
		Mux: mux.NewRouter(),
	}

	return s
}

func (s *Service) Init(config interface{}) (cdsclient.ServiceConfig, error) {
	var cfg cdsclient.ServiceConfig
	sConfig, ok := config.(Configuration)
	if !ok {
		return cfg, sdk.WithStack(fmt.Errorf("invalid CDN service configuration"))
	}

	cfg.Host = sConfig.API.HTTP.URL
	cfg.Token = sConfig.API.Token
	cfg.InsecureSkipVerifyTLS = sConfig.API.HTTP.Insecure
	cfg.RequestSecondsTimeout = sConfig.API.RequestTimeout

	return cfg, nil
}

// ApplyConfiguration apply an object of type CDN.Configuration after checking it
func (s *Service) ApplyConfiguration(config interface{}) error {
	if err := s.CheckConfiguration(config); err != nil {
		return err
	}
	var ok bool
	s.Cfg, ok = config.(Configuration)
	if !ok {
		return fmt.Errorf("invalid configuration")
	}
	s.ServiceName = s.Cfg.Name
	s.ServiceType = sdk.TypeCDN
	s.HTTPURL = s.Cfg.URL
	s.MaxHeartbeatFailures = s.Cfg.API.MaxHeartbeatFailures

	if s.Cfg.Cache.LruSize == 0 {
		s.Cfg.Cache.LruSize = defaultLruSize
	}
	if s.Cfg.Log.StepMaxSize == 0 {
		s.Cfg.Log.StepMaxSize = defaultStepMaxSize
	}
	if s.Cfg.Log.StepLinesRateLimit == 0 {
		s.Cfg.Log.StepLinesRateLimit = defaultStepLinesRateLimit
	}
	if s.Cfg.TCP.GlobalTCPRateLimit == 0 {
		s.Cfg.TCP.GlobalTCPRateLimit = defaultGlobalTCPRateLimit
	}

	return nil
}

// CheckConfiguration checks the validity of the configuration object
func (s *Service) CheckConfiguration(config interface{}) error {
	sConfig, ok := config.(Configuration)
	if !ok {
		return fmt.Errorf("invalid configuration")
	}

	if sConfig.URL == "" {
		return fmt.Errorf("your CDS configuration seems to be empty. Please use environment variables, file or Consul to set your configuration")
	}
	if sConfig.Name == "" {
		return fmt.Errorf("please enter a name in your CDN configuration")
	}

	return nil
}

func (s *Service) Start(ctx context.Context, cfg cdsclient.ServiceConfig) error {
	if err := s.Common.Start(ctx, cfg); err != nil {
		return err
	}

	var err error
	log.Info(ctx, "Initializing redis cache on %s...", s.Cfg.Cache.Redis.Host)
	s.Cache, err = cache.New(s.Cfg.Cache.Redis.Host, s.Cfg.Cache.Redis.Password, s.Cfg.Cache.TTL)
	if err != nil {
		return sdk.WrapError(err, "cannot connect to redis instance")
	}

	if s.Cfg.EnableLogProcessing {
		log.Info(ctx, "Initializing database connection...")
		// Intialize database
		s.DBConnectionFactory, err = database.Init(
			ctx,
			s.Cfg.Database.User,
			s.Cfg.Database.Role,
			s.Cfg.Database.Password,
			s.Cfg.Database.Name,
			s.Cfg.Database.Schema,
			s.Cfg.Database.Host,
			s.Cfg.Database.Port,
			s.Cfg.Database.SSLMode,
			s.Cfg.Database.ConnectTimeout,
			s.Cfg.Database.Timeout,
			s.Cfg.Database.MaxConn)
		if err != nil {
			return sdk.WrapError(err, "cannot connect to database")
		}

		log.Info(ctx, "Setting up database keys...")
		s.Mapper = gorpmapper.New()
		encryptionKeyConfig := s.Cfg.Database.EncryptionKey.GetKeys(gorpmapper.KeyEcnryptionIdentifier)
		signatureKeyConfig := s.Cfg.Database.SignatureKey.GetKeys(gorpmapper.KeySignIdentifier)
		if err := s.Mapper.ConfigureKeys(&signatureKeyConfig, &encryptionKeyConfig); err != nil {
			return sdk.WrapError(err, "cannot setup database keys")
		}

		// Init dao packages
		item.InitDBMapping(s.Mapper)
		storage.InitDBMapping(s.Mapper)

		log.Info(ctx, "Initializing lru connection...")
		s.LogCache, err = lru.NewRedisLRU(s.mustDBWithCtx(ctx), s.Cfg.Cache.LruSize, s.Cfg.Cache.Redis.Host, s.Cfg.Cache.Redis.Password)
		if err != nil {
			return sdk.WrapError(err, "cannot connect to redis instance for lru")
		}

		var err error
		// Init storage units
		s.Units, err = storage.Init(ctx, s.Mapper, s.Cache, s.mustDBWithCtx(ctx), s.GoRoutines, s.Cfg.Units)
		if err != nil {
			return err
		}

		s.Units.Start(ctx, s.GoRoutines)

		s.GoRoutines.Run(ctx, "service.cdn-gc-items", func(ctx context.Context) {
			s.itemsGC(ctx)
		})
		s.GoRoutines.Run(ctx, "service.cdn-purge-items", func(ctx context.Context) {
			s.itemPurge(ctx)
		})

		// Start CDS Backend migration
		for _, st := range s.Units.Storages {
			cdsStorage, ok := st.(*cds.CDS)
			if !ok {
				continue
			}
			s.GoRoutines.Exec(ctx, "cdn-cds-backend-migration", func(ctx context.Context) {
				if err := s.listenCDSSync(ctx, cdsStorage); err != nil {
					log.Error(ctx, "unable to listen pubsub for cds sync: %v", err)
				}
			})
		}

		s.GoRoutines.Run(ctx, "service.log-cache-eviction", func(ctx context.Context) {
			s.LogCache.Evict(ctx)
		})
	}

	return nil
}

// Serve will start the http api server
func (s *Service) Serve(c context.Context) error {
	ctx, cancel := context.WithCancel(c)
	defer cancel()

	if err := s.initMetrics(ctx); err != nil {
		return err
	}

	if err := s.runTCPLogServer(ctx); err != nil {
		return err
	}

	log.Info(ctx, "Initializing HTTP router")
	s.initRouter(ctx)
	if err := s.initWebsocket(); err != nil {
		return err
	}
	server := &http.Server{
		Addr:           fmt.Sprintf("%s:%d", s.Cfg.HTTP.Addr, s.Cfg.HTTP.Port),
		Handler:        s.Router.Mux,
		MaxHeaderBytes: 1 << 20,
	}

	// Gracefully shutdown the http server
	s.GoRoutines.Exec(ctx, "service.httpserver-shutdown", func(ctx context.Context) {
		<-ctx.Done()
		log.Info(ctx, "CDN> Shutdown HTTP Server")
		_ = server.Shutdown(ctx)
	})

	// Start the http server
	log.Info(ctx, "CDN> Starting HTTP Server on port %d", s.Cfg.HTTP.Port)
	if err := server.ListenAndServe(); err != nil {
		log.Fatal(ctx, "CDN> Cannot start cds-cdn: %v", err)
	}
	return ctx.Err()
}

func (s *Service) mustDBWithCtx(ctx context.Context) *gorp.DbMap {
	db := s.DBConnectionFactory.GetDBMap(s.Mapper)()
	db = db.WithContext(ctx).(*gorp.DbMap)
	if db == nil {
		panic(fmt.Errorf("database unavailable"))
	}
	return db
}
<|endoftext|>


<commit_before>package sv_file_content

import (
	"github.com/watermint/toolbox/domain/model/mo_file"
	"github.com/watermint/toolbox/domain/model/mo_path"
	"github.com/watermint/toolbox/infra/api/api_context"
	"github.com/watermint/toolbox/infra/api/api_util"
	"go.uber.org/zap"
	"io"
	"os"
	"path/filepath"
)

type Upload interface {
	Add(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error)
	Overwrite(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error)
	Update(destPath mo_path.Path, filePath string, revision string) (entry mo_file.Entry, err error)
}

type UploadOpt func(o *UploadOpts) *UploadOpts
type UploadOpts struct {
	ChunkSize int64
	Mute      bool
}

const (
	DefaultChunkSize = 150 * 1_048_576
)

func NewUpload(ctx api_context.Context, opts ...UploadOpt) Upload {
	uo := &UploadOpts{
		ChunkSize: DefaultChunkSize,
		Mute:      false,
	}
	for _, o := range opts {
		o(uo)
	}
	return &uploadImpl{
		ctx: ctx,
		uo:  uo,
	}
}
func ChunkSizeKb(chunkSizeKb int) UploadOpt {
	return func(o *UploadOpts) *UploadOpts {
		o.ChunkSize = int64(chunkSizeKb * 1024)
		return o
	}
}

type uploadImpl struct {
	ctx api_context.Context
	uo  *UploadOpts
}

func (z *uploadImpl) Add(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error) {
	return z.upload(destPath, filePath, "add", "")
}

func (z *uploadImpl) Overwrite(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error) {
	return z.upload(destPath, filePath, "overwrite", "")
}

func (z *uploadImpl) Update(destPath mo_path.Path, filePath string, revision string) (entry mo_file.Entry, err error) {
	return z.upload(destPath, filePath, "update", revision)
}

func (z *uploadImpl) upload(destPath mo_path.Path, filePath string, mode string, revision string) (entry mo_file.Entry, err error) {
	info, err := os.Lstat(filePath)
	if err != nil {
		return nil, err
	}
	if info.Size() < z.uo.ChunkSize {
		return z.uploadSingle(info, destPath, filePath, mode, revision)
	} else {
		return z.uploadChunked(info, destPath, filePath, mode, revision)
	}
}

type UploadParamMode struct {
	Tag    string `json:".tag"`
	Update string `json:"update,omitempty"`
}

type UploadParams struct {
	Path           string           `json:"path"`
	Mode           *UploadParamMode `json:"mode"`
	Mute           bool             `json:"mute"`
	ClientModified string           `json:"client_modified"`
	Autorename     bool             `json:"autorename"`
}

func UploadPath(destPath mo_path.Path, f os.FileInfo) mo_path.Path {
	return destPath.ChildPath(filepath.Base(f.Name()))
}

func (z *uploadImpl) makeParams(info os.FileInfo, destPath mo_path.Path, mode string, revision string) *UploadParams {
	upm := &UploadParamMode{
		Tag:    mode,
		Update: "",
	}
	up := &UploadParams{
		Path:           UploadPath(destPath, info).Path(),
		Mode:           upm,
		Mute:           false,
		ClientModified: api_util.RebaseAsString(info.ModTime()),
	}
	switch mode {
	case "update":
		upm.Update = revision
	case "add":
		up.Autorename = true
	}
	return up
}

func (z *uploadImpl) uploadSingle(info os.FileInfo, destPath mo_path.Path, filePath string, mode string, revision string) (entry mo_file.Entry, err error) {
	l := z.ctx.Log().With(zap.String("filePath", filePath), zap.Int64("size", info.Size()))
	l.Debug("Uploading file")

	r, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	res, err := z.ctx.Upload("files/upload", r).
		Param(z.makeParams(info, destPath, mode, revision)).Call()
	if err != nil {
		return nil, err
	}
	entry = &mo_file.Metadata{}
	if err := res.Model(entry); err != nil {
		return nil, err
	}
	return entry, nil
}

func (z *uploadImpl) uploadChunked(info os.FileInfo, destPath mo_path.Path, filePath string, mode string, revision string) (entry mo_file.Entry, err error) {
	l := z.ctx.Log().With(zap.String("filePath", filePath), zap.Int64("size", info.Size()))

	total := info.Size()
	var written int64
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}

	type SessionId struct {
		SessionId string `json:"session_id"`
	}
	type CursorInfo struct {
		SessionId string `json:"session_id"`
		Offset    int64  `json:"offset"`
	}
	type AppendInfo struct {
		Cursor *CursorInfo `json:"cursor"`
	}
	type CommitInfo struct {
		Cursor *CursorInfo   `json:"cursor"`
		Commit *UploadParams `json:"commit"`
	}

	l.Debug("Upload session start")
	r := io.LimitReader(f, z.uo.ChunkSize)
	res, err := z.ctx.Upload("files/upload_session/start", r).Call()
	if err != nil {
		return nil, err
	}
	sid := &SessionId{}
	if j, err := res.Json(); err != nil {
		return nil, err
	} else {
		sid.SessionId = j.Get("session_id").String()
	}
	written += z.uo.ChunkSize
	l = l.With(zap.String("sessionId", sid.SessionId))

	for (total - written) > z.uo.ChunkSize {
		l.Debug("Append chunk", zap.Int64("written", written))
		ai := &AppendInfo{
			Cursor: &CursorInfo{
				SessionId: sid.SessionId,
				Offset:    written,
			},
		}
		r = io.LimitReader(f, z.uo.ChunkSize)
		_, err := z.ctx.Upload("files/upload_session/append_v2", r).Param(ai).Call()
		if err != nil {
			return nil, err
		}
		written += z.uo.ChunkSize
	}

	l.Debug("Finish")
	ci := &CommitInfo{
		Cursor: &CursorInfo{
			SessionId: sid.SessionId,
			Offset:    written,
		},
		Commit: z.makeParams(info, destPath, mode, revision),
	}
	res, err = z.ctx.Upload("files/upload_session/finish", f).Param(ci).Call()
	if err != nil {
		return nil, err
	}
	entry = &mo_file.Metadata{}
	if err := res.Model(entry); err != nil {
		return nil, err
	}
	return entry, nil
}
<commit_msg>fix #220 : verify chunk size<commit_after>package sv_file_content

import (
	"github.com/watermint/toolbox/domain/model/mo_file"
	"github.com/watermint/toolbox/domain/model/mo_path"
	"github.com/watermint/toolbox/infra/api/api_context"
	"github.com/watermint/toolbox/infra/api/api_util"
	"go.uber.org/zap"
	"io"
	"os"
	"path/filepath"
	"sync"
)

type Upload interface {
	Add(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error)
	Overwrite(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error)
	Update(destPath mo_path.Path, filePath string, revision string) (entry mo_file.Entry, err error)
}

type UploadOpt func(o *UploadOpts) *UploadOpts
type UploadOpts struct {
	ChunkSize int64
	Mute      bool
}

const (
	MaxChunkSize     = 150 * 1_048_576 // 150MB
	DefaultChunkSize = MaxChunkSize
)

var (
	warnExceededChunkSize = sync.Once{}
)

func NewUpload(ctx api_context.Context, opts ...UploadOpt) Upload {
	uo := &UploadOpts{
		ChunkSize: DefaultChunkSize,
		Mute:      false,
	}
	for _, o := range opts {
		o(uo)
	}
	if uo.ChunkSize > MaxChunkSize {
		warnExceededChunkSize.Do(func() {
			ctx.Log().Warn("Chunk size exceed maximum size, chunk size will be adjusted to maximum size", zap.Int64("givenChunkSize", uo.ChunkSize))
		})
		uo.ChunkSize = MaxChunkSize
	}
	return &uploadImpl{
		ctx: ctx,
		uo:  uo,
	}
}
func ChunkSizeKb(chunkSizeKb int) UploadOpt {
	return func(o *UploadOpts) *UploadOpts {
		o.ChunkSize = int64(chunkSizeKb * 1024)
		return o
	}
}

type uploadImpl struct {
	ctx api_context.Context
	uo  *UploadOpts
}

func (z *uploadImpl) Add(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error) {
	return z.upload(destPath, filePath, "add", "")
}

func (z *uploadImpl) Overwrite(destPath mo_path.Path, filePath string) (entry mo_file.Entry, err error) {
	return z.upload(destPath, filePath, "overwrite", "")
}

func (z *uploadImpl) Update(destPath mo_path.Path, filePath string, revision string) (entry mo_file.Entry, err error) {
	return z.upload(destPath, filePath, "update", revision)
}

func (z *uploadImpl) upload(destPath mo_path.Path, filePath string, mode string, revision string) (entry mo_file.Entry, err error) {
	info, err := os.Lstat(filePath)
	if err != nil {
		return nil, err
	}
	if info.Size() < z.uo.ChunkSize {
		return z.uploadSingle(info, destPath, filePath, mode, revision)
	} else {
		return z.uploadChunked(info, destPath, filePath, mode, revision)
	}
}

type UploadParamMode struct {
	Tag    string `json:".tag"`
	Update string `json:"update,omitempty"`
}

type UploadParams struct {
	Path           string           `json:"path"`
	Mode           *UploadParamMode `json:"mode"`
	Mute           bool             `json:"mute"`
	ClientModified string           `json:"client_modified"`
	Autorename     bool             `json:"autorename"`
}

func UploadPath(destPath mo_path.Path, f os.FileInfo) mo_path.Path {
	return destPath.ChildPath(filepath.Base(f.Name()))
}

func (z *uploadImpl) makeParams(info os.FileInfo, destPath mo_path.Path, mode string, revision string) *UploadParams {
	upm := &UploadParamMode{
		Tag:    mode,
		Update: "",
	}
	up := &UploadParams{
		Path:           UploadPath(destPath, info).Path(),
		Mode:           upm,
		Mute:           false,
		ClientModified: api_util.RebaseAsString(info.ModTime()),
	}
	switch mode {
	case "update":
		upm.Update = revision
	case "add":
		up.Autorename = true
	}
	return up
}

func (z *uploadImpl) uploadSingle(info os.FileInfo, destPath mo_path.Path, filePath string, mode string, revision string) (entry mo_file.Entry, err error) {
	l := z.ctx.Log().With(zap.String("filePath", filePath), zap.Int64("size", info.Size()))
	l.Debug("Uploading file")

	r, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	res, err := z.ctx.Upload("files/upload", r).
		Param(z.makeParams(info, destPath, mode, revision)).Call()
	if err != nil {
		return nil, err
	}
	entry = &mo_file.Metadata{}
	if err := res.Model(entry); err != nil {
		return nil, err
	}
	return entry, nil
}

func (z *uploadImpl) uploadChunked(info os.FileInfo, destPath mo_path.Path, filePath string, mode string, revision string) (entry mo_file.Entry, err error) {
	l := z.ctx.Log().With(zap.String("filePath", filePath), zap.Int64("size", info.Size()))

	total := info.Size()
	var written int64
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}

	type SessionId struct {
		SessionId string `json:"session_id"`
	}
	type CursorInfo struct {
		SessionId string `json:"session_id"`
		Offset    int64  `json:"offset"`
	}
	type AppendInfo struct {
		Cursor *CursorInfo `json:"cursor"`
	}
	type CommitInfo struct {
		Cursor *CursorInfo   `json:"cursor"`
		Commit *UploadParams `json:"commit"`
	}

	l.Debug("Upload session start")
	r := io.LimitReader(f, z.uo.ChunkSize)
	res, err := z.ctx.Upload("files/upload_session/start", r).Call()
	if err != nil {
		return nil, err
	}
	sid := &SessionId{}
	if j, err := res.Json(); err != nil {
		return nil, err
	} else {
		sid.SessionId = j.Get("session_id").String()
	}
	written += z.uo.ChunkSize
	l = l.With(zap.String("sessionId", sid.SessionId))

	for (total - written) > z.uo.ChunkSize {
		l.Debug("Append chunk", zap.Int64("written", written))
		ai := &AppendInfo{
			Cursor: &CursorInfo{
				SessionId: sid.SessionId,
				Offset:    written,
			},
		}
		r = io.LimitReader(f, z.uo.ChunkSize)
		_, err := z.ctx.Upload("files/upload_session/append_v2", r).Param(ai).Call()
		if err != nil {
			return nil, err
		}
		written += z.uo.ChunkSize
	}

	l.Debug("Finish")
	ci := &CommitInfo{
		Cursor: &CursorInfo{
			SessionId: sid.SessionId,
			Offset:    written,
		},
		Commit: z.makeParams(info, destPath, mode, revision),
	}
	res, err = z.ctx.Upload("files/upload_session/finish", f).Param(ci).Call()
	if err != nil {
		return nil, err
	}
	entry = &mo_file.Metadata{}
	if err := res.Model(entry); err != nil {
		return nil, err
	}
	return entry, nil
}
<|endoftext|>


<commit_before>package exec

import (
	"fmt"
	"gosubli.me/mg"
	"gosubli.me/sink"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	zChunk  = []byte{}
	zChunks = [][]byte{}

	caseCache = struct {
		sync.Mutex
		m map[string]*regexp.Regexp
	}{m: map[string]*regexp.Regexp{}}
)

type StreamRes struct {
	Chunks [][]byte
	Stream string
	End    bool
}

type Res struct {
	Attrs  []Attr
	Chunks [][]byte
	Ok     bool
	Dur    string
	Mem    string
}

type Attr map[string]string

type Switch struct {
	Case        []string
	Negative    bool
	Discard     bool
	Fallthrough bool

	cases []*regexp.Regexp
	attrs []Attr
}

type Exec struct {
	Cid    string
	Stream string

	Wd            string
	Fn            string
	Input         string
	DiscardStdout bool
	DiscardStderr bool
	Env           map[string]string
	Cmd           string
	Args          []string
	Switch        []Switch
	SwitchOk      bool

	filter   func([]byte) []byte
	fini     func()
	switches []*Switch
	brk      *mg.Broker
	sink     *sink.Chan
	buf      struct {
		s     [][]byte
		n     int
		merge bool
	}
}

func (e *Exec) Call() (interface{}, string) {
	e.initSwitches()
	cl, err := e.cl()

	if e.fini != nil {
		defer e.fini()
	}

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

	wg := &sync.WaitGroup{}
	wg.Add(1)
	go e.stream(wg)
	var c *exec.Cmd
	var dur mg.MsDuration

	for _, c = range cl {
		err, dur = procs.Run(e.Cid, c)
		if err != nil {
			break
		}
	}

	e.sink.Close()
	wg.Wait()

	res := Res{
		Chunks: e.chunks(),
		Mem:    maxRss(c.ProcessState),
		Dur:    dur.String(),
	}
	res.Attrs, res.Ok = e.cmdAttrsOk(c)

	return res, mg.Err(err)
}

func (e *Exec) cl() ([]*exec.Cmd, error) {
	switch e.Cmd {
	case ".play":
		return playCmd(e)
	default:
		return []*exec.Cmd{mkCmd(e, e.Input, e.Cmd, e.Args...)}, nil
	}
}

func (e *Exec) stream(wg *sync.WaitGroup) {
	defer wg.Done()

	tck := time.NewTicker(500 * time.Millisecond)
	for {
		select {
		case s, ok := <-e.sink.C:
			if e.filter != nil {
				s = e.filter(s)
			}

			eof := !ok
			s = e.sw(s)
			e.put(s, eof)
			if eof {
				return
			}
		case <-tck.C:
			if len(e.buf.s) != 0 {
				e.flush(false)
			}
		}
	}
}

func (e *Exec) put(s []byte, eof bool) {
	switch {
	case crSuffix(s):
		e.buf.merge = false
		e.buf.s = append(e.buf.s, s)
	case !e.buf.merge || len(e.buf.s) == 0:
		e.buf.merge = true
		e.buf.s = append(e.buf.s, zChunk)
		fallthrough
	default:
		i := len(e.buf.s) - 1
		e.buf.s[i] = append(e.buf.s[i], s...)
	}

	e.buf.n += len(s)
	if eof || e.buf.n >= 32*1024 {
		e.flush(eof)
	}
}

func (e *Exec) chunks() [][]byte {
	if e.buf.s != nil {
		return e.buf.s
	}
	// be nice to the client by not sending them `null` when they expect an array
	return zChunks
}

func (e *Exec) flush(eof bool) {
	if e.Stream == "" {
		return
	}

	e.brk.Send(mg.Response{
		Token: e.Stream,
		Data: StreamRes{
			Chunks: e.chunks(),
			End:    eof,
		},
	})

	for i, s := range e.buf.s {
		e.buf.s[i] = s[:0]
	}
	e.buf.s = e.buf.s[:0]
	e.buf.n = 0
}

func (e *Exec) initSwitches() error {
	for i, _ := range e.Switch {
		sw := &e.Switch[i]
		for _, s := range sw.Case {
			if s != "" {
				rx, err := rx(s)
				if err != nil {
					return err
				}
				sw.cases = append(sw.cases, rx)
			}
		}
		e.switches = append(e.switches, sw)
	}
	return nil
}

func (e *Exec) cmdAttrsOk(c *exec.Cmd) ([]Attr, bool) {
	attrs, swOk := e.collectAttrs()
	if e.SwitchOk && len(e.Switch) > 0 {
		return attrs, swOk
	}
	return attrs, c.ProcessState != nil && c.ProcessState.Success()
}

func (e *Exec) collectAttrs() ([]Attr, bool) {
	attrs := []Attr{}
	ok := true
	for _, sw := range e.switches {
		if len(sw.attrs) > 0 {
			ok = false
			attrs = append(attrs, sw.attrs...)
		}
	}
	return attrs, ok
}

func (e *Exec) sw(s []byte) []byte {
	discard := false

	for _, sw := range e.switches {
		attr, matched := sw.match(s)
		if len(attr) != 0 {
			sw.attrs = append(sw.attrs, attr)
		}

		if matched != sw.Negative {
			if sw.Discard {
				discard = true
			}

			if !sw.Fallthrough {
				break
			}
		}
	}

	if discard {
		return nil
	}
	return s
}

func (s *Switch) match(p []byte) (Attr, bool) {
	if len(s.cases) == 0 {
		return nil, true
	}

	for _, rx := range s.cases {
		mt := rx.FindSubmatch(p)
		if len(mt) > 0 {
			// p (and by extension mt) is owned by the caller so make sure it gets copied
			names := rx.SubexpNames()
			attr := Attr{}
			if len(mt) == len(names) {
				for i, k := range names {
					if k != "" {
						attr[k] = string(mt[i])
					}
				}
			}
			return attr, true
		}
	}

	return nil, false
}

func rx(s string) (*regexp.Regexp, error) {
	caseCache.Lock()
	defer caseCache.Unlock()

	if rx, ok := caseCache.m[s]; ok {
		return rx, nil
	}

	rx, err := regexp.Compile(s)
	if err != nil {
		return nil, fmt.Errorf("cannot compile regexp `%v`: %v", s, err.Error())
	}

	caseCache.m[s] = rx
	return rx, nil
}

func crSuffix(s []byte) bool {
	i := len(s) - 1
	return i >= 0 && s[i] == '\r'
}

func mkCmd(e *Exec, input string, cmd string, args ...string) *exec.Cmd {
	c := exec.Command(cmd, args...)
	c.Dir = e.Wd
	c.Env = mg.Env(e.Env)
	if !e.DiscardStdout {
		c.Stdout = e.sink
	}
	if !e.DiscardStderr {
		c.Stderr = e.sink
	}
	if input != "" {
		c.Stdin = strings.NewReader(input)
	}
	return c
}

func init() {
	mg.Register("exec", func(b *mg.Broker) mg.Caller {
		return &Exec{
			brk:  b,
			sink: sink.NewBufferedChan(100),
		}
	})
}
<commit_msg>* allow multiple match per chunk<commit_after>package exec

import (
	"fmt"
	"gosubli.me/mg"
	"gosubli.me/sink"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	zChunk  = []byte{}
	zChunks = [][]byte{}

	caseCache = struct {
		sync.Mutex
		m map[string]*regexp.Regexp
	}{m: map[string]*regexp.Regexp{}}
)

type StreamRes struct {
	Chunks [][]byte
	Stream string
	End    bool
}

type Res struct {
	Attrs  []Attr
	Chunks [][]byte
	Ok     bool
	Dur    string
	Mem    string
}

type Attr map[string]string

type Switch struct {
	Case        []string
	Negative    bool
	Discard     bool
	Fallthrough bool

	cases []*regexp.Regexp
	attrs []Attr
}

type Exec struct {
	Cid    string
	Stream string

	Wd            string
	Fn            string
	Input         string
	DiscardStdout bool
	DiscardStderr bool
	Env           map[string]string
	Cmd           string
	Args          []string
	Switch        []Switch
	SwitchOk      bool

	filter   func([]byte) []byte
	fini     func()
	switches []*Switch
	brk      *mg.Broker
	sink     *sink.Chan
	buf      struct {
		s     [][]byte
		n     int
		merge bool
	}
}

func (e *Exec) Call() (interface{}, string) {
	e.initSwitches()
	cl, err := e.cl()

	if e.fini != nil {
		defer e.fini()
	}

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

	wg := &sync.WaitGroup{}
	wg.Add(1)
	go e.stream(wg)
	var c *exec.Cmd
	var dur mg.MsDuration

	for _, c = range cl {
		err, dur = procs.Run(e.Cid, c)
		if err != nil {
			break
		}
	}

	e.sink.Close()
	wg.Wait()

	res := Res{
		Chunks: e.chunks(),
		Mem:    maxRss(c.ProcessState),
		Dur:    dur.String(),
	}
	res.Attrs, res.Ok = e.cmdAttrsOk(c)

	return res, mg.Err(err)
}

func (e *Exec) cl() ([]*exec.Cmd, error) {
	switch e.Cmd {
	case ".play":
		return playCmd(e)
	default:
		return []*exec.Cmd{mkCmd(e, e.Input, e.Cmd, e.Args...)}, nil
	}
}

func (e *Exec) stream(wg *sync.WaitGroup) {
	defer wg.Done()

	tck := time.NewTicker(500 * time.Millisecond)
	for {
		select {
		case s, ok := <-e.sink.C:
			if e.filter != nil {
				s = e.filter(s)
			}

			eof := !ok
			s = e.sw(s)
			e.put(s, eof)
			if eof {
				return
			}
		case <-tck.C:
			if len(e.buf.s) != 0 {
				e.flush(false)
			}
		}
	}
}

func (e *Exec) put(s []byte, eof bool) {
	switch {
	case crSuffix(s):
		e.buf.merge = false
		e.buf.s = append(e.buf.s, s)
	case !e.buf.merge || len(e.buf.s) == 0:
		e.buf.merge = true
		e.buf.s = append(e.buf.s, zChunk)
		fallthrough
	default:
		i := len(e.buf.s) - 1
		e.buf.s[i] = append(e.buf.s[i], s...)
	}

	e.buf.n += len(s)
	if eof || e.buf.n >= 32*1024 {
		e.flush(eof)
	}
}

func (e *Exec) chunks() [][]byte {
	if e.buf.s != nil {
		return e.buf.s
	}
	// be nice to the client by not sending them `null` when they expect an array
	return zChunks
}

func (e *Exec) flush(eof bool) {
	if e.Stream == "" {
		return
	}

	e.brk.Send(mg.Response{
		Token: e.Stream,
		Data: StreamRes{
			Chunks: e.chunks(),
			End:    eof,
		},
	})

	for i, s := range e.buf.s {
		e.buf.s[i] = s[:0]
	}
	e.buf.s = e.buf.s[:0]
	e.buf.n = 0
}

func (e *Exec) initSwitches() error {
	for i, _ := range e.Switch {
		sw := &e.Switch[i]
		for _, s := range sw.Case {
			if s != "" {
				rx, err := rx(s)
				if err != nil {
					return err
				}
				sw.cases = append(sw.cases, rx)
			}
		}
		e.switches = append(e.switches, sw)
	}
	return nil
}

func (e *Exec) cmdAttrsOk(c *exec.Cmd) ([]Attr, bool) {
	attrs, swOk := e.collectAttrs()
	if e.SwitchOk && len(e.Switch) > 0 {
		return attrs, swOk
	}
	return attrs, c.ProcessState != nil && c.ProcessState.Success()
}

func (e *Exec) collectAttrs() ([]Attr, bool) {
	attrs := []Attr{}
	ok := true
	for _, sw := range e.switches {
		if len(sw.attrs) > 0 {
			ok = false
			attrs = append(attrs, sw.attrs...)
		}
	}
	return attrs, ok
}

func (e *Exec) sw(s []byte) []byte {
	discard := false

	for _, sw := range e.switches {
		al, matched := sw.match(s)
		sw.attrs = append(sw.attrs, al...)

		if matched != sw.Negative {
			if sw.Discard {
				discard = true
			}

			if !sw.Fallthrough {
				break
			}
		}
	}

	if discard {
		return nil
	}
	return s
}

func (s *Switch) match(p []byte) ([]Attr, bool) {
	if len(s.cases) == 0 {
		return nil, true
	}

	for _, rx := range s.cases {
		ml := rx.FindAllSubmatch(p, -1)
		if len(ml) > 0 {
			// p (and by extension ml, mt) is owned by the caller so make sure it gets copied
			names := rx.SubexpNames()
			al := []Attr{}
			for _, mt := range ml {
				attr := Attr{}
				if len(mt) == len(names) {
					for i, k := range names {
						if k != "" {
							attr[k] = string(mt[i])
						}
					}
				}
				if len(mt) != 0 {
					al = append(al, attr)
				}
			}
			return al, true
		}
	}

	return nil, false
}

func rx(s string) (*regexp.Regexp, error) {
	caseCache.Lock()
	defer caseCache.Unlock()

	if rx, ok := caseCache.m[s]; ok {
		return rx, nil
	}

	rx, err := regexp.Compile(s)
	if err != nil {
		return nil, fmt.Errorf("cannot compile regexp `%v`: %v", s, err.Error())
	}

	caseCache.m[s] = rx
	return rx, nil
}

func crSuffix(s []byte) bool {
	i := len(s) - 1
	return i >= 0 && s[i] == '\r'
}

func mkCmd(e *Exec, input string, cmd string, args ...string) *exec.Cmd {
	c := exec.Command(cmd, args...)
	c.Dir = e.Wd
	c.Env = mg.Env(e.Env)
	if !e.DiscardStdout {
		c.Stdout = e.sink
	}
	if !e.DiscardStderr {
		c.Stderr = e.sink
	}
	if input != "" {
		c.Stdin = strings.NewReader(input)
	}
	return c
}

func init() {
	mg.Register("exec", func(b *mg.Broker) mg.Caller {
		return &Exec{
			brk:  b,
			sink: sink.NewBufferedChan(100),
		}
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors

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

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

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

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/knative/serving/cmd/util"
	"github.com/knative/serving/pkg/autoscaler"
	"k8s.io/apimachinery/pkg/util/wait"

	"github.com/knative/pkg/logging/logkey"

	"github.com/knative/pkg/configmap"
	"github.com/knative/pkg/signals"
	"github.com/knative/serving/pkg/activator"
	activatorhandler "github.com/knative/serving/pkg/activator/handler"
	activatorutil "github.com/knative/serving/pkg/activator/util"
	clientset "github.com/knative/serving/pkg/client/clientset/versioned"
	"github.com/knative/serving/pkg/http/h2c"
	"github.com/knative/serving/pkg/logging"
	"github.com/knative/serving/pkg/metrics"
	"github.com/knative/serving/pkg/system"
	"go.uber.org/zap"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/knative/serving/pkg/websocket"
)

const (
	maxUploadBytes = 32e6 // 32MB - same as app engine
	component      = "activator"

	maxRetries             = 18 // the sum of all retries would add up to 1 minute
	minRetryInterval       = 100 * time.Millisecond
	exponentialBackoffBase = 1.3

	// Add a little buffer space between request handling and stat
	// reporting so that latency in the stat pipeline doesn't
	// interfere with request handling.
	statReportingQueueLength = 10

	// Add enough buffer to not block request serving on stats collection
	requestCountingQueueLength = 100
)

var (
	logger *zap.SugaredLogger

	statSink *websocket.ManagedConnection
	statChan = make(chan *autoscaler.StatMessage, statReportingQueueLength)
	reqChan  = make(chan activatorhandler.ReqEvent, requestCountingQueueLength)
)

func statReporter() {
	for {
		sm := <-statChan
		if statSink == nil {
			logger.Error("Stat sink not connected.")
			continue
		}
		err := statSink.Send(sm)
		if err != nil {
			logger.Error("Error while sending stat", zap.Error(err))
		}
	}
}

func main() {
	flag.Parse()
	cm, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatalf("Error loading logging configuration: %v", err)
	}
	config, err := logging.NewConfigFromMap(cm)
	if err != nil {
		log.Fatalf("Error parsing logging configuration: %v", err)
	}
	createdLogger, atomicLevel := logging.NewLoggerFromConfig(config, component)
	logger = createdLogger.With(zap.String(logkey.ControllerType, "activator"))
	defer logger.Sync()

	logger.Info("Starting the knative activator")

	clusterConfig, err := rest.InClusterConfig()
	if err != nil {
		logger.Fatal("Error getting in cluster configuration", zap.Error(err))
	}
	kubeClient, err := kubernetes.NewForConfig(clusterConfig)
	if err != nil {
		logger.Fatal("Error building new kubernetes client", zap.Error(err))
	}
	servingClient, err := clientset.NewForConfig(clusterConfig)
	if err != nil {
		logger.Fatal("Error building serving clientset", zap.Error(err))
	}

	reporter, err := activator.NewStatsReporter()
	if err != nil {
		logger.Fatal("Failed to create stats reporter", zap.Error(err))
	}

	a := activator.NewRevisionActivator(kubeClient, servingClient, logger, reporter)
	a = activator.NewDedupingActivator(a)

	// Retry on 503's for up to 60 seconds. The reason is there is
	// a small delay for k8s to include the ready IP in service.
	// https://github.com/knative/serving/issues/660#issuecomment-384062553
	shouldRetry := activatorutil.RetryStatus(http.StatusServiceUnavailable)
	backoffSettings := wait.Backoff{
		Duration: minRetryInterval,
		Factor:   exponentialBackoffBase,
		Steps:    maxRetries,
	}
	rt := activatorutil.NewRetryRoundTripper(activatorutil.AutoTransport, logger, backoffSettings, shouldRetry)

	// Open a websocket connection to the autoscaler
	autoscalerEndpoint := fmt.Sprintf("ws://%s.%s.svc.cluster.local:%s", "autoscaler", system.Namespace, "8080")
	logger.Infof("Connecting to autoscaler at %s", autoscalerEndpoint)
	statSink = websocket.NewDurableSendingConnection(autoscalerEndpoint)
	go statReporter()

	podName := util.GetRequiredEnvOrFatal("POD_NAME", logger)
	activatorhandler.NewConcurrencyReporter(podName, activatorhandler.Channels{
		ReqChan:    reqChan,
		StatChan:   statChan,
		ReportChan: time.NewTicker(time.Second).C,
	})

	ah := &activatorhandler.FilteringHandler{
		NextHandler: activatorhandler.NewRequestEventHandler(reqChan,
			&activatorhandler.EnforceMaxContentLengthHandler{
				MaxContentLengthBytes: maxUploadBytes,
				NextHandler: &activatorhandler.ActivationHandler{
					Activator: a,
					Transport: rt,
					Logger:    logger,
					Reporter:  reporter,
				},
			},
		),
	}

	// set up signals so we handle the first shutdown signal gracefully
	stopCh := signals.SetupSignalHandler()
	go func() {
		<-stopCh
		a.Shutdown()
	}()

	// Watch the logging config map and dynamically update logging levels.
	configMapWatcher := configmap.NewInformedWatcher(kubeClient, system.Namespace)
	configMapWatcher.Watch(logging.ConfigName, logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))
	// Watch the observability config map and dynamically update metrics exporter.
	configMapWatcher.Watch(metrics.ObservabilityConfigName, metrics.UpdateExporterFromConfigMap(component, logger))
	if err = configMapWatcher.Start(stopCh); err != nil {
		logger.Fatalf("failed to start configuration manager: %v", err)
	}

	h2c.ListenAndServe(":8080", ah)
}
<commit_msg>Fix activator blocking on SIGTERM (#2571)<commit_after>/*
Copyright 2018 The Knative Authors

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

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

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

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/knative/serving/cmd/util"
	"github.com/knative/serving/pkg/autoscaler"
	"k8s.io/apimachinery/pkg/util/wait"

	"github.com/knative/pkg/logging/logkey"

	"github.com/knative/pkg/configmap"
	"github.com/knative/pkg/signals"
	"github.com/knative/serving/pkg/activator"
	activatorhandler "github.com/knative/serving/pkg/activator/handler"
	activatorutil "github.com/knative/serving/pkg/activator/util"
	clientset "github.com/knative/serving/pkg/client/clientset/versioned"
	"github.com/knative/serving/pkg/http/h2c"
	"github.com/knative/serving/pkg/logging"
	"github.com/knative/serving/pkg/metrics"
	"github.com/knative/serving/pkg/system"
	"go.uber.org/zap"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/knative/serving/pkg/websocket"
)

const (
	maxUploadBytes = 32e6 // 32MB - same as app engine
	component      = "activator"

	maxRetries             = 18 // the sum of all retries would add up to 1 minute
	minRetryInterval       = 100 * time.Millisecond
	exponentialBackoffBase = 1.3

	// Add a little buffer space between request handling and stat
	// reporting so that latency in the stat pipeline doesn't
	// interfere with request handling.
	statReportingQueueLength = 10

	// Add enough buffer to not block request serving on stats collection
	requestCountingQueueLength = 100
)

var (
	logger *zap.SugaredLogger

	statSink *websocket.ManagedConnection
	statChan = make(chan *autoscaler.StatMessage, statReportingQueueLength)
	reqChan  = make(chan activatorhandler.ReqEvent, requestCountingQueueLength)
)

func statReporter(stopCh <-chan struct{}) {
	for {
		select {
		case sm := <-statChan:
			if statSink == nil {
				logger.Error("Stat sink not connected.")
				continue
			}
			err := statSink.Send(sm)
			if err != nil {
				logger.Error("Error while sending stat", zap.Error(err))
			}
		case <-stopCh:
			return
		}
	}
}

func main() {
	flag.Parse()
	cm, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatalf("Error loading logging configuration: %v", err)
	}
	config, err := logging.NewConfigFromMap(cm)
	if err != nil {
		log.Fatalf("Error parsing logging configuration: %v", err)
	}
	createdLogger, atomicLevel := logging.NewLoggerFromConfig(config, component)
	logger = createdLogger.With(zap.String(logkey.ControllerType, "activator"))
	defer logger.Sync()

	logger.Info("Starting the knative activator")

	clusterConfig, err := rest.InClusterConfig()
	if err != nil {
		logger.Fatal("Error getting in cluster configuration", zap.Error(err))
	}
	kubeClient, err := kubernetes.NewForConfig(clusterConfig)
	if err != nil {
		logger.Fatal("Error building new kubernetes client", zap.Error(err))
	}
	servingClient, err := clientset.NewForConfig(clusterConfig)
	if err != nil {
		logger.Fatal("Error building serving clientset", zap.Error(err))
	}

	reporter, err := activator.NewStatsReporter()
	if err != nil {
		logger.Fatal("Failed to create stats reporter", zap.Error(err))
	}

	a := activator.NewRevisionActivator(kubeClient, servingClient, logger, reporter)
	a = activator.NewDedupingActivator(a)

	// Retry on 503's for up to 60 seconds. The reason is there is
	// a small delay for k8s to include the ready IP in service.
	// https://github.com/knative/serving/issues/660#issuecomment-384062553
	shouldRetry := activatorutil.RetryStatus(http.StatusServiceUnavailable)
	backoffSettings := wait.Backoff{
		Duration: minRetryInterval,
		Factor:   exponentialBackoffBase,
		Steps:    maxRetries,
	}
	rt := activatorutil.NewRetryRoundTripper(activatorutil.AutoTransport, logger, backoffSettings, shouldRetry)

	// set up signals so we handle the first shutdown signal gracefully
	stopCh := signals.SetupSignalHandler()

	// Open a websocket connection to the autoscaler
	autoscalerEndpoint := fmt.Sprintf("ws://%s.%s.svc.cluster.local:%s", "autoscaler", system.Namespace, "8080")
	logger.Infof("Connecting to autoscaler at %s", autoscalerEndpoint)
	statSink = websocket.NewDurableSendingConnection(autoscalerEndpoint)
	go statReporter(stopCh)

	podName := util.GetRequiredEnvOrFatal("POD_NAME", logger)
	activatorhandler.NewConcurrencyReporter(podName, activatorhandler.Channels{
		ReqChan:    reqChan,
		StatChan:   statChan,
		ReportChan: time.NewTicker(time.Second).C,
	})

	ah := &activatorhandler.FilteringHandler{
		NextHandler: activatorhandler.NewRequestEventHandler(reqChan,
			&activatorhandler.EnforceMaxContentLengthHandler{
				MaxContentLengthBytes: maxUploadBytes,
				NextHandler: &activatorhandler.ActivationHandler{
					Activator: a,
					Transport: rt,
					Logger:    logger,
					Reporter:  reporter,
				},
			},
		),
	}

	// Watch the logging config map and dynamically update logging levels.
	configMapWatcher := configmap.NewInformedWatcher(kubeClient, system.Namespace)
	configMapWatcher.Watch(logging.ConfigName, logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))
	// Watch the observability config map and dynamically update metrics exporter.
	configMapWatcher.Watch(metrics.ObservabilityConfigName, metrics.UpdateExporterFromConfigMap(component, logger))
	if err = configMapWatcher.Start(stopCh); err != nil {
		logger.Fatalf("failed to start configuration manager: %v", err)
	}

	srv := h2c.NewServer(":8080", ah)
	go func() {
		if err := srv.ListenAndServe(); err != nil {
			logger.Errorf("Error running HTTP server: %v", err)
		}
	}()

	<-stopCh
	a.Shutdown()
	srv.Shutdown(context.TODO())
}
<|endoftext|>


<commit_before>/*
Copyright 2013 Google 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.
*/

// Hacks for running camput as a child process on Android.

package main

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sync"
	"time"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/schema"
)

var detectOnce sync.Once
var onAndroidCache bool

func dirExists(f string) bool {
	fi, err := os.Stat(f)
	return err == nil && fi.IsDir()
}

func initOnAndroid() {
	// Good enough heuristic. Suggestions welcome.
	onAndroidCache = dirExists("/data/data") && dirExists("/system/etc")
}

func onAndroid() bool {
	detectOnce.Do(initOnAndroid)
	return onAndroidCache
}

var pingRx = regexp.MustCompile(`\((.+?)\)`)

func androidLookupHost(host string) string {
	// Android has no "dig" or "host" tool, so use "ping -c 1". Ghetto.
	// $ ping -c 1 google.com
	// PING google.com (74.125.224.64) 56(84) bytes of data.
	c := make(chan string, 1)
	go func() {
		cmd := exec.Command("/system/bin/ping", "-c", "1", host)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			panic(err)
		}
		defer stdout.Close()
		if err := cmd.Start(); err != nil {
			log.Printf("Error resolving %q with ping: %v", host, err)
			c <- host
			return
		}
		defer func() {
			if p := cmd.Process; p != nil {
				p.Kill()
			}
		}()
		br := bufio.NewReader(stdout)
		line, err := br.ReadString('\n')
		if err != nil {
			c <- host
			return
		}
		if m := pingRx.FindStringSubmatch(line); m != nil {
			c <- m[1]
			return
		}
		c <- host
	}()
	return <-c
}

func dialFunc() func(network, addr string) (net.Conn, error) {
	if !onAndroid() {
		return nil // use default
	}
	return func(network, addr string) (net.Conn, error) {
		// Temporary laziness hack, avoiding doing a
		// cross-compiled Android cgo build.
		// Without cgo, package net uses
		// /etc/resolv.conf (not available on
		// Android).  We really want a cgo binary to
		// use Android's DNS cache, but it's kinda
		// hard/impossible to cross-compile for now.
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			return nil, fmt.Errorf("couldn't split %q", addr)
		}
		return net.Dial(network, net.JoinHostPort(androidLookupHost(host), port))
	}
}

func tlsClientConfig() *tls.Config {
	const certDir = "/system/etc/security/cacerts"
	if fi, err := os.Stat(certDir); err != nil || !fi.IsDir() {
		return nil
	}
	pool := x509.NewCertPool()
	cfg := &tls.Config{RootCAs: pool}

	f, err := os.Open(certDir)
	if err != nil {
		return nil
	}
	defer f.Close()
	names, _ := f.Readdirnames(-1)
	for _, name := range names {
		if pem, err := ioutil.ReadFile(filepath.Join(certDir, name)); err == nil {
			pool.AppendCertsFromPEM(pem)
		}
	}
	return cfg
}

var androidOutput = os.Getenv("CAMPUT_ANDROID_OUTPUT") != ""

var androidOutMu sync.Mutex

func androidf(format string, args ...interface{}) {
	androidOutMu.Lock()
	defer androidOutMu.Unlock()
	fmt.Printf(format, args...)
}

func noteFileUploaded(fullPath string) {
	if !androidOutput {
		return
	}
	androidf("FILE_UPLOADED %s\n", fullPath)
}

type allStats struct {
	total, skipped, uploaded stats
}

var lastStatBroadcast allStats

func printAndroidCamputStatus(t *TreeUpload) {
	bcast := allStats{t.total, t.skipped, t.uploaded}
	if bcast == lastStatBroadcast {
		return
	}
	lastStatBroadcast = bcast

	androidf("STATS nfile=%d nbyte=%d skfile=%d skbyte=%d upfile=%d upbyte=%d\n",
		t.total.files, t.total.bytes,
		t.skipped.files, t.skipped.bytes,
		t.uploaded.files, t.uploaded.bytes)
}

// androidStatusRecevier is a blobserver.StatReceiver wrapper that
// reports the full filename path and size of uploaded blobs.
// The android app wrapping camput watches stdout for this, for progress bars.
type androidStatusRecevier struct {
	sr blobserver.StatReceiver
	path string
}

func (asr androidStatusRecevier) noteChunkOnServer(sb blobref.SizedBlobRef) {
	androidf("CHUNK_UPLOADED %d %s %s\n", sb.Size, sb.BlobRef, asr.path)
}

func (asr androidStatusRecevier) ReceiveBlob(blob *blobref.BlobRef, source io.Reader) (blobref.SizedBlobRef, error) {
	// Sniff the first 1KB of it and don't print the stats if it looks like it was just a schema
	// blob.  We won't update the progress bar for that yet.
	var buf [1024]byte
	contents := buf[:0]
	sb, err := asr.sr.ReceiveBlob(blob, io.TeeReader(source, writeUntilSliceFull{&contents}))
	if err == nil && !schema.LikelySchemaBlob(contents) {
		asr.noteChunkOnServer(sb)
	}
	return sb, err
}

func (asr androidStatusRecevier) StatBlobs(dest chan<- blobref.SizedBlobRef, blobs []*blobref.BlobRef, wait time.Duration) error {
	midc := make(chan blobref.SizedBlobRef)
	errc := make(chan error, 1)
	go func() {
		err := asr.sr.StatBlobs(midc, blobs, wait)
		errc <- err
		close(midc)
	}()
	for sb := range midc {
		asr.noteChunkOnServer(sb)
		dest <- sb
	}
	return <-errc
}

type writeUntilSliceFull struct {
	s *[]byte
}

func (w writeUntilSliceFull) Write(p []byte) (n int, err error) {
	s := *w.s
	l := len(s)
	growBy := cap(s) - l
	if growBy > len(p) {
		growBy = len(p)
	}
	s = s[0:l+growBy]
	copy(s[l:], p)
	*w.s = s
	return len(p), nil
}

<commit_msg>android: TODO notes for camput<commit_after>/*
Copyright 2013 Google 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.
*/

// Hacks for running camput as a child process on Android.

package main

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sync"
	"time"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/blobserver"
	"camlistore.org/pkg/schema"
)

var detectOnce sync.Once
var onAndroidCache bool

func dirExists(f string) bool {
	fi, err := os.Stat(f)
	return err == nil && fi.IsDir()
}

func initOnAndroid() {
	// Good enough heuristic. Suggestions welcome.
	onAndroidCache = dirExists("/data/data") && dirExists("/system/etc")
}

func onAndroid() bool {
	detectOnce.Do(initOnAndroid)
	return onAndroidCache
}

var pingRx = regexp.MustCompile(`\((.+?)\)`)

func androidLookupHost(host string) string {
	// Android has no "dig" or "host" tool, so use "ping -c 1". Ghetto.
	// $ ping -c 1 google.com
	// PING google.com (74.125.224.64) 56(84) bytes of data.
	c := make(chan string, 1)
	go func() {
		cmd := exec.Command("/system/bin/ping", "-c", "1", host)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			panic(err)
		}
		defer stdout.Close()
		if err := cmd.Start(); err != nil {
			log.Printf("Error resolving %q with ping: %v", host, err)
			c <- host
			return
		}
		defer func() {
			if p := cmd.Process; p != nil {
				p.Kill()
			}
		}()
		br := bufio.NewReader(stdout)
		line, err := br.ReadString('\n')
		if err != nil {
			c <- host
			return
		}
		if m := pingRx.FindStringSubmatch(line); m != nil {
			c <- m[1]
			return
		}
		c <- host
	}()
	return <-c
}

func dialFunc() func(network, addr string) (net.Conn, error) {
	if !onAndroid() {
		return nil // use default
	}
	return func(network, addr string) (net.Conn, error) {
		// Temporary laziness hack, avoiding doing a
		// cross-compiled Android cgo build.
		// Without cgo, package net uses
		// /etc/resolv.conf (not available on
		// Android).  We really want a cgo binary to
		// use Android's DNS cache, but it's kinda
		// hard/impossible to cross-compile for now.
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			return nil, fmt.Errorf("couldn't split %q", addr)
		}
		return net.Dial(network, net.JoinHostPort(androidLookupHost(host), port))
	}
}

func tlsClientConfig() *tls.Config {
	const certDir = "/system/etc/security/cacerts"
	if fi, err := os.Stat(certDir); err != nil || !fi.IsDir() {
		return nil
	}
	pool := x509.NewCertPool()
	cfg := &tls.Config{RootCAs: pool}

	f, err := os.Open(certDir)
	if err != nil {
		return nil
	}
	defer f.Close()
	names, _ := f.Readdirnames(-1)
	for _, name := range names {
		if pem, err := ioutil.ReadFile(filepath.Join(certDir, name)); err == nil {
			pool.AppendCertsFromPEM(pem)
		}
	}
	return cfg
}

var androidOutput = os.Getenv("CAMPUT_ANDROID_OUTPUT") != ""

var androidOutMu sync.Mutex

func androidf(format string, args ...interface{}) {
	androidOutMu.Lock()
	defer androidOutMu.Unlock()
	fmt.Printf(format, args...)
}

func noteFileUploaded(fullPath string) {
	if !androidOutput {
		return
	}
	androidf("FILE_UPLOADED %s\n", fullPath)
}

type allStats struct {
	total, skipped, uploaded stats
}

var lastStatBroadcast allStats

func printAndroidCamputStatus(t *TreeUpload) {
	bcast := allStats{t.total, t.skipped, t.uploaded}
	if bcast == lastStatBroadcast {
		return
	}
	lastStatBroadcast = bcast

	androidf("STATS nfile=%d nbyte=%d skfile=%d skbyte=%d upfile=%d upbyte=%d\n",
		t.total.files, t.total.bytes,
		t.skipped.files, t.skipped.bytes,
		t.uploaded.files, t.uploaded.bytes)
}

// androidStatusRecevier is a blobserver.StatReceiver wrapper that
// reports the full filename path and size of uploaded blobs.
// The android app wrapping camput watches stdout for this, for progress bars.
type androidStatusRecevier struct {
	sr blobserver.StatReceiver
	path string
}

func (asr androidStatusRecevier) noteChunkOnServer(sb blobref.SizedBlobRef) {
	androidf("CHUNK_UPLOADED %d %s %s\n", sb.Size, sb.BlobRef, asr.path)
}

func (asr androidStatusRecevier) ReceiveBlob(blob *blobref.BlobRef, source io.Reader) (blobref.SizedBlobRef, error) {
	// Sniff the first 1KB of it and don't print the stats if it looks like it was just a schema
	// blob.  We won't update the progress bar for that yet.
	var buf [1024]byte
	contents := buf[:0]
	sb, err := asr.sr.ReceiveBlob(blob, io.TeeReader(source, writeUntilSliceFull{&contents}))
	if err == nil && !schema.LikelySchemaBlob(contents) {
		asr.noteChunkOnServer(sb)
	}
	return sb, err
}

func (asr androidStatusRecevier) StatBlobs(dest chan<- blobref.SizedBlobRef, blobs []*blobref.BlobRef, wait time.Duration) error {
	midc := make(chan blobref.SizedBlobRef)
	errc := make(chan error, 1)
	go func() {
		err := asr.sr.StatBlobs(midc, blobs, wait)
		errc <- err
		close(midc)
	}()
	for sb := range midc {
		asr.noteChunkOnServer(sb)
		dest <- sb
	}
	return <-errc
}

type writeUntilSliceFull struct {
	s *[]byte
}

func (w writeUntilSliceFull) Write(p []byte) (n int, err error) {
	s := *w.s
	l := len(s)
	growBy := cap(s) - l
	if growBy > len(p) {
		growBy = len(p)
	}
	s = s[0:l+growBy]
	copy(s[l:], p)
	*w.s = s
	return len(p), nil
}

// TODO: show in the Android UI stats about:
// * TCP connects/disconnects.
// * TCP bytes written.
// * TCP bytes received.
// * HTTP requests started.
// * HTTP requests sent.
// * HTTP headers received.
// * HTTP bodies received.
// * Byte up/down counts and summary of each HTTP request?
// For example:
//
// In-flight:
//   (17) blob-upload (123123423 up)
//
// Recently-completed: (last 10 or something)
//   (2) blob-upload (1 blobs) (483234 up / 1239 down, 200 OK)
//   (3) blob-stat (2 blobs) (483 up / 1239 down, 204 No Content)
//   (1) blob-stat (1 blobs) (483 up / 1239 down)
<|endoftext|>


<commit_before>// Copyright (c) 2019 The Jaeger Authors.
// Copyright (c) 2017 Uber Technologies, 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 main

import (
	"fmt"
	"io"
	"log"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/uber/jaeger-lib/metrics"
	jexpvar "github.com/uber/jaeger-lib/metrics/expvar"
	"github.com/uber/jaeger-lib/metrics/fork"
	_ "go.uber.org/automaxprocs"
	"go.uber.org/zap"

	"github.com/jaegertracing/jaeger/cmd/collector/app"
	"github.com/jaegertracing/jaeger/cmd/docs"
	"github.com/jaegertracing/jaeger/cmd/env"
	"github.com/jaegertracing/jaeger/cmd/flags"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/version"
	ss "github.com/jaegertracing/jaeger/plugin/sampling/strategystore"
	"github.com/jaegertracing/jaeger/plugin/storage"
	"github.com/jaegertracing/jaeger/ports"
)

const serviceName = "jaeger-collector"

func main() {
	svc := flags.NewService(ports.CollectorAdminHTTP)

	storageFactory, err := storage.NewFactory(storage.FactoryConfigFromEnvAndCLI(os.Args, os.Stderr))
	if err != nil {
		log.Fatalf("Cannot initialize storage factory: %v", err)
	}
	strategyStoreFactory, err := ss.NewFactory(ss.FactoryConfigFromEnv())
	if err != nil {
		log.Fatalf("Cannot initialize sampling strategy store factory: %v", err)
	}

	v := viper.New()
	command := &cobra.Command{
		Use:   "jaeger-collector",
		Short: "Jaeger collector receives and processes traces from Jaeger agents and clients",
		Long:  `Jaeger collector receives traces from Jaeger agents and runs them through a processing pipeline.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := svc.Start(v); err != nil {
				return err
			}
			logger := svc.Logger // shortcut
			baseFactory := svc.MetricsFactory.Namespace(metrics.NSOptions{Name: "jaeger"})
			metricsFactory := fork.New("internal",
				jexpvar.NewFactory(10), // backend for internal opts
				baseFactory.Namespace(metrics.NSOptions{Name: "collector"}))

			storageFactory.InitFromViper(v)
			if err := storageFactory.Initialize(baseFactory, logger); err != nil {
				logger.Fatal("Failed to init storage factory", zap.Error(err))
			}
			spanWriter, err := storageFactory.CreateSpanWriter()
			if err != nil {
				logger.Fatal("Failed to create span writer", zap.Error(err))
			}

			strategyStoreFactory.InitFromViper(v)
			if err := strategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
				logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
			}
			strategyStore, err := strategyStoreFactory.CreateStrategyStore()
			if err != nil {
				logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
			}

			c := app.New(&app.CollectorParams{
				ServiceName:    serviceName,
				Logger:         logger,
				MetricsFactory: metricsFactory,
				SpanWriter:     spanWriter,
				StrategyStore:  strategyStore,
				HealthCheck:    svc.HC(),
			})
			collectorOpts := new(app.CollectorOptions).InitFromViper(v)
			c.Start(collectorOpts)

			svc.RunAndThen(func() {
				if err := c.Close(); err != nil {
					logger.Error("failed to cleanly close the collector", zap.Error(err))
				}
				if closer, ok := spanWriter.(io.Closer); ok {
					err := closer.Close()
					if err != nil {
						logger.Error("failed to close span writer", zap.Error(err))
					}
				}
				if err := storageFactory.Close(); err != nil {
					logger.Error("Failed to close storage factory", zap.Error(err))
				}

			})
			return nil
		},
	}

	command.AddCommand(version.Command())
	command.AddCommand(env.Command())
	command.AddCommand(docs.Command(v))

	config.AddFlags(
		v,
		command,
		svc.AddFlags,
		app.AddFlags,
		storageFactory.AddFlags,
		strategyStoreFactory.AddFlags,
	)

	if err := command.Execute(); err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
}
<commit_msg>Handle collector Start error (#2647)<commit_after>// Copyright (c) 2019 The Jaeger Authors.
// Copyright (c) 2017 Uber Technologies, 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 main

import (
	"fmt"
	"io"
	"log"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/uber/jaeger-lib/metrics"
	jexpvar "github.com/uber/jaeger-lib/metrics/expvar"
	"github.com/uber/jaeger-lib/metrics/fork"
	_ "go.uber.org/automaxprocs"
	"go.uber.org/zap"

	"github.com/jaegertracing/jaeger/cmd/collector/app"
	"github.com/jaegertracing/jaeger/cmd/docs"
	"github.com/jaegertracing/jaeger/cmd/env"
	"github.com/jaegertracing/jaeger/cmd/flags"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/version"
	ss "github.com/jaegertracing/jaeger/plugin/sampling/strategystore"
	"github.com/jaegertracing/jaeger/plugin/storage"
	"github.com/jaegertracing/jaeger/ports"
)

const serviceName = "jaeger-collector"

func main() {
	svc := flags.NewService(ports.CollectorAdminHTTP)

	storageFactory, err := storage.NewFactory(storage.FactoryConfigFromEnvAndCLI(os.Args, os.Stderr))
	if err != nil {
		log.Fatalf("Cannot initialize storage factory: %v", err)
	}
	strategyStoreFactory, err := ss.NewFactory(ss.FactoryConfigFromEnv())
	if err != nil {
		log.Fatalf("Cannot initialize sampling strategy store factory: %v", err)
	}

	v := viper.New()
	command := &cobra.Command{
		Use:   "jaeger-collector",
		Short: "Jaeger collector receives and processes traces from Jaeger agents and clients",
		Long:  `Jaeger collector receives traces from Jaeger agents and runs them through a processing pipeline.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := svc.Start(v); err != nil {
				return err
			}
			logger := svc.Logger // shortcut
			baseFactory := svc.MetricsFactory.Namespace(metrics.NSOptions{Name: "jaeger"})
			metricsFactory := fork.New("internal",
				jexpvar.NewFactory(10), // backend for internal opts
				baseFactory.Namespace(metrics.NSOptions{Name: "collector"}))

			storageFactory.InitFromViper(v)
			if err := storageFactory.Initialize(baseFactory, logger); err != nil {
				logger.Fatal("Failed to init storage factory", zap.Error(err))
			}
			spanWriter, err := storageFactory.CreateSpanWriter()
			if err != nil {
				logger.Fatal("Failed to create span writer", zap.Error(err))
			}

			strategyStoreFactory.InitFromViper(v)
			if err := strategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
				logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
			}
			strategyStore, err := strategyStoreFactory.CreateStrategyStore()
			if err != nil {
				logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
			}

			c := app.New(&app.CollectorParams{
				ServiceName:    serviceName,
				Logger:         logger,
				MetricsFactory: metricsFactory,
				SpanWriter:     spanWriter,
				StrategyStore:  strategyStore,
				HealthCheck:    svc.HC(),
			})
			collectorOpts := new(app.CollectorOptions).InitFromViper(v)
			if err := c.Start(collectorOpts); err != nil {
				logger.Fatal("Failed to start collector", zap.Error(err))
			}

			svc.RunAndThen(func() {
				if err := c.Close(); err != nil {
					logger.Error("failed to cleanly close the collector", zap.Error(err))
				}
				if closer, ok := spanWriter.(io.Closer); ok {
					err := closer.Close()
					if err != nil {
						logger.Error("failed to close span writer", zap.Error(err))
					}
				}
				if err := storageFactory.Close(); err != nil {
					logger.Error("Failed to close storage factory", zap.Error(err))
				}

			})
			return nil
		},
	}

	command.AddCommand(version.Command())
	command.AddCommand(env.Command())
	command.AddCommand(docs.Command(v))

	config.AddFlags(
		v,
		command,
		svc.AddFlags,
		app.AddFlags,
		storageFactory.AddFlags,
		strategyStoreFactory.AddFlags,
	)

	if err := command.Execute(); err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"crypto/tls"
	"crypto/x509"
	"github.com/mattn/go-isatty"
	"github.com/mitchellh/mapstructure"
	"github.com/pkg/errors"
	"github.com/zrepl/zrepl/logger"
	"io/ioutil"
	"os"
	"time"
)

type LoggingConfig struct {
	Outlets logger.Outlets
}

type MetadataFlags int64

const (
	MetadataTime MetadataFlags = 1 << iota
	MetadataLevel

	MetadataNone MetadataFlags = 0
	MetadataAll  MetadataFlags = ^0
)

func StdoutOutlet() logger.Outlet {
	formatter := &HumanFormatter{}
	writer := os.Stdout
	formatter.SetMetadataFlags(MetadataAll)
	return WriterOutlet{
		formatter, writer,
	}
}

func parseLogging(i interface{}) (c *LoggingConfig, err error) {

	c = &LoggingConfig{}
	c.Outlets = logger.NewOutlets()

	var asList []interface{}
	if err = mapstructure.Decode(i, &asList); err != nil {
		return nil, errors.Wrap(err, "mapstructure error")
	}
	if len(asList) == 0 {
		// Default config
		out := WriterOutlet{&HumanFormatter{}, os.Stdout}
		c.Outlets.Add(out, logger.Warn)
		return
	}

	var syslogOutlets, stdoutOutlets int
	for lei, le := range asList {

		outlet, minLevel, err := parseOutlet(le)
		if err != nil {
			return nil, errors.Wrapf(err, "cannot parse outlet #%d", lei)
		}
		var _ logger.Outlet = WriterOutlet{}
		var _ logger.Outlet = &SyslogOutlet{}
		switch outlet.(type) {
		case *SyslogOutlet:
			syslogOutlets++
		case WriterOutlet:
			stdoutOutlets++
		}

		c.Outlets.Add(outlet, minLevel)

	}

	if syslogOutlets > 1 {
		return nil, errors.Errorf("can only define one 'syslog' outlet")
	}
	if stdoutOutlets > 1 {
		return nil, errors.Errorf("can only define one 'stdout' outlet")
	}

	return c, nil

}

func parseLogFormat(i interface{}) (f EntryFormatter, err error) {
	var is string
	switch j := i.(type) {
	case string:
		is = j
	default:
		return nil, errors.Errorf("invalid log format: wrong type: %T", i)
	}

	switch is {
	case "human":
		return &HumanFormatter{}, nil
	case "logfmt":
		return &LogfmtFormatter{}, nil
	case "json":
		return &JSONFormatter{}, nil
	default:
		return nil, errors.Errorf("invalid log format: '%s'", is)
	}

}

func parseOutlet(i interface{}) (o logger.Outlet, level logger.Level, err error) {

	var in struct {
		Outlet string
		Level  string
		Format string
	}
	if err = mapstructure.Decode(i, &in); err != nil {
		err = errors.Wrap(err, "mapstructure error")
		return
	}
	if in.Outlet == "" || in.Level == "" || in.Format == "" {
		err = errors.Errorf("must specify 'outlet', 'level' and 'format' field")
		return
	}

	minLevel, err := logger.ParseLevel(in.Level)
	if err != nil {
		err = errors.Wrap(err, "cannot parse 'level' field")
		return
	}
	formatter, err := parseLogFormat(in.Format)
	if err != nil {
		err = errors.Wrap(err, "cannot parse")
		return
	}

	switch in.Outlet {
	case "stdout":
		o, err = parseStdoutOutlet(i, formatter)
	case "tcp":
		o, err = parseTCPOutlet(i, formatter)
	case "syslog":
		o, err = parseSyslogOutlet(i, formatter)
	default:
		err = errors.Errorf("unknown outlet type '%s'", in.Outlet)
	}
	return o, minLevel, err

}

func parseStdoutOutlet(i interface{}, formatter EntryFormatter) (WriterOutlet, error) {

	var in struct {
		Time bool
	}
	if err := mapstructure.Decode(i, &in); err != nil {
		return WriterOutlet{}, errors.Wrap(err, "invalid structure for stdout outlet")
	}

	flags := MetadataAll
	writer := os.Stdout
	if !isatty.IsTerminal(writer.Fd()) && !in.Time {
		flags &= ^MetadataTime
	}

	formatter.SetMetadataFlags(flags)
	return WriterOutlet{
		formatter,
		os.Stdout,
	}, nil
}

func parseTCPOutlet(i interface{}, formatter EntryFormatter) (out *TCPOutlet, err error) {

	out = &TCPOutlet{}
	out.Formatter = formatter
	out.Formatter.SetMetadataFlags(MetadataAll)

	var in struct {
		Net           string
		Address       string
		RetryInterval string `mapstructure:"retry_interval"`
		TLS           *struct {
			CA   string
			Cert string
			Key  string
		}
	}
	if err = mapstructure.Decode(i, &in); err != nil {
		return nil, errors.Wrap(err, "mapstructure error")
	}

	out.RetryInterval, err = time.ParseDuration(in.RetryInterval)
	if err != nil {
		return nil, errors.Wrap(err, "cannot parse 'retry_interval'")
	}

	out.Net, out.Address = in.Net, in.Address

	if in.TLS != nil {

		cert, err := tls.LoadX509KeyPair(in.TLS.Cert, in.TLS.Key)
		if err != nil {
			return nil, errors.Wrap(err, "cannot load client cert")
		}

		var rootCAs *x509.CertPool
		if in.TLS.CA == "" {
			if rootCAs, err = x509.SystemCertPool(); err != nil {
				return nil, errors.Wrap(err, "cannot open system cert pool")
			}
		} else {
			rootCAs = x509.NewCertPool()
			rootCAPEM, err := ioutil.ReadFile(in.TLS.CA)
			if err != nil {
				return nil, errors.Wrap(err, "cannot load CA cert")
			}
			if !rootCAs.AppendCertsFromPEM(rootCAPEM) {
				return nil, errors.New("cannot parse CA cert")
			}
		}
		if err != nil && in.TLS.CA == "" {
			return nil, errors.Wrap(err, "cannot load root ca pool")
		}

		out.TLS = &tls.Config{
			Certificates: []tls.Certificate{cert},
			RootCAs:      rootCAs,
		}

		out.TLS.BuildNameToCertificate()
	}

	return

}

func parseSyslogOutlet(i interface{}, formatter EntryFormatter) (out *SyslogOutlet, err error) {

	var in struct {
		RetryInterval string `mapstructure:"retry_interval"`
	}
	if err = mapstructure.Decode(i, &in); err != nil {
		return nil, errors.Wrap(err, "mapstructure error")
	}

	out = &SyslogOutlet{}
	out.Formatter = formatter
	out.Formatter.SetMetadataFlags(MetadataNone)

	out.RetryInterval = 0 // default to 0 as we assume local syslog will just work
	if in.RetryInterval != "" {
		out.RetryInterval, err = time.ParseDuration(in.RetryInterval)
		if err != nil {
			return nil, errors.Wrap(err, "cannot parse 'retry_interval'")
		}
	}

	return
}
<commit_msg>fixup: remove unused StdoutOutlet function<commit_after>package cmd

import (
	"crypto/tls"
	"crypto/x509"
	"github.com/mattn/go-isatty"
	"github.com/mitchellh/mapstructure"
	"github.com/pkg/errors"
	"github.com/zrepl/zrepl/logger"
	"io/ioutil"
	"os"
	"time"
)

type LoggingConfig struct {
	Outlets logger.Outlets
}

type MetadataFlags int64

const (
	MetadataTime MetadataFlags = 1 << iota
	MetadataLevel

	MetadataNone MetadataFlags = 0
	MetadataAll  MetadataFlags = ^0
)

func parseLogging(i interface{}) (c *LoggingConfig, err error) {

	c = &LoggingConfig{}
	c.Outlets = logger.NewOutlets()

	var asList []interface{}
	if err = mapstructure.Decode(i, &asList); err != nil {
		return nil, errors.Wrap(err, "mapstructure error")
	}
	if len(asList) == 0 {
		// Default config
		out := WriterOutlet{&HumanFormatter{}, os.Stdout}
		c.Outlets.Add(out, logger.Warn)
		return
	}

	var syslogOutlets, stdoutOutlets int
	for lei, le := range asList {

		outlet, minLevel, err := parseOutlet(le)
		if err != nil {
			return nil, errors.Wrapf(err, "cannot parse outlet #%d", lei)
		}
		var _ logger.Outlet = WriterOutlet{}
		var _ logger.Outlet = &SyslogOutlet{}
		switch outlet.(type) {
		case *SyslogOutlet:
			syslogOutlets++
		case WriterOutlet:
			stdoutOutlets++
		}

		c.Outlets.Add(outlet, minLevel)

	}

	if syslogOutlets > 1 {
		return nil, errors.Errorf("can only define one 'syslog' outlet")
	}
	if stdoutOutlets > 1 {
		return nil, errors.Errorf("can only define one 'stdout' outlet")
	}

	return c, nil

}

func parseLogFormat(i interface{}) (f EntryFormatter, err error) {
	var is string
	switch j := i.(type) {
	case string:
		is = j
	default:
		return nil, errors.Errorf("invalid log format: wrong type: %T", i)
	}

	switch is {
	case "human":
		return &HumanFormatter{}, nil
	case "logfmt":
		return &LogfmtFormatter{}, nil
	case "json":
		return &JSONFormatter{}, nil
	default:
		return nil, errors.Errorf("invalid log format: '%s'", is)
	}

}

func parseOutlet(i interface{}) (o logger.Outlet, level logger.Level, err error) {

	var in struct {
		Outlet string
		Level  string
		Format string
	}
	if err = mapstructure.Decode(i, &in); err != nil {
		err = errors.Wrap(err, "mapstructure error")
		return
	}
	if in.Outlet == "" || in.Level == "" || in.Format == "" {
		err = errors.Errorf("must specify 'outlet', 'level' and 'format' field")
		return
	}

	minLevel, err := logger.ParseLevel(in.Level)
	if err != nil {
		err = errors.Wrap(err, "cannot parse 'level' field")
		return
	}
	formatter, err := parseLogFormat(in.Format)
	if err != nil {
		err = errors.Wrap(err, "cannot parse")
		return
	}

	switch in.Outlet {
	case "stdout":
		o, err = parseStdoutOutlet(i, formatter)
	case "tcp":
		o, err = parseTCPOutlet(i, formatter)
	case "syslog":
		o, err = parseSyslogOutlet(i, formatter)
	default:
		err = errors.Errorf("unknown outlet type '%s'", in.Outlet)
	}
	return o, minLevel, err

}

func parseStdoutOutlet(i interface{}, formatter EntryFormatter) (WriterOutlet, error) {

	var in struct {
		Time bool
	}
	if err := mapstructure.Decode(i, &in); err != nil {
		return WriterOutlet{}, errors.Wrap(err, "invalid structure for stdout outlet")
	}

	flags := MetadataAll
	writer := os.Stdout
	if !isatty.IsTerminal(writer.Fd()) && !in.Time {
		flags &= ^MetadataTime
	}

	formatter.SetMetadataFlags(flags)
	return WriterOutlet{
		formatter,
		os.Stdout,
	}, nil
}

func parseTCPOutlet(i interface{}, formatter EntryFormatter) (out *TCPOutlet, err error) {

	out = &TCPOutlet{}
	out.Formatter = formatter
	out.Formatter.SetMetadataFlags(MetadataAll)

	var in struct {
		Net           string
		Address       string
		RetryInterval string `mapstructure:"retry_interval"`
		TLS           *struct {
			CA   string
			Cert string
			Key  string
		}
	}
	if err = mapstructure.Decode(i, &in); err != nil {
		return nil, errors.Wrap(err, "mapstructure error")
	}

	out.RetryInterval, err = time.ParseDuration(in.RetryInterval)
	if err != nil {
		return nil, errors.Wrap(err, "cannot parse 'retry_interval'")
	}

	out.Net, out.Address = in.Net, in.Address

	if in.TLS != nil {

		cert, err := tls.LoadX509KeyPair(in.TLS.Cert, in.TLS.Key)
		if err != nil {
			return nil, errors.Wrap(err, "cannot load client cert")
		}

		var rootCAs *x509.CertPool
		if in.TLS.CA == "" {
			if rootCAs, err = x509.SystemCertPool(); err != nil {
				return nil, errors.Wrap(err, "cannot open system cert pool")
			}
		} else {
			rootCAs = x509.NewCertPool()
			rootCAPEM, err := ioutil.ReadFile(in.TLS.CA)
			if err != nil {
				return nil, errors.Wrap(err, "cannot load CA cert")
			}
			if !rootCAs.AppendCertsFromPEM(rootCAPEM) {
				return nil, errors.New("cannot parse CA cert")
			}
		}
		if err != nil && in.TLS.CA == "" {
			return nil, errors.Wrap(err, "cannot load root ca pool")
		}

		out.TLS = &tls.Config{
			Certificates: []tls.Certificate{cert},
			RootCAs:      rootCAs,
		}

		out.TLS.BuildNameToCertificate()
	}

	return

}

func parseSyslogOutlet(i interface{}, formatter EntryFormatter) (out *SyslogOutlet, err error) {

	var in struct {
		RetryInterval string `mapstructure:"retry_interval"`
	}
	if err = mapstructure.Decode(i, &in); err != nil {
		return nil, errors.Wrap(err, "mapstructure error")
	}

	out = &SyslogOutlet{}
	out.Formatter = formatter
	out.Formatter.SetMetadataFlags(MetadataNone)

	out.RetryInterval = 0 // default to 0 as we assume local syslog will just work
	if in.RetryInterval != "" {
		out.RetryInterval, err = time.ParseDuration(in.RetryInterval)
		if err != nil {
			return nil, errors.Wrap(err, "cannot parse 'retry_interval'")
		}
	}

	return
}
<|endoftext|>