<commit_before>package jobsupervisor

import (
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/cloudfoundry/bosh-agent/internal/github.com/pivotal-golang/clock"
	"github.com/cloudfoundry/bosh-agent/internal/github.com/pivotal/go-smtpd/smtpd"

	boshalert "github.com/cloudfoundry/bosh-agent/agent/alert"
	bosherr "github.com/cloudfoundry/bosh-agent/internal/github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-agent/internal/github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-agent/internal/github.com/cloudfoundry/bosh-utils/system"
	boshmonit "github.com/cloudfoundry/bosh-agent/jobsupervisor/monit"
	boshdir "github.com/cloudfoundry/bosh-agent/settings/directories"
)

const monitJobSupervisorLogTag = "monitJobSupervisor"

type monitJobSupervisor struct {
	fs                    boshsys.FileSystem
	runner                boshsys.CmdRunner
	client                boshmonit.Client
	logger                boshlog.Logger
	dirProvider           boshdir.Provider
	jobFailuresServerPort int
	reloadOptions         MonitReloadOptions
	timeService           clock.Clock
}

type MonitReloadOptions struct {
	// Number of times `monit reload` will be executed
	MaxTries int

	// Number of times monit incarnation will be checked
	// for difference after executing `monit reload`
	MaxCheckTries int

	// Length of time between checking for incarnation difference
	DelayBetweenCheckTries time.Duration
}

func NewMonitJobSupervisor(
	fs boshsys.FileSystem,
	runner boshsys.CmdRunner,
	client boshmonit.Client,
	logger boshlog.Logger,
	dirProvider boshdir.Provider,
	jobFailuresServerPort int,
	reloadOptions MonitReloadOptions,
	timeService clock.Clock,
) JobSupervisor {
	return &monitJobSupervisor{
		fs:                    fs,
		runner:                runner,
		client:                client,
		logger:                logger,
		dirProvider:           dirProvider,
		jobFailuresServerPort: jobFailuresServerPort,
		reloadOptions:         reloadOptions,
		timeService:           timeService,
	}
}

func (m monitJobSupervisor) Reload() error {
	var currentIncarnation int

	oldIncarnation, err := m.getIncarnation()
	if err != nil {
		return bosherr.WrapError(err, "Getting monit incarnation")
	}

	// Monit process could be started in the same second as `monit reload` runs
	// so it's ideal for MaxCheckTries * DelayBetweenCheckTries to be greater than 1 sec
	// because monit incarnation id is just a timestamp with 1 sec resolution.
	for reloadI := 0; reloadI < m.reloadOptions.MaxTries; reloadI++ {
		// Exit code or output cannot be trusted
		_, _, _, err := m.runner.RunCommand("monit", "reload")
		if err != nil {
			m.logger.Error(monitJobSupervisorLogTag, "Failed to reload monit %s", err.Error())
		}

		for checkI := 0; checkI < m.reloadOptions.MaxCheckTries; checkI++ {
			currentIncarnation, err = m.getIncarnation()
			if err != nil {
				return bosherr.WrapError(err, "Getting monit incarnation")
			}

			// Incarnation id can decrease or increase because
			// monit uses time(...) and system time can be changed
			if oldIncarnation != currentIncarnation {
				return nil
			}

			m.logger.Debug(
				monitJobSupervisorLogTag,
				"Waiting for monit to reload: before=%d after=%d",
				oldIncarnation, currentIncarnation,
			)

			time.Sleep(m.reloadOptions.DelayBetweenCheckTries)
		}
	}

	return bosherr.Errorf(
		"Failed to reload monit: before=%d after=%d",
		oldIncarnation, currentIncarnation,
	)
}

func (m monitJobSupervisor) Start() error {
	services, err := m.client.ServicesInGroup("vcap")
	if err != nil {
		return bosherr.WrapError(err, "Getting vcap services")
	}

	for _, service := range services {
		m.logger.Debug(monitJobSupervisorLogTag, "Starting service %s", service)
		err = m.client.StartService(service)
		if err != nil {
			return bosherr.WrapErrorf(err, "Starting service %s", service)
		}
	}

	return nil
}

func (m monitJobSupervisor) Stop() error {
	serviceNames, err := m.client.ServicesInGroup("vcap")
	if err != nil {
		return bosherr.WrapError(err, "Getting vcap services")
	}

	for _, serviceName := range serviceNames {
		m.logger.Debug(monitJobSupervisorLogTag, "Stopping service '%s'", serviceName)
		err = m.client.StopService(serviceName)
		if err != nil {
			return bosherr.WrapErrorf(err, "Stopping service '%s'", serviceName)
		}
	}

	servicesToBeStopped := []string{}
	timer := m.timeService.NewTimer(10 * time.Minute)

	for {
		select {
		case <-timer.C():
			return bosherr.Errorf("Timed out waiting for services '%s' to stop after 10 minutes", strings.Join(servicesToBeStopped, ", "))
		default:
		}

		monitStatus, err := m.client.Status()
		if err != nil {
			return bosherr.WrapErrorf(err, "Getting monit status")
		}

		services := monitStatus.ServicesInGroup("vcap")
		servicesToBeStopped = []string{}

		for _, service := range services {
			if service.Monitored || service.Pending {
				servicesToBeStopped = append(servicesToBeStopped, service.Name)
			}

			if service.Errored {
				return bosherr.Errorf("Stopping service '%s' errored with message '%s'", service.Name, service.StatusMessage)
			}
		}

		if len(servicesToBeStopped) == 0 {
			return nil
		}

		m.timeService.Sleep(500 * time.Millisecond)
	}
}

func (m monitJobSupervisor) Unmonitor() error {
	services, err := m.client.ServicesInGroup("vcap")
	if err != nil {
		return bosherr.WrapError(err, "Getting vcap services")
	}

	for _, service := range services {
		m.logger.Debug(monitJobSupervisorLogTag, "Unmonitoring service %s", service)
		err := m.client.UnmonitorService(service)
		if err != nil {
			return bosherr.WrapErrorf(err, "Unmonitoring service %s", service)
		}
	}

	return nil
}

func (m monitJobSupervisor) Status() (status string) {
	status = "running"
	m.logger.Debug(monitJobSupervisorLogTag, "Getting monit status")
	monitStatus, err := m.client.Status()
	if err != nil {
		status = "unknown"
		return
	}

	for _, service := range monitStatus.ServicesInGroup("vcap") {
		if service.Status == "starting" {
			return "starting"
		}
		if !service.Monitored || service.Status != "running" {
			status = "failing"
		}
	}

	return
}

func (m monitJobSupervisor) getIncarnation() (int, error) {
	monitStatus, err := m.client.Status()
	if err != nil {
		return -1, err
	}

	return monitStatus.GetIncarnation()
}

func (m monitJobSupervisor) AddJob(jobName string, jobIndex int, configPath string) error {
	targetFilename := fmt.Sprintf("%04d_%s.monitrc", jobIndex, jobName)
	targetConfigPath := filepath.Join(m.dirProvider.MonitJobsDir(), targetFilename)

	configContent, err := m.fs.ReadFile(configPath)
	if err != nil {
		return bosherr.WrapError(err, "Reading job config from file")
	}

	err = m.fs.WriteFile(targetConfigPath, configContent)
	if err != nil {
		return bosherr.WrapError(err, "Writing to job config file")
	}

	return nil
}

func (m monitJobSupervisor) RemoveAllJobs() error {
	return m.fs.RemoveAll(m.dirProvider.MonitJobsDir())
}

func (m monitJobSupervisor) MonitorJobFailures(handler JobFailureHandler) (err error) {
	alertHandler := func(smtpd.Connection, smtpd.MailAddress) (env smtpd.Envelope, err error) {
		env = &alertEnvelope{
			new(smtpd.BasicEnvelope),
			handler,
			new(boshalert.MonitAlert),
		}
		return
	}

	serv := &smtpd.Server{
		Addr:      fmt.Sprintf("127.0.0.1:%d", m.jobFailuresServerPort),
		OnNewMail: alertHandler,
	}

	err = serv.ListenAndServe()
	if err != nil {
		err = bosherr.WrapError(err, "Listen for SMTP")
	}
	return
}
<commit_msg>Adds a logger debug to Stop loop.<commit_after>package jobsupervisor

import (
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/cloudfoundry/bosh-agent/internal/github.com/pivotal-golang/clock"
	"github.com/cloudfoundry/bosh-agent/internal/github.com/pivotal/go-smtpd/smtpd"

	boshalert "github.com/cloudfoundry/bosh-agent/agent/alert"
	bosherr "github.com/cloudfoundry/bosh-agent/internal/github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-agent/internal/github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-agent/internal/github.com/cloudfoundry/bosh-utils/system"
	boshmonit "github.com/cloudfoundry/bosh-agent/jobsupervisor/monit"
	boshdir "github.com/cloudfoundry/bosh-agent/settings/directories"
)

const monitJobSupervisorLogTag = "monitJobSupervisor"

type monitJobSupervisor struct {
	fs                    boshsys.FileSystem
	runner                boshsys.CmdRunner
	client                boshmonit.Client
	logger                boshlog.Logger
	dirProvider           boshdir.Provider
	jobFailuresServerPort int
	reloadOptions         MonitReloadOptions
	timeService           clock.Clock
}

type MonitReloadOptions struct {
	// Number of times `monit reload` will be executed
	MaxTries int

	// Number of times monit incarnation will be checked
	// for difference after executing `monit reload`
	MaxCheckTries int

	// Length of time between checking for incarnation difference
	DelayBetweenCheckTries time.Duration
}

func NewMonitJobSupervisor(
	fs boshsys.FileSystem,
	runner boshsys.CmdRunner,
	client boshmonit.Client,
	logger boshlog.Logger,
	dirProvider boshdir.Provider,
	jobFailuresServerPort int,
	reloadOptions MonitReloadOptions,
	timeService clock.Clock,
) JobSupervisor {
	return &monitJobSupervisor{
		fs:                    fs,
		runner:                runner,
		client:                client,
		logger:                logger,
		dirProvider:           dirProvider,
		jobFailuresServerPort: jobFailuresServerPort,
		reloadOptions:         reloadOptions,
		timeService:           timeService,
	}
}

func (m monitJobSupervisor) Reload() error {
	var currentIncarnation int

	oldIncarnation, err := m.getIncarnation()
	if err != nil {
		return bosherr.WrapError(err, "Getting monit incarnation")
	}

	// Monit process could be started in the same second as `monit reload` runs
	// so it's ideal for MaxCheckTries * DelayBetweenCheckTries to be greater than 1 sec
	// because monit incarnation id is just a timestamp with 1 sec resolution.
	for reloadI := 0; reloadI < m.reloadOptions.MaxTries; reloadI++ {
		// Exit code or output cannot be trusted
		_, _, _, err := m.runner.RunCommand("monit", "reload")
		if err != nil {
			m.logger.Error(monitJobSupervisorLogTag, "Failed to reload monit %s", err.Error())
		}

		for checkI := 0; checkI < m.reloadOptions.MaxCheckTries; checkI++ {
			currentIncarnation, err = m.getIncarnation()
			if err != nil {
				return bosherr.WrapError(err, "Getting monit incarnation")
			}

			// Incarnation id can decrease or increase because
			// monit uses time(...) and system time can be changed
			if oldIncarnation != currentIncarnation {
				return nil
			}

			m.logger.Debug(
				monitJobSupervisorLogTag,
				"Waiting for monit to reload: before=%d after=%d",
				oldIncarnation, currentIncarnation,
			)

			time.Sleep(m.reloadOptions.DelayBetweenCheckTries)
		}
	}

	return bosherr.Errorf(
		"Failed to reload monit: before=%d after=%d",
		oldIncarnation, currentIncarnation,
	)
}

func (m monitJobSupervisor) Start() error {
	services, err := m.client.ServicesInGroup("vcap")
	if err != nil {
		return bosherr.WrapError(err, "Getting vcap services")
	}

	for _, service := range services {
		m.logger.Debug(monitJobSupervisorLogTag, "Starting service %s", service)
		err = m.client.StartService(service)
		if err != nil {
			return bosherr.WrapErrorf(err, "Starting service %s", service)
		}
	}

	return nil
}

func (m monitJobSupervisor) Stop() error {
	serviceNames, err := m.client.ServicesInGroup("vcap")
	if err != nil {
		return bosherr.WrapError(err, "Getting vcap services")
	}

	for _, serviceName := range serviceNames {
		m.logger.Debug(monitJobSupervisorLogTag, "Stopping service '%s'", serviceName)
		err = m.client.StopService(serviceName)
		if err != nil {
			return bosherr.WrapErrorf(err, "Stopping service '%s'", serviceName)
		}
	}

	servicesToBeStopped := []string{}
	timer := m.timeService.NewTimer(10 * time.Minute)

	for {
		m.logger.Debug(monitJobSupervisorLogTag, "Waiting for services to stop")
		select {
		case <-timer.C():
			return bosherr.Errorf("Timed out waiting for services '%s' to stop after 10 minutes", strings.Join(servicesToBeStopped, ", "))
		default:
		}

		monitStatus, err := m.client.Status()
		if err != nil {
			return bosherr.WrapErrorf(err, "Getting monit status")
		}

		services := monitStatus.ServicesInGroup("vcap")
		servicesToBeStopped = []string{}

		for _, service := range services {
			if service.Monitored || service.Pending {
				servicesToBeStopped = append(servicesToBeStopped, service.Name)
			}

			if service.Errored {
				return bosherr.Errorf("Stopping service '%s' errored with message '%s'", service.Name, service.StatusMessage)
			}
		}

		if len(servicesToBeStopped) == 0 {
			return nil
		}

		m.timeService.Sleep(500 * time.Millisecond)
	}
}

func (m monitJobSupervisor) Unmonitor() error {
	services, err := m.client.ServicesInGroup("vcap")
	if err != nil {
		return bosherr.WrapError(err, "Getting vcap services")
	}

	for _, service := range services {
		m.logger.Debug(monitJobSupervisorLogTag, "Unmonitoring service %s", service)
		err := m.client.UnmonitorService(service)
		if err != nil {
			return bosherr.WrapErrorf(err, "Unmonitoring service %s", service)
		}
	}

	return nil
}

func (m monitJobSupervisor) Status() (status string) {
	status = "running"
	m.logger.Debug(monitJobSupervisorLogTag, "Getting monit status")
	monitStatus, err := m.client.Status()
	if err != nil {
		status = "unknown"
		return
	}

	for _, service := range monitStatus.ServicesInGroup("vcap") {
		if service.Status == "starting" {
			return "starting"
		}
		if !service.Monitored || service.Status != "running" {
			status = "failing"
		}
	}

	return
}

func (m monitJobSupervisor) getIncarnation() (int, error) {
	monitStatus, err := m.client.Status()
	if err != nil {
		return -1, err
	}

	return monitStatus.GetIncarnation()
}

func (m monitJobSupervisor) AddJob(jobName string, jobIndex int, configPath string) error {
	targetFilename := fmt.Sprintf("%04d_%s.monitrc", jobIndex, jobName)
	targetConfigPath := filepath.Join(m.dirProvider.MonitJobsDir(), targetFilename)

	configContent, err := m.fs.ReadFile(configPath)
	if err != nil {
		return bosherr.WrapError(err, "Reading job config from file")
	}

	err = m.fs.WriteFile(targetConfigPath, configContent)
	if err != nil {
		return bosherr.WrapError(err, "Writing to job config file")
	}

	return nil
}

func (m monitJobSupervisor) RemoveAllJobs() error {
	return m.fs.RemoveAll(m.dirProvider.MonitJobsDir())
}

func (m monitJobSupervisor) MonitorJobFailures(handler JobFailureHandler) (err error) {
	alertHandler := func(smtpd.Connection, smtpd.MailAddress) (env smtpd.Envelope, err error) {
		env = &alertEnvelope{
			new(smtpd.BasicEnvelope),
			handler,
			new(boshalert.MonitAlert),
		}
		return
	}

	serv := &smtpd.Server{
		Addr:      fmt.Sprintf("127.0.0.1:%d", m.jobFailuresServerPort),
		OnNewMail: alertHandler,
	}

	err = serv.ListenAndServe()
	if err != nil {
		err = bosherr.WrapError(err, "Listen for SMTP")
	}
	return
}
<|endoftext|>


<commit_before>// A webserver to interface range cluster

package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
	// our packages
	"rangeexpr"
	"rangestore"
	"rangestore/etcdstore"
	"rangestore/filestore"
)

// globals
var store string      // name of the store
var params string     // path for filestore, server string for etcd, etc
var slowlog int       // in ms, log queries slower than these
var etcdroot string   // where does the yarge root start in etcd (useful for shared cluster)
var serveraddr string // server address
var fast bool         // is fast lookup okay
var roptimize bool    // do we have reverse lookup optimization
var debug bool        // debug
var help bool         // help

// future need to closure the function with more data to be passed?
func genericHandlerV1(fn func(http.ResponseWriter, *http.Request, interface{}), s interface{}) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) { // returns void, we don't care!
		fn(w, r, s)
	}
}

// handler for each request
// * log the request
// * log slow queries (clientip will the key to track a request)
func requestHandler(w http.ResponseWriter, r *http.Request, s interface{}) {
	var query string
	var err error

	var remoteaddr = fmt.Sprintf("%s:%s", r.Header.Get("X-Real-IP"), r.Header.Get("X-Real-Port"))
	if remoteaddr == ":" {
		remoteaddr = r.RemoteAddr
	}

	query, err = url.QueryUnescape(r.URL.RawQuery)
	if err != nil {
		log.Println("EROR> [%s] Request: [%s] Error: %s", remoteaddr, r.URL.RawQuery, err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if debug {
		log.Printf("DBUG> [%s] %s", remoteaddr, query)
	}

	// exapand the query
	var results *[]string
	var errs []error

	defer func() {
		if _r := recover(); _r != nil {
			errs = append(errs, errors.New(fmt.Sprintf("[%s] %s [Panicked while Expanding Query]", remoteaddr, query)))
			results = &[]string{}
			log.Printf("[%s] %s [Panicked while Expanding Query]", remoteaddr, query)
		}
	}()

	// measure how long it took
	t0 := time.Now()
	// do the expand
	results, errs = expandQuery(query, s)
	t1 := time.Now()

	timetaken := time.Duration(t1.Sub(t0)) / time.Microsecond

	// return the results to the client
	// set the headers if we have errors
	if len(errs) > 0 {
		w.Header().Set("Range-Err-Count", fmt.Sprintf("%d", len(errs)))
		var _errs = make([]string, 0)
		for _, i := range errs {
			_errs = append(_errs, fmt.Sprintf("%s", i))
		}
		http.Error(w, strings.Join(_errs, ","), http.StatusInternalServerError)
		return
	}

	// set header with time taken to process thne request
	w.Header().Set("Range-Expand-Microsecond", fmt.Sprintf("%d", timetaken))
	// write the results
	_, err = fmt.Fprintf(w, "%s", strings.Join(*results, "\n"))
	if err != nil {
		log.Printf("ERROR> [%s] %s (Writing back to Client Failed [Reason: %s])\n", remoteaddr, query, err)
	}

	isSlow := timetaken > time.Duration(slowlog)*time.Microsecond
	//	log.Println(timetaken, time.Duration(slowlog)*time.Microsecond)
	if debug || isSlow {
		if isSlow {
			log.Printf("INFO> [SLOWQUERY] [%s] %s Result %s [Time: %v]", remoteaddr, query, *results, timetaken)
		} else {
			log.Printf("DBUG> [%s] %s Result %s [Time: %v]", remoteaddr, query, *results, timetaken)
		}
	}

	return
}

func expandQuery(query string, s interface{}) (*[]string, []error) {
	// parse the query
	var yr *rangeexpr.RangeExpr

	yr = &rangeexpr.RangeExpr{Buffer: query}
	// initialize
	yr.Init()
	yr.Expression.Init(query)
	// parse the query
	if err := yr.Parse(); err != nil {
		return &[]string{}, []error{errors.New("Parse Error")}
	}
	// build AST
	yr.Execute()

	// evaluate AST
	return yr.Evaluate(s)
}

func startServer(store interface{}) {
	// handling deploy requests
	http.HandleFunc("/v1/range/", genericHandlerV1(requestHandler, store))
	log.Printf("Range WebServer Started [%s]", serveraddr)
	http.ListenAndServe(serveraddr, nil)
	return
}

// init function to set up whatever state is required
// for real program execution
func init() {
	parseFlags()
	// handle help
	if help == true {
		printHelp()
		os.Exit(0)
	}
	return
}

func main() {
	// set log to get the code location
	log.SetFlags(log.Lshortfile)
	// create an connection to store
	var _store interface{}
	var err error
	switch store {
	case "teststore":
		_store, err = rangestore.ConnectTestStore("Test Store") // this can never return error
	case "filestore":
		var path = params
		var depth = -1
		_store, err = filestore.ConnectFileStore(path, depth, fast)
	case "etcdstore":
		var hosts = []string{params}
		_store, err = etcdstore.ConnectEtcdStore(hosts, roptimize, fast, etcdroot)
	default:
		log.Fatalf(`Unknown store [%s] (Supports only "filestore", "teststore", "etcdstore"\n`, store)
	}
	// if error, exit
	if err != nil {
		log.Fatalf("Error in Connecting to Store", err)
		return
	}

	startServer(_store)
}

// parse the flags
func parseFlags() {
	flag.StringVar(&store, "store", "teststore", "Store Name")
	flag.StringVar(&params, "params", "", "Store Parameters")
	flag.IntVar(&slowlog, "slowlog", 3, "Microseconds definition of Slow Query")
	flag.StringVar(&etcdroot, "etcdroot", "", "Root for Range in Etcd Cluster")
	flag.BoolVar(&fast, "fast", true, "Fast Lookup, return the first result")
	flag.BoolVar(&roptimize, "roptimize", true, "Reverse Lookup Optimization")
	flag.StringVar(&serveraddr, "serveraddr", "0.0.0.0:9999", "Server Address")
	flag.BoolVar(&debug, "debug", false, "Debug")
	flag.BoolVar(&help, "help", false, "Good Ol' Help")

	// parse the options
	flag.Parse()

	return
}

// print Help
func printHelp() {
	fmt.Println(
		`Usage: rangerserver [OPTIONS]
 --store ................ Store Name, it can be "teststore", "filestore" or "etcdstore" (default: "teststore")
 --params ............... Parameters for Store, (default: filestore - /var/yarge/, etcdstore - http://127.0.0.1:4001)
 --slowlog .............. Any Query that takes more than this param in ns will be logged (default: 10ns)
 --etcdroot ............. The yarge node root in etcd, useful for shared cluster (default: "")
 --fast ................. Enable Fast Lookup, return the first result for reverse lookups
 --roptimize ............ Enable reverse lookup optimization
 --serveraddr ........... Server Listening Port (default: 0.0.0.0:9999)
 --debug ................ Debug
 --help ................. Good Ol' Help`,
	)
	return
}
<commit_msg>close the connection to store<commit_after>// A webserver to interface range cluster

package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
	// our packages
	"rangeexpr"
	"rangestore"
	"rangestore/etcdstore"
	"rangestore/filestore"
)

// globals
var store string      // name of the store
var params string     // path for filestore, server string for etcd, etc
var slowlog int       // in ms, log queries slower than these
var etcdroot string   // where does the yarge root start in etcd (useful for shared cluster)
var serveraddr string // server address
var fast bool         // is fast lookup okay
var roptimize bool    // do we have reverse lookup optimization
var debug bool        // debug
var help bool         // help

// future need to closure the function with more data to be passed?
func genericHandlerV1(fn func(http.ResponseWriter, *http.Request, interface{}), s interface{}) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) { // returns void, we don't care!
		fn(w, r, s)
	}
}

// handler for each request
// * log the request
// * log slow queries (clientip will the key to track a request)
func requestHandler(w http.ResponseWriter, r *http.Request, s interface{}) {
	var query string
	var err error

	var remoteaddr = fmt.Sprintf("%s:%s", r.Header.Get("X-Real-IP"), r.Header.Get("X-Real-Port"))
	if remoteaddr == ":" {
		remoteaddr = r.RemoteAddr
	}

	query, err = url.QueryUnescape(r.URL.RawQuery)
	if err != nil {
		log.Println("EROR> [%s] Request: [%s] Error: %s", remoteaddr, r.URL.RawQuery, err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if debug {
		log.Printf("DBUG> [%s] %s", remoteaddr, query)
	}

	// exapand the query
	var results *[]string
	var errs []error

	defer func() {
		if _r := recover(); _r != nil {
			errs = append(errs, errors.New(fmt.Sprintf("[%s] %s [Panicked while Expanding Query]", remoteaddr, query)))
			results = &[]string{}
			log.Printf("[%s] %s [Panicked while Expanding Query]", remoteaddr, query)
		}
	}()

	// measure how long it took
	t0 := time.Now()
	// do the expand
	results, errs = expandQuery(query, s)
	t1 := time.Now()

	timetaken := time.Duration(t1.Sub(t0)) / time.Microsecond

	// return the results to the client
	// set the headers if we have errors
	if len(errs) > 0 {
		w.Header().Set("Range-Err-Count", fmt.Sprintf("%d", len(errs)))
		var _errs = make([]string, 0)
		for _, i := range errs {
			_errs = append(_errs, fmt.Sprintf("%s", i))
		}
		http.Error(w, strings.Join(_errs, ","), http.StatusInternalServerError)
		return
	}

	// set header with time taken to process thne request
	w.Header().Set("Range-Expand-Microsecond", fmt.Sprintf("%d", timetaken))
	// write the results
	_, err = fmt.Fprintf(w, "%s", strings.Join(*results, "\n"))
	if err != nil {
		log.Printf("ERROR> [%s] %s (Writing back to Client Failed [Reason: %s])\n", remoteaddr, query, err)
	}

	isSlow := timetaken > time.Duration(slowlog)*time.Microsecond
	//	log.Println(timetaken, time.Duration(slowlog)*time.Microsecond)
	if debug || isSlow {
		if isSlow {
			log.Printf("INFO> [SLOWQUERY] [%s] %s Result %s [Time: %v]", remoteaddr, query, *results, timetaken)
		} else {
			log.Printf("DBUG> [%s] %s Result %s [Time: %v]", remoteaddr, query, *results, timetaken)
		}
	}

	return
}

func expandQuery(query string, s interface{}) (*[]string, []error) {
	// parse the query
	var yr *rangeexpr.RangeExpr

	yr = &rangeexpr.RangeExpr{Buffer: query}
	// initialize
	yr.Init()
	yr.Expression.Init(query)
	// parse the query
	if err := yr.Parse(); err != nil {
		return &[]string{}, []error{errors.New("Parse Error")}
	}
	// build AST
	yr.Execute()

	// evaluate AST
	return yr.Evaluate(s)
}

func startServer(store interface{}) {
	// handling deploy requests
	http.HandleFunc("/v1/range/", genericHandlerV1(requestHandler, store))
	log.Printf("Range WebServer Started [%s]", serveraddr)
	if err := http.ListenAndServe(serveraddr, nil); err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
	return
}

// init function to set up whatever state is required
// for real program execution
func init() {
	parseFlags()
	// handle help
	if help == true {
		printHelp()
		os.Exit(0)
	}
	return
}

func main() {
	// set log to get the code location
	log.SetFlags(log.Lshortfile)
	// create an connection to store
	var _store interface{}
	var err error
	switch store {
	case "teststore":
		_store, err = rangestore.ConnectTestStore("Test Store") // this can never return error
		defer func() {
			_store.(*rangestore.TestStore).DisconnectTestStore()
		}()
	case "filestore":
		var path = params
		var depth = -1
		_store, err = filestore.ConnectFileStore(path, depth, fast)
		defer func() {
			_store.(*filestore.FileStore).DisconnectFileStore()
		}()
	case "etcdstore":
		var hosts = []string{params}
		_store, err = etcdstore.ConnectEtcdStore(hosts, roptimize, fast, etcdroot)
		defer func() {
			_store.(*etcdstore.EtcdStore).DisconnectEtcdStore()
		}()
	default:
		log.Fatalf(`Unknown store [%s] (Supports only "filestore", "teststore", "etcdstore"\n`, store)
	}
	// if error, exit
	if err != nil {
		log.Fatalf("Error in Connecting to Store", err)
		return
	}

	startServer(_store)
}

// parse the flags
func parseFlags() {
	flag.StringVar(&store, "store", "teststore", "Store Name")
	flag.StringVar(&params, "params", "", "Store Parameters")
	flag.IntVar(&slowlog, "slowlog", 3, "Microseconds definition of Slow Query")
	flag.StringVar(&etcdroot, "etcdroot", "", "Root for Range in Etcd Cluster")
	flag.BoolVar(&fast, "fast", false, "Fast Lookup, return the first result")
	flag.BoolVar(&roptimize, "roptimize", true, "Reverse Lookup Optimization")
	flag.StringVar(&serveraddr, "serveraddr", "0.0.0.0:9999", "Server Address")
	flag.BoolVar(&debug, "debug", false, "Debug")
	flag.BoolVar(&help, "help", false, "Good Ol' Help")

	// parse the options
	flag.Parse()

	return
}

// print Help
func printHelp() {
	fmt.Println(
		`Usage: rangerserver [OPTIONS]
 --store ................ Store Name, it can be "teststore", "filestore" or "etcdstore" (default: "teststore")
 --params ............... Parameters for Store, (default: filestore - /var/yarge/, etcdstore - http://127.0.0.1:4001)
 --slowlog .............. Any Query that takes more than this param in ns will be logged (default: 10ns)
 --etcdroot ............. The yarge node root in etcd, useful for shared cluster (default: "")
 --fast ................. Enable Fast Lookup, return the first result for reverse lookups
 --roptimize ............ Enable reverse lookup optimization
 --serveraddr ........... Server Listening Port (default: 0.0.0.0:9999)
 --debug ................ Debug
 --help ................. Good Ol' Help`,
	)
	return
}
<|endoftext|>


<commit_before>package lib

// Version is the github-nippou version
const Version = "4.1.8"
<commit_msg>:up: Bump version to 4.1.9<commit_after>package lib

// Version is the github-nippou version
const Version = "4.1.9"
<|endoftext|>


<commit_before>// Copyright 2016 The LUCI 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 prpc

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"

	"github.com/golang/protobuf/proto"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/grpc/grpcutil"
	"go.chromium.org/luci/server/router"
)

var (
	// Describe the permitted Access Control requests.
	allowHeaders = strings.Join([]string{"Origin", "Content-Type", "Accept", "Authorization"}, ", ")
	allowMethods = strings.Join([]string{"OPTIONS", "POST"}, ", ")

	// allowPreflightCacheAgeSecs is the amount of time to enable the browser to
	// cache the preflight access control response, in seconds.
	//
	// 600 seconds is 10 minutes.
	allowPreflightCacheAgeSecs = "600"

	// exposeHeaders lists the whitelisted non-standard response headers that the
	// client may accept.
	exposeHeaders = strings.Join([]string{HeaderGRPCCode}, ", ")

	// NoAuthentication can be used in place of an Authenticator to explicitly
	// specify that your Server will skip authentication.
	//
	// Use it with Server.Authenticator or RegisterDefaultAuth.
	NoAuthentication Authenticator = nullAuthenticator{}
)

// AllowOriginAll returns true unconditionally.
// It can be used as Server.AccessControl.
// It must NOT be used in combination with cookie-based authentication.
func AllowOriginAll(c context.Context, origin string) bool {
	return true
}

// Override is a pRPC method-specific override which may optionally handle the
// entire pRPC method call. If it returns true, the override is assumed to
// have fully handled the pRPC method call and processing of the request does
// not continue. In this case it's the override's responsibility to adhere to
// all pRPC semantics. However if it returns false, processing continues as
// normal, allowing the override to act as a preprocessor. In this case it's
// the override's responsibility to ensure it hasn't done anything that will
// be incompatible with pRPC semantics (such as writing garbage to the response
// writer in the router context).
type Override func(*router.Context) bool

// Server is a pRPC server to serve RPC requests.
// Zero value is valid.
type Server struct {
	// Authenticator, if not nil, specifies how to authenticate requests.
	//
	// If nil, the default authenticator set by RegisterDefaultAuth will be used.
	// If the default authenticator is also nil, all request handlers will panic.
	//
	// If you want to disable the authentication (e.g for tests), explicitly set
	// Authenticator to NoAuthentication.
	Authenticator Authenticator

	// AccessControl, if not nil, is a callback that is invoked per request to
	// determine if permissive access control headers should be added to the
	// response.
	//
	// This callback includes the request Context and the origin header supplied
	// by the client. If nil, or if it returns false, no headers will be written.
	// Otherwise, access control headers for the specified origin will be
	// included in the response.
	AccessControl func(c context.Context, origin string) bool

	// UnaryServerInterceptor provides a hook to intercept the execution of
	// a unary RPC on the server. It is the responsibility of the interceptor to
	// invoke handler to complete the RPC.
	UnaryServerInterceptor grpc.UnaryServerInterceptor

	mu        sync.Mutex
	services  map[string]*service
	overrides map[string]map[string]Override
}

type service struct {
	methods map[string]grpc.MethodDesc
	impl    interface{}
}

// RegisterService registers a service implementation.
// Called from the generated code.
//
// desc must contain description of the service, its message types
// and all transitive dependencies.
//
// Panics if a service of the same name is already registered.
func (s *Server) RegisterService(desc *grpc.ServiceDesc, impl interface{}) {
	serv := &service{
		impl:    impl,
		methods: make(map[string]grpc.MethodDesc, len(desc.Methods)),
	}
	for _, m := range desc.Methods {
		serv.methods[m.MethodName] = m
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if s.services == nil {
		s.services = map[string]*service{}
	} else if _, ok := s.services[desc.ServiceName]; ok {
		panic(fmt.Errorf("service %q is already registered", desc.ServiceName))
	}

	s.services[desc.ServiceName] = serv
}

// RegisterOverride registers an overriding function.
//
// Panics if an override for the given service method is already registered.
func (s *Server) RegisterOverride(serviceName, methodName string, fn Override) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.overrides == nil {
		s.overrides = map[string]map[string]Override{}
	}
	if _, ok := s.overrides[serviceName]; !ok {
		s.overrides[serviceName] = map[string]Override{}
	}
	if _, ok := s.overrides[serviceName][methodName]; ok {
		panic(fmt.Errorf("method %q of service %q is already overridden", methodName, serviceName))
	}

	s.overrides[serviceName][methodName] = fn
}

// authenticate forces authentication set by RegisterDefaultAuth.
func (s *Server) authenticate() router.Middleware {
	a := s.Authenticator
	if a == nil {
		a = GetDefaultAuth()
		if a == nil {
			panic("prpc: no custom Authenticator was provided and default authenticator was not registered.\n" +
				"Either explicitly set `Server.Authenticator = NoAuthentication`, or use RegisterDefaultAuth()")
		}
	}

	return func(c *router.Context, next router.Handler) {

		ctx, err := a.Authenticate(c.Context, c.Request)
		if err == nil {
			c.Context = ctx
			next(c)
			return
		}

		format, perr := responseFormat(c.Request.Header.Get(headerAccept))
		switch {
		case perr != nil:
			writeError(c.Context, c.Writer, perr, FormatBinary)
		case transient.Tag.In(err):
			writeError(c.Context, c.Writer, withCode(err, codes.Internal), format)
		default:
			writeError(c.Context, c.Writer, withCode(err, codes.Unauthenticated), format)
		}
	}
}

// InstallHandlers installs HTTP handlers at /prpc/:service/:method.
//
// See https://godoc.org/go.chromium.org/luci/grpc/prpc#hdr-Protocol
// for pRPC protocol.
//
// The authenticator in 'base' is always replaced by pRPC specific one. For more
// details about the authentication see Server.Authenticator doc.
func (s *Server) InstallHandlers(r *router.Router, base router.MiddlewareChain) {
	s.mu.Lock()
	defer s.mu.Unlock()

	rr := r.Subrouter("/prpc/:service/:method")
	rr.Use(base.Extend(s.authenticate()))

	rr.POST("", router.MiddlewareChain{}, s.handlePOST)
	rr.OPTIONS("", router.MiddlewareChain{}, s.handleOPTIONS)
}

// handle handles RPCs.
// See https://godoc.org/go.chromium.org/luci/grpc/prpc#hdr-Protocol
// for pRPC protocol.
func (s *Server) handlePOST(c *router.Context) {
	serviceName := c.Params.ByName("service")
	methodName := c.Params.ByName("method")

	if methods, ok := s.overrides[serviceName]; ok {
		if handler, ok := methods[methodName]; ok {
			if handler(c) {
				return
			}
		}
	}

	s.setAccessControlHeaders(c, false)
	res := s.call(c, serviceName, methodName)

	c.Writer.Header().Set("X-Content-Type-Options", "nosniff")

	if res.err != nil {
		writeError(c.Context, c.Writer, res.err, res.fmt)
		return
	}
	writeMessage(c.Context, c.Writer, res.out, res.fmt)
}

func (s *Server) handleOPTIONS(c *router.Context) {
	s.setAccessControlHeaders(c, true)
	c.Writer.WriteHeader(http.StatusOK)
}

var requestContextKey = "context key with *requestContext"

type requestContext struct {
	// additional headers that will be sent in the response
	header http.Header
}

// SetHeader sets the header metadata.
// When called multiple times, all the provided metadata will be merged.
//
// If ctx is not a pRPC server context, then SetHeader calls grpc.SetHeader
// such that calling prpc.SetHeader works for both pRPC and gRPC.
func SetHeader(ctx context.Context, md metadata.MD) error {
	if rctx, ok := ctx.Value(&requestContextKey).(*requestContext); ok {
		for k, vs := range md {
			if strings.HasPrefix(k, "X-Prpc-") || k == headerContentType {
				return errors.Reason("reserved header key %q", k).Err()
			}
			for _, v := range vs {
				rctx.header.Add(metaToHeader(k, v))
			}
		}
		return nil
	}

	return grpc.SetHeader(ctx, md)
}

type response struct {
	out proto.Message
	fmt Format
	err error
}

func (s *Server) call(c *router.Context, serviceName, methodName string) (r response) {
	service := s.services[serviceName]
	if service == nil {
		r.err = status.Errorf(
			codes.Unimplemented,
			"service %q is not implemented",
			serviceName)
		return
	}

	method, ok := service.methods[methodName]
	if !ok {
		r.err = status.Errorf(
			codes.Unimplemented,
			"method %q in service %q is not implemented",
			methodName, serviceName)
		return
	}

	var perr *protocolError
	r.fmt, perr = responseFormat(c.Request.Header.Get(headerAccept))
	if perr != nil {
		r.err = perr
		return
	}

	methodCtx, err := parseHeader(c.Context, c.Request.Header, c.Request.Host)
	if err != nil {
		r.err = withStatus(err, http.StatusBadRequest)
		return
	}

	methodCtx = context.WithValue(methodCtx, &requestContextKey, &requestContext{header: c.Writer.Header()})

	out, err := method.Handler(service.impl, methodCtx, func(in interface{}) error {
		if in == nil {
			return grpcutil.Errf(codes.Internal, "input message is nil")
		}
		// Do not collapse it to one line. There is implicit err type conversion.
		if perr := readMessage(c.Request, in.(proto.Message)); perr != nil {
			return perr
		}
		return nil
	}, s.UnaryServerInterceptor)

	switch {
	case err != nil:
		r.err = err
	case out == nil:
		r.err = status.Error(codes.Internal, "service returned nil message")
	default:
		r.out = out.(proto.Message)
	}
	return
}

func (s *Server) setAccessControlHeaders(c *router.Context, preflight bool) {
	// Don't write out access control headers if the origin is unspecified.
	const originHeader = "Origin"
	origin := c.Request.Header.Get(originHeader)
	if origin == "" || s.AccessControl == nil || !s.AccessControl(c.Context, origin) {
		return
	}

	h := c.Writer.Header()
	h.Add("Access-Control-Allow-Origin", origin)
	h.Add("Vary", originHeader)
	h.Add("Access-Control-Allow-Credentials", "true")

	if preflight {
		h.Add("Access-Control-Allow-Headers", allowHeaders)
		h.Add("Access-Control-Allow-Methods", allowMethods)
		h.Add("Access-Control-Max-Age", allowPreflightCacheAgeSecs)
	} else {
		h.Add("Access-Control-Expose-Headers", exposeHeaders)
	}
}

// ServiceNames returns a sorted list of full names of all registered services.
func (s *Server) ServiceNames() []string {
	s.mu.Lock()
	defer s.mu.Unlock()

	names := make([]string, 0, len(s.services))
	for name := range s.services {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}
<commit_msg>prpc: make prpc.Server truly concurrency-safe.<commit_after>// Copyright 2016 The LUCI 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 prpc

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"

	"github.com/golang/protobuf/proto"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/grpc/grpcutil"
	"go.chromium.org/luci/server/router"
)

var (
	// Describe the permitted Access Control requests.
	allowHeaders = strings.Join([]string{"Origin", "Content-Type", "Accept", "Authorization"}, ", ")
	allowMethods = strings.Join([]string{"OPTIONS", "POST"}, ", ")

	// allowPreflightCacheAgeSecs is the amount of time to enable the browser to
	// cache the preflight access control response, in seconds.
	//
	// 600 seconds is 10 minutes.
	allowPreflightCacheAgeSecs = "600"

	// exposeHeaders lists the whitelisted non-standard response headers that the
	// client may accept.
	exposeHeaders = strings.Join([]string{HeaderGRPCCode}, ", ")

	// NoAuthentication can be used in place of an Authenticator to explicitly
	// specify that your Server will skip authentication.
	//
	// Use it with Server.Authenticator or RegisterDefaultAuth.
	NoAuthentication Authenticator = nullAuthenticator{}
)

// AllowOriginAll returns true unconditionally.
// It can be used as Server.AccessControl.
// It must NOT be used in combination with cookie-based authentication.
func AllowOriginAll(c context.Context, origin string) bool {
	return true
}

// Override is a pRPC method-specific override which may optionally handle the
// entire pRPC method call. If it returns true, the override is assumed to
// have fully handled the pRPC method call and processing of the request does
// not continue. In this case it's the override's responsibility to adhere to
// all pRPC semantics. However if it returns false, processing continues as
// normal, allowing the override to act as a preprocessor. In this case it's
// the override's responsibility to ensure it hasn't done anything that will
// be incompatible with pRPC semantics (such as writing garbage to the response
// writer in the router context).
type Override func(*router.Context) bool

// Server is a pRPC server to serve RPC requests.
// Zero value is valid.
type Server struct {
	// Authenticator, if not nil, specifies how to authenticate requests.
	//
	// If nil, the default authenticator set by RegisterDefaultAuth will be used.
	// If the default authenticator is also nil, all request handlers will panic.
	//
	// If you want to disable the authentication (e.g for tests), explicitly set
	// Authenticator to NoAuthentication.
	Authenticator Authenticator

	// AccessControl, if not nil, is a callback that is invoked per request to
	// determine if permissive access control headers should be added to the
	// response.
	//
	// This callback includes the request Context and the origin header supplied
	// by the client. If nil, or if it returns false, no headers will be written.
	// Otherwise, access control headers for the specified origin will be
	// included in the response.
	AccessControl func(c context.Context, origin string) bool

	// UnaryServerInterceptor provides a hook to intercept the execution of
	// a unary RPC on the server. It is the responsibility of the interceptor to
	// invoke handler to complete the RPC.
	UnaryServerInterceptor grpc.UnaryServerInterceptor

	mu        sync.RWMutex
	services  map[string]*service
	overrides map[string]map[string]Override
}

type service struct {
	methods map[string]grpc.MethodDesc
	impl    interface{}
}

// RegisterService registers a service implementation.
// Called from the generated code.
//
// desc must contain description of the service, its message types
// and all transitive dependencies.
//
// Panics if a service of the same name is already registered.
func (s *Server) RegisterService(desc *grpc.ServiceDesc, impl interface{}) {
	serv := &service{
		impl:    impl,
		methods: make(map[string]grpc.MethodDesc, len(desc.Methods)),
	}
	for _, m := range desc.Methods {
		serv.methods[m.MethodName] = m
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if s.services == nil {
		s.services = map[string]*service{}
	} else if _, ok := s.services[desc.ServiceName]; ok {
		panic(fmt.Errorf("service %q is already registered", desc.ServiceName))
	}

	s.services[desc.ServiceName] = serv
}

// RegisterOverride registers an overriding function.
//
// Panics if an override for the given service method is already registered.
func (s *Server) RegisterOverride(serviceName, methodName string, fn Override) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.overrides == nil {
		s.overrides = map[string]map[string]Override{}
	}
	if _, ok := s.overrides[serviceName]; !ok {
		s.overrides[serviceName] = map[string]Override{}
	}
	if _, ok := s.overrides[serviceName][methodName]; ok {
		panic(fmt.Errorf("method %q of service %q is already overridden", methodName, serviceName))
	}

	s.overrides[serviceName][methodName] = fn
}

// authenticate forces authentication set by RegisterDefaultAuth.
func (s *Server) authenticate() router.Middleware {
	a := s.Authenticator
	if a == nil {
		a = GetDefaultAuth()
		if a == nil {
			panic("prpc: no custom Authenticator was provided and default authenticator was not registered.\n" +
				"Either explicitly set `Server.Authenticator = NoAuthentication`, or use RegisterDefaultAuth()")
		}
	}

	return func(c *router.Context, next router.Handler) {

		ctx, err := a.Authenticate(c.Context, c.Request)
		if err == nil {
			c.Context = ctx
			next(c)
			return
		}

		format, perr := responseFormat(c.Request.Header.Get(headerAccept))
		switch {
		case perr != nil:
			writeError(c.Context, c.Writer, perr, FormatBinary)
		case transient.Tag.In(err):
			writeError(c.Context, c.Writer, withCode(err, codes.Internal), format)
		default:
			writeError(c.Context, c.Writer, withCode(err, codes.Unauthenticated), format)
		}
	}
}

// InstallHandlers installs HTTP handlers at /prpc/:service/:method.
//
// See https://godoc.org/go.chromium.org/luci/grpc/prpc#hdr-Protocol
// for pRPC protocol.
//
// The authenticator in 'base' is always replaced by pRPC specific one. For more
// details about the authentication see Server.Authenticator doc.
func (s *Server) InstallHandlers(r *router.Router, base router.MiddlewareChain) {
	s.mu.Lock()
	defer s.mu.Unlock()

	rr := r.Subrouter("/prpc/:service/:method")
	rr.Use(base.Extend(s.authenticate()))

	rr.POST("", router.MiddlewareChain{}, s.handlePOST)
	rr.OPTIONS("", router.MiddlewareChain{}, s.handleOPTIONS)
}

// handle handles RPCs.
// See https://godoc.org/go.chromium.org/luci/grpc/prpc#hdr-Protocol
// for pRPC protocol.
func (s *Server) handlePOST(c *router.Context) {
	serviceName := c.Params.ByName("service")
	methodName := c.Params.ByName("method")

	override, service, method, methodFound := s.lookup(serviceName, methodName)
	// Override takes precedence over notImplementedErr.
	if override != nil && override(c) {
		return
	}

	s.setAccessControlHeaders(c, false)
	res := response{}
	switch {
	case service == nil:
		res.err = status.Errorf(
			codes.Unimplemented,
			"service %q is not implemented",
			serviceName)
	case !methodFound:
		res.err = status.Errorf(
			codes.Unimplemented,
			"method %q in service %q is not implemented",
			methodName, serviceName)
	default:
		s.call(c, service, method, &res)
	}
	c.Writer.Header().Set("X-Content-Type-Options", "nosniff")

	if res.err != nil {
		writeError(c.Context, c.Writer, res.err, res.fmt)
		return
	}
	writeMessage(c.Context, c.Writer, res.out, res.fmt)
}

func (s *Server) handleOPTIONS(c *router.Context) {
	s.setAccessControlHeaders(c, true)
	c.Writer.WriteHeader(http.StatusOK)
}

var requestContextKey = "context key with *requestContext"

type requestContext struct {
	// additional headers that will be sent in the response
	header http.Header
}

// SetHeader sets the header metadata.
// When called multiple times, all the provided metadata will be merged.
//
// If ctx is not a pRPC server context, then SetHeader calls grpc.SetHeader
// such that calling prpc.SetHeader works for both pRPC and gRPC.
func SetHeader(ctx context.Context, md metadata.MD) error {
	if rctx, ok := ctx.Value(&requestContextKey).(*requestContext); ok {
		for k, vs := range md {
			if strings.HasPrefix(k, "X-Prpc-") || k == headerContentType {
				return errors.Reason("reserved header key %q", k).Err()
			}
			for _, v := range vs {
				rctx.header.Add(metaToHeader(k, v))
			}
		}
		return nil
	}

	return grpc.SetHeader(ctx, md)
}

type response struct {
	out proto.Message
	fmt Format
	err error
}

func (s *Server) lookup(serviceName, methodName string) (override Override, service *service, method grpc.MethodDesc, methodFound bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if methods, ok := s.overrides[serviceName]; ok {
		override = methods[methodName]
	}
	service = s.services[serviceName]
	if service == nil {
		return
	}
	method, methodFound = service.methods[methodName]
	return
}

func (s *Server) call(c *router.Context, service *service, method grpc.MethodDesc, r *response) {
	var perr *protocolError
	r.fmt, perr = responseFormat(c.Request.Header.Get(headerAccept))
	if perr != nil {
		r.err = perr
		return
	}

	methodCtx, err := parseHeader(c.Context, c.Request.Header, c.Request.Host)
	if err != nil {
		r.err = withStatus(err, http.StatusBadRequest)
		return
	}

	methodCtx = context.WithValue(methodCtx, &requestContextKey, &requestContext{header: c.Writer.Header()})

	out, err := method.Handler(service.impl, methodCtx, func(in interface{}) error {
		if in == nil {
			return grpcutil.Errf(codes.Internal, "input message is nil")
		}
		// Do not collapse it to one line. There is implicit err type conversion.
		if perr := readMessage(c.Request, in.(proto.Message)); perr != nil {
			return perr
		}
		return nil
	}, s.UnaryServerInterceptor)

	switch {
	case err != nil:
		r.err = err
	case out == nil:
		r.err = status.Error(codes.Internal, "service returned nil message")
	default:
		r.out = out.(proto.Message)
	}
	return
}

func (s *Server) setAccessControlHeaders(c *router.Context, preflight bool) {
	// Don't write out access control headers if the origin is unspecified.
	const originHeader = "Origin"
	origin := c.Request.Header.Get(originHeader)
	if origin == "" || s.AccessControl == nil || !s.AccessControl(c.Context, origin) {
		return
	}

	h := c.Writer.Header()
	h.Add("Access-Control-Allow-Origin", origin)
	h.Add("Vary", originHeader)
	h.Add("Access-Control-Allow-Credentials", "true")

	if preflight {
		h.Add("Access-Control-Allow-Headers", allowHeaders)
		h.Add("Access-Control-Allow-Methods", allowMethods)
		h.Add("Access-Control-Max-Age", allowPreflightCacheAgeSecs)
	} else {
		h.Add("Access-Control-Expose-Headers", exposeHeaders)
	}
}

// ServiceNames returns a sorted list of full names of all registered services.
func (s *Server) ServiceNames() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	names := make([]string, 0, len(s.services))
	for name := range s.services {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}
<|endoftext|>


<commit_before><commit_msg>make new threads appear when using null cache<commit_after><|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"

	"github.com/gin-gonic/contrib/gzip"
	"github.com/gin-gonic/gin"
	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
	"github.com/nu7hatch/gouuid"

	"github.com/Depado/goploader/server/conf"
	"github.com/Depado/goploader/server/models"
	"github.com/Depado/goploader/server/monitoring"
	"github.com/Depado/goploader/server/utils"
)

var db gorm.DB

func index(c *gin.Context) {
	log.Printf("[INFO][%s]\tIssued a GET request\n", c.ClientIP())
	c.HTML(http.StatusOK, "index.html", gin.H{})
}

func create(c *gin.Context) {
	var err error
	var u *uuid.UUID
	remote := c.ClientIP()
	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 20000000)

	fd, h, err := c.Request.FormFile("file")
	if err != nil {
		log.Printf("[ERROR][%s]\tDuring reading file : %s", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	defer fd.Close()
	if u, err = uuid.NewV4(); err != nil {
		log.Printf("[ERROR][%s]\tDuring creation of uuid : %s\n", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	path := path.Join(conf.C.UploadDir, u.String())
	file, err := os.Create(path)
	if err != nil {
		log.Printf("[ERROR][%s]\tDuring file creation : %s\n", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	defer file.Close()
	wr, err := io.Copy(file, bufio.NewReaderSize(fd, 512))
	if err != nil {
		log.Printf("[ERROR][%s]\tDuring writing file : %s\n", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	db.Create(&models.ResourceEntry{Key: u.String(), Name: h.Filename})
	log.Printf("[INFO][%s]\tCreated %s file and entry (%v bytes written)\n", remote, u.String(), wr)
	c.Writer.WriteHeader(201)
	c.Writer.Write([]byte("https://" + conf.C.NameServer + "/v/" + u.String() + "\n"))
}

func view(c *gin.Context) {
	id := c.Param("uuid")
	re := models.ResourceEntry{}
	remote := c.ClientIP()

	db.Where(&models.ResourceEntry{Key: id}).First(&re)
	if re.Key == "" {
		log.Printf("[INFO][%s]\tNot found : %s", remote, id)
		c.AbortWithStatus(http.StatusNotFound)
		return
	}
	log.Printf("[INFO][%s]\tFetched %s file and entry\n", remote, id)
	f, err := os.Open(conf.C.UploadDir + re.Key)
	if err != nil {
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	c.Header("Content-Disposition", "filename=\""+re.Name+"\"")
	http.ServeContent(c.Writer, c.Request, re.Key, re.CreatedAt, f)
}

func main() {
	var err error

	confPath := flag.String("c", "conf.yml", "Local path to configuration file.")
	flag.Parse()

	if err = conf.Load(*confPath); err != nil {
		log.Fatal(err)
	}
	if err = utils.EnsureDir(conf.C.UploadDir); err != nil {
		log.Fatal(err)
	}
	if db, err = gorm.Open("sqlite3", conf.C.DB); err != nil {
		log.Fatal(err)
	}
	db.AutoMigrate(&models.ResourceEntry{})
	go monitoring.Monit(&db)

	log.Printf("[INFO][System]\tStarted goploader server on port %d\n", conf.C.Port)
	gin.SetMode(gin.ReleaseMode)
	r := gin.Default()
	r.Use(gzip.Gzip(gzip.DefaultCompression))
	r.LoadHTMLGlob("templates/*")
	r.Static("/static", "./assets")
	r.Static("/favicon.ico", "./assets/favicon.ico")
	r.GET("/", index)
	r.POST("/", create)
	r.GET("/v/:uuid", view)
	r.Run(fmt.Sprintf(":%d", conf.C.Port))
}
<commit_msg>Server code using uniuri<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"

	"github.com/dchest/uniuri"
	"github.com/gin-gonic/contrib/gzip"
	"github.com/gin-gonic/gin"
	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"

	"github.com/Depado/goploader/server/conf"
	"github.com/Depado/goploader/server/models"
	"github.com/Depado/goploader/server/monitoring"
	"github.com/Depado/goploader/server/utils"
)

var db gorm.DB

func index(c *gin.Context) {
	log.Printf("[INFO][%s]\tIssued a GET request\n", c.ClientIP())
	c.HTML(http.StatusOK, "index.html", gin.H{})
}

func create(c *gin.Context) {
	var err error
	remote := c.ClientIP()
	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, conf.C.LimitSize*1000000)

	fd, h, err := c.Request.FormFile("file")
	if err != nil {
		log.Printf("[ERROR][%s]\tDuring reading file : %s", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	defer fd.Close()

	u := uniuri.NewLen(conf.C.UniURILength)
	path := path.Join(conf.C.UploadDir, u)
	file, err := os.Create(path)
	if err != nil {
		log.Printf("[ERROR][%s]\tDuring file creation : %s\n", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	defer file.Close()
	wr, err := io.Copy(file, bufio.NewReaderSize(fd, 512))
	if err != nil {
		log.Printf("[ERROR][%s]\tDuring writing file : %s\n", remote, err)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	db.Create(&models.ResourceEntry{Key: u, Name: h.Filename})

	log.Printf("[INFO][%s]\tCreated %s file and entry (%v bytes written)\n", remote, u, wr)
	c.Writer.WriteHeader(201)
	c.Writer.Write([]byte("https://" + conf.C.NameServer + "/v/" + u + "\n"))
}

func view(c *gin.Context) {
	id := c.Param("uniuri")
	re := models.ResourceEntry{}
	remote := c.ClientIP()

	db.Where(&models.ResourceEntry{Key: id}).First(&re)
	if re.Key == "" {
		log.Printf("[INFO][%s]\tNot found : %s", remote, id)
		c.AbortWithStatus(http.StatusNotFound)
		return
	}
	log.Printf("[INFO][%s]\tFetched %s file and entry\n", remote, id)
	f, err := os.Open(conf.C.UploadDir + re.Key)
	if err != nil {
		log.Printf("[ERROR][%s]\tWhile opening %s file\n", remote, id)
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	c.Header("Content-Disposition", "filename=\""+re.Name+"\"")
	http.ServeContent(c.Writer, c.Request, re.Key, re.CreatedAt, f)
}

func main() {
	var err error

	confPath := flag.String("c", "conf.yml", "Local path to configuration file.")
	flag.Parse()

	if err = conf.Load(*confPath); err != nil {
		log.Fatal(err)
	}
	if err = utils.EnsureDir(conf.C.UploadDir); err != nil {
		log.Fatal(err)
	}
	if db, err = gorm.Open("sqlite3", conf.C.DB); err != nil {
		log.Fatal(err)
	}
	db.AutoMigrate(&models.ResourceEntry{})

	go monitoring.Monit(&db)

	log.Printf("[INFO][System]\tStarted goploader server on port %d\n", conf.C.Port)
	gin.SetMode(gin.ReleaseMode)
	// Default router
	r := gin.Default()
	// Middlewares Initialization
	r.Use(gzip.Gzip(gzip.DefaultCompression))
	// Templates and static files
	r.LoadHTMLGlob("templates/*")
	r.Static("/static", "./assets")
	r.Static("/favicon.ico", "./assets/favicon.ico")
	// Routes
	r.GET("/", index)
	r.POST("/", create)
	r.GET("/v/:uniuri", view)
	// Run
	r.Run(fmt.Sprintf(":%d", conf.C.Port))
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"github.com/ActiveState/log"
	"github.com/apcera/nats"
	"time"
)

// NewNatsClient connects to the NATS server of the Stackato cluster
func NewNatsClient(retries int) *nats.EncodedConn {
	servers, err := getNatsServers()
	if err != nil {
		log.Fatalf("Unable to get Nats URI: %v", err)
	}
	log.Infof("Connecting to NATS servers %s\n", servers)

	var nc *nats.Conn
	opts := nats.DefaultOptions
	opts.Servers = servers
	// opts.Secure = true

	for attempt := 0; attempt < retries; attempt++ {
		nc, err = opts.Connect()
		if err != nil {
			if (attempt + 1) == retries {
				log.Fatal(err)
			}
			log.Warnf("NATS connection error (%v); retrying after 1 second..",
				err)
			time.Sleep(time.Second)
		}
	}

	log.Infof("Connected to NATS servers %s\n", servers)
	client, err := nats.NewEncodedConn(nc, "json")
	if err != nil {
		log.Fatal(err)
	}

	// Diagnosing Bug #97856 by periodically checking if we are still
	// connected to NATS.
	go func() {
		log.Info("Periodically checking NATS connectivity")
		for _ = range time.Tick(1 * time.Minute) {
			if nc.IsClosed() {
				log.Fatal("Connection to NATS has been closed (in the last minute)")
			}
		}
	}()

	return client
}

func getNatsServers() ([]string, error) {
	var ipaddr string
	// Use non-lookback address on a micro cloud to connect from docker
	// container to host.
	if InsideDocker() && GetClusterConfig().IsMicro() {
		var err error
		ipaddr, err = GetDockerHostIp()
		if ipaddr == "" {
			return nil, err
		}
	} else {
		ipaddr = GetClusterConfig().MbusIp
	}

	// HACK: Ideally we should be reading NatsUri from
	// cloud_controller config (mbus). we take a shortcut here in
	// order to not have to create a separate ConfDis instance for
	// cloud_controller config (and having to watch it). This will
	// have to change if we switch to clustered version of NATS.
	uri := fmt.Sprintf("nats://%s:4222/", ipaddr)

	return []string{uri}, nil
}
<commit_msg>nats: connect to all nats servers in the cluster<commit_after>package server

import (
	"fmt"
	"github.com/ActiveState/log"
	"github.com/apcera/nats"
	"time"
)

// NewNatsClient connects to the NATS server of the Stackato cluster
func NewNatsClient(retries int) *nats.EncodedConn {
	servers, err := getNatsServers()
	if err != nil {
		log.Fatalf("Unable to get Nats URI: %v", err)
	}
	log.Infof("Connecting to NATS servers %s\n", servers)

	var nc *nats.Conn
	opts := nats.DefaultOptions
	opts.Servers = servers
	// opts.Secure = true

	for attempt := 0; attempt < retries; attempt++ {
		nc, err = opts.Connect()
		if err != nil {
			if (attempt + 1) == retries {
				log.Fatal(err)
			}
			log.Warnf("NATS connection error (%v); retrying after 1 second..",
				err)
			time.Sleep(time.Second)
		}
	}

	log.Infof("Connected to NATS servers %s\n", servers)
	client, err := nats.NewEncodedConn(nc, "json")
	if err != nil {
		log.Fatal(err)
	}

	// Diagnosing Bug #97856 by periodically checking if we are still
	// connected to NATS.
	go func() {
		log.Info("Periodically checking NATS connectivity")
		for _ = range time.Tick(1 * time.Minute) {
			if nc.IsClosed() {
				log.Fatal("Connection to NATS has been closed (in the last minute)")
			}
		}
	}()

	return client
}

func getNatsServers() ([]string, error) {
	ipaddrs := []string{}

	// Use non-lookback address on a micro cloud to connect from docker
	// container to host.
	if InsideDocker() && GetClusterConfig().IsMicro() {
		ipaddr, err := GetDockerHostIp()
		if ipaddr == "" {
			return nil, err
		}
		ipaddrs = append(ipaddrs, ipaddr)
	} else {
		ipaddrs = getNodesWithNatsRunning()
	}

	// HACK: Ideally we should be reading NatsUri from
	// cloud_controller config (mbus). we take a shortcut here in
	// order to not have to create a separate ConfDis instance for
	// cloud_controller config (and having to watch it). This will
	// have to change if we switch to clustered version of NATS.
	uris := []string{}
	for _, ipaddr := range ipaddrs {
		uris = append(uris, fmt.Sprintf("nats://%s:4222/", ipaddr))
	}

	return uris, nil
}

func getNodesWithNatsRunning() []string {
	nodes := []string{}
	for ipaddr, info := range *GetNodeConfig() {
		for role, _ := range info.Roles {
			if role == "nats" || role == "primary" {
				nodes = append(nodes, ipaddr)
				break
			}
		}
	}
	return nodes
}
<|endoftext|>


<commit_before>package handlerutils

import (
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"unicode/utf8"

	"github.com/prasannavl/mchain/hconv"

	"github.com/prasannavl/mchain"
)

var parseUrl = url.Parse

// Redirect to a fixed URL
type redirectHandler struct {
	url  string
	code int
}

func (rh *redirectHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	Redirect(w, r, rh.url, rh.code)
}

// RedirectHandler returns a request handler that redirects
// each request it receives to the given url using the given
// status code.
//
// The provided code should be in the 3xx range and is usually
// StatusMovedPermanently, StatusFound or StatusSeeOther.
func HttpRedirectHandler(url string, code int) http.Handler {
	return &redirectHandler{url, code}
}

func RedirectHandler(url string, code int) mchain.Handler {
	// Use recoverPanic as false, since no panic is raised in
	// this module, and if a panic occurs there's something
	// seriously wrong, and it's okay to bubble it.
	return hconv.FromHttp(&redirectHandler{url, code}, false)
}

// Redirect replies to the request with a redirect to url,
// which may be a path relative to the request path.
//
// The provided code should be in the 3xx range and is usually
// StatusMovedPermanently, StatusFound or StatusSeeOther.
func Redirect(w http.ResponseWriter, r *http.Request, url string, code int) {
	// parseURL is just url.Parse (url is shadowed for godoc).
	if u, err := parseUrl(url); err == nil {
		// If url was relative, make absolute by
		// combining with request path.
		// The browser would probably do this for us,
		// but doing it ourselves is more reliable.

		// NOTE(rsc): RFC 2616 says that the Location
		// line must be an absolute URI, like
		// "http://www.google.com/redirect/",
		// not a path like "/redirect/".
		// Unfortunately, we don't know what to
		// put in the host name section to get the
		// client to connect to us again, so we can't
		// know the right absolute URI to send back.
		// Because of this problem, no one pays attention
		// to the RFC; they all send back just a new path.
		// So do we.
		if u.Scheme == "" && u.Host == "" {
			oldpath := r.URL.Path
			if oldpath == "" { // should not happen, but avoid a crash if it does
				oldpath = "/"
			}

			// no leading http://server
			if url == "" || url[0] != '/' {
				// make relative path absolute
				olddir, _ := path.Split(oldpath)
				url = olddir + url
			}

			var query string
			if i := strings.Index(url, "?"); i != -1 {
				url, query = url[:i], url[i:]
			}

			// clean up but preserve trailing slash
			trailing := strings.HasSuffix(url, "/")
			url = path.Clean(url)
			if trailing && !strings.HasSuffix(url, "/") {
				url += "/"
			}
			url += query
		}
	}

	w.Header().Set("Location", hexEscapeNonASCII(url))
	w.WriteHeader(code)
}

func hexEscapeNonASCII(s string) string {
	newLen := 0
	for i := 0; i < len(s); i++ {
		if s[i] >= utf8.RuneSelf {
			newLen += 3
		} else {
			newLen++
		}
	}
	if newLen == len(s) {
		return s
	}
	b := make([]byte, 0, newLen)
	for i := 0; i < len(s); i++ {
		if s[i] >= utf8.RuneSelf {
			b = append(b, '%')
			b = strconv.AppendInt(b, int64(s[i]), 16)
		} else {
			b = append(b, s[i])
		}
	}
	return string(b)
}

// UnsafeRedirect  does not convert relative paths to absolute paths, or clean paths
// like Redirect does.
func UnsafeRedirect(w http.ResponseWriter, r *http.Request, newPath string, redirectStatus int) {
	w.Header().Set("Location", PathWithOptionalURLQuery(newPath, r.URL.RawQuery))
	w.WriteHeader(redirectStatus)
}

func PathWithOptionalURLQuery(newPath string, rawQuery string) string {
	if strings.ContainsRune(newPath, '?') {
		return  newPath
	}
	if rawQuery != "" {
		newPath += "?" + rawQuery
	}
	return newPath
}
<commit_msg>add: unsafe redirect handlers<commit_after>package handlerutils

import (
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"unicode/utf8"

	"github.com/prasannavl/mchain/hconv"

	"github.com/prasannavl/mchain"
)

var parseUrl = url.Parse

// Redirect to a fixed URL
type redirectHandler struct {
	url  string
	code int
}

func (rh *redirectHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	Redirect(w, r, rh.url, rh.code)
}

// HttpRedirectHandler returns a request handler that redirects
// each request it receives to the given url using the given
// status code.
//
// The provided code should be in the 3xx range and is usually
// StatusMovedPermanently, StatusFound or StatusSeeOther.
func HttpRedirectHandler(url string, code int) http.Handler {
	return &redirectHandler{url, code}
}

func RedirectHandler(url string, code int) mchain.Handler {
	// Use recoverPanic as false, since no panic is raised in
	// this module, and if a panic occurs there's something
	// seriously wrong, and it's okay to bubble it.
	return hconv.FromHttp(&redirectHandler{url, code}, false)
}

// Redirect replies to the request with a redirect to url,
// which may be a path relative to the request path.
//
// The provided code should be in the 3xx range and is usually
// StatusMovedPermanently, StatusFound or StatusSeeOther.
func Redirect(w http.ResponseWriter, r *http.Request, url string, code int) {
	// parseURL is just url.Parse (url is shadowed for godoc).
	if u, err := parseUrl(url); err == nil {
		// If url was relative, make absolute by
		// combining with request path.
		// The browser would probably do this for us,
		// but doing it ourselves is more reliable.

		// NOTE(rsc): RFC 2616 says that the Location
		// line must be an absolute URI, like
		// "http://www.google.com/redirect/",
		// not a path like "/redirect/".
		// Unfortunately, we don't know what to
		// put in the host name section to get the
		// client to connect to us again, so we can't
		// know the right absolute URI to send back.
		// Because of this problem, no one pays attention
		// to the RFC; they all send back just a new path.
		// So do we.
		if u.Scheme == "" && u.Host == "" {
			oldpath := r.URL.Path
			if oldpath == "" { // should not happen, but avoid a crash if it does
				oldpath = "/"
			}

			// no leading http://server
			if url == "" || url[0] != '/' {
				// make relative path absolute
				olddir, _ := path.Split(oldpath)
				url = olddir + url
			}

			var query string
			if i := strings.Index(url, "?"); i != -1 {
				url, query = url[:i], url[i:]
			}

			// clean up but preserve trailing slash
			trailing := strings.HasSuffix(url, "/")
			url = path.Clean(url)
			if trailing && !strings.HasSuffix(url, "/") {
				url += "/"
			}
			url += query
		}
	}

	w.Header().Set("Location", hexEscapeNonASCII(url))
	w.WriteHeader(code)
}

func hexEscapeNonASCII(s string) string {
	newLen := 0
	for i := 0; i < len(s); i++ {
		if s[i] >= utf8.RuneSelf {
			newLen += 3
		} else {
			newLen++
		}
	}
	if newLen == len(s) {
		return s
	}
	b := make([]byte, 0, newLen)
	for i := 0; i < len(s); i++ {
		if s[i] >= utf8.RuneSelf {
			b = append(b, '%')
			b = strconv.AppendInt(b, int64(s[i]), 16)
		} else {
			b = append(b, s[i])
		}
	}
	return string(b)
}

// UnsafeRedirect  does not convert relative paths to absolute paths, or clean paths
// like Redirect does.
func UnsafeRedirect(w http.ResponseWriter, r *http.Request, location string, code int) {
	w.Header().Set("Location", PathWithOptionalURLQuery(location, r.URL.RawQuery))
	w.WriteHeader(code)
}

// Redirect to a fixed URL
type unsafeRedirectHandler struct {
	location  string
	code int
}

func (rh *unsafeRedirectHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	UnsafeRedirect(w, r, rh.location, rh.code)
}

func HttpUnsafeRedirectHandler(location string, code int) http.Handler {
	return &unsafeRedirectHandler{location, code}
}

func UnsafeRedirectHandler(location string, code int) mchain.Handler {
	return hconv.FromHttp(&unsafeRedirectHandler{location, code}, false)
}

func PathWithOptionalURLQuery(path string, rawQuery string) string {
	if strings.ContainsRune(path, '?') {
		return  path
	}
	if rawQuery != "" {
		path += "?" + rawQuery
	}
	return path
}
<|endoftext|>


<commit_before>package signal

import (
	"fmt"
	"os"
	gosignal "os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// Trap sets up a simplified signal "trap", appropriate for common
// behavior expected from a vanilla unix command-line tool in general
// (and the Docker engine in particular).
//
// * If SIGINT or SIGTERM are received, `cleanup` is called, then the process is terminated.
// * If SIGINT or SIGTERM are received 3 times before cleanup is complete, then cleanup is
//   skipped and the process is terminated immediately (allows force quit of stuck daemon)
// * A SIGQUIT always causes an exit without cleanup, with a goroutine dump preceding exit.
// * Ignore SIGPIPE events. These are generated by systemd when journald is restarted while
//   the docker daemon is not restarted and also running under systemd.
//   Fixes https://github.com/docker/docker/issues/19728
//
func Trap(cleanup func()) {
	c := make(chan os.Signal, 1)
	// we will handle INT, TERM, QUIT, SIGPIPE here
	signals := []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGPIPE}
	gosignal.Notify(c, signals...)
	go func() {
		interruptCount := uint32(0)
		for sig := range c {
			if sig == syscall.SIGPIPE {
				continue
			}

			go func(sig os.Signal) {
				logrus.Infof("Processing signal '%v'", sig)
				switch sig {
				case os.Interrupt, syscall.SIGTERM:
					if atomic.LoadUint32(&interruptCount) < 3 {
						// Initiate the cleanup only once
						if atomic.AddUint32(&interruptCount, 1) == 1 {
							// Call the provided cleanup handler
							cleanup()
							os.Exit(0)
						} else {
							return
						}
					} else {
						// 3 SIGTERM/INT signals received; force exit without cleanup
						logrus.Info("Forcing docker daemon shutdown without cleanup; 3 interrupts received")
					}
				case syscall.SIGQUIT:
					DumpStacks("")
					logrus.Info("Forcing docker daemon shutdown without cleanup on SIGQUIT")
				}
				//for the SIGINT/TERM, and SIGQUIT non-clean shutdown case, exit with 128 + signal #
				os.Exit(128 + int(sig.(syscall.Signal)))
			}(sig)
		}
	}()
}

const stacksLogNameTemplate = "goroutine-stacks-%s.log"

// DumpStacks appends the runtime stack into file in dir and returns full path
// to that file.
func DumpStacks(dir string) (string, error) {
	var (
		buf       []byte
		stackSize int
	)
	bufferLen := 16384
	for stackSize == len(buf) {
		buf = make([]byte, bufferLen)
		stackSize = runtime.Stack(buf, true)
		bufferLen *= 2
	}
	buf = buf[:stackSize]
	var f *os.File
	if dir != "" {
		path := filepath.Join(dir, fmt.Sprintf(stacksLogNameTemplate, strings.Replace(time.Now().Format(time.RFC3339), ":", "", -1)))
		var err error
		f, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0666)
		if err != nil {
			return "", errors.Wrap(err, "failed to open file to write the goroutine stacks")
		}
		defer f.Close()
		defer f.Sync()
	} else {
		f = os.Stderr
	}
	if _, err := f.Write(buf); err != nil {
		return "", errors.Wrap(err, "failed to write goroutine stacks")
	}
	return f.Name(), nil
}
<commit_msg>Remove the logrus from pkg/signal<commit_after>package signal

import (
	"fmt"
	"os"
	gosignal "os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/pkg/errors"
)

// Trap sets up a simplified signal "trap", appropriate for common
// behavior expected from a vanilla unix command-line tool in general
// (and the Docker engine in particular).
//
// * If SIGINT or SIGTERM are received, `cleanup` is called, then the process is terminated.
// * If SIGINT or SIGTERM are received 3 times before cleanup is complete, then cleanup is
//   skipped and the process is terminated immediately (allows force quit of stuck daemon)
// * A SIGQUIT always causes an exit without cleanup, with a goroutine dump preceding exit.
// * Ignore SIGPIPE events. These are generated by systemd when journald is restarted while
//   the docker daemon is not restarted and also running under systemd.
//   Fixes https://github.com/docker/docker/issues/19728
//
func Trap(cleanup func(), logger interface {
	Info(args ...interface{})
}) {
	c := make(chan os.Signal, 1)
	// we will handle INT, TERM, QUIT, SIGPIPE here
	signals := []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGPIPE}
	gosignal.Notify(c, signals...)
	go func() {
		interruptCount := uint32(0)
		for sig := range c {
			if sig == syscall.SIGPIPE {
				continue
			}

			go func(sig os.Signal) {
				logger.Info(fmt.Sprintf("Processing signal '%v'", sig))
				switch sig {
				case os.Interrupt, syscall.SIGTERM:
					if atomic.LoadUint32(&interruptCount) < 3 {
						// Initiate the cleanup only once
						if atomic.AddUint32(&interruptCount, 1) == 1 {
							// Call the provided cleanup handler
							cleanup()
							os.Exit(0)
						} else {
							return
						}
					} else {
						// 3 SIGTERM/INT signals received; force exit without cleanup
						logger.Info("Forcing docker daemon shutdown without cleanup; 3 interrupts received")
					}
				case syscall.SIGQUIT:
					DumpStacks("")
					logger.Info("Forcing docker daemon shutdown without cleanup on SIGQUIT")
				}
				//for the SIGINT/TERM, and SIGQUIT non-clean shutdown case, exit with 128 + signal #
				os.Exit(128 + int(sig.(syscall.Signal)))
			}(sig)
		}
	}()
}

const stacksLogNameTemplate = "goroutine-stacks-%s.log"

// DumpStacks appends the runtime stack into file in dir and returns full path
// to that file.
func DumpStacks(dir string) (string, error) {
	var (
		buf       []byte
		stackSize int
	)
	bufferLen := 16384
	for stackSize == len(buf) {
		buf = make([]byte, bufferLen)
		stackSize = runtime.Stack(buf, true)
		bufferLen *= 2
	}
	buf = buf[:stackSize]
	var f *os.File
	if dir != "" {
		path := filepath.Join(dir, fmt.Sprintf(stacksLogNameTemplate, strings.Replace(time.Now().Format(time.RFC3339), ":", "", -1)))
		var err error
		f, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0666)
		if err != nil {
			return "", errors.Wrap(err, "failed to open file to write the goroutine stacks")
		}
		defer f.Close()
		defer f.Sync()
	} else {
		f = os.Stderr
	}
	if _, err := f.Write(buf); err != nil {
		return "", errors.Wrap(err, "failed to write goroutine stacks")
	}
	return f.Name(), nil
}
<|endoftext|>


<commit_before>package internal

import (
	"io"
	// "log"
	"math"
	"strings"
	"time"

	"github.com/pkg/errors"
)

type node struct {
	lastupdate time.Time
	alive      bool
}

func (n *node) isQuarantineExpired(quarantine float64) bool {
	return math.Abs(time.Now().Sub(n.lastupdate).Minutes()) > quarantine
	//	d := time.Now().Sub(n.lastupdate).Minutes()
	//	log.Println("Mins from last =", d, " quarantine =", quarantine)
	//	return d > quarantine
}

type pool struct {
	data map[string]node
}

// GlobalPool is global pool
var GlobalPool = pool{data: make(map[string]node)}

func (p *pool) UpdateStatus(server string, alive bool) {
	p.data[server] = node{lastupdate: time.Now(), alive: alive}
	//	log.Println("Status of server", server, "set alive =", alive)
}

func (p *pool) GetNextServer(server string, servers []string, quarantine float64) (string, error) {
	if len(servers) == 0 {
		return "", errors.New("Your server list is empty")
	}

	start := -1
	size := len(servers)
	count := size
	if len(server) > 0 {
		for i, s := range servers {
			if server == s {
				start = i + 1
				break
			}
		}
		if start == -1 {
			return "", errors.New(strings.Join([]string{"Server not found in your server list: ", server}, ""))
		}
		count--
	} else {
		start = 0
	}

	// try to find node with no status, alive nodes or where quarantine is expired:
	for i := 0; i < count; i++ {
		index := (i + start) % size
		s := servers[index]
		node, found := p.data[s]
		if !found {
			return s, nil
		}
		if node.alive {
			return s, nil
		}
		if node.isQuarantineExpired(quarantine) {
			return s, nil
		}
	}
	return "", io.EOF
}
<commit_msg>extract node from if there are no alive nodes only<commit_after>package internal

import (
	"io"
	// "log"
	"math"
	"strings"
	"time"

	"github.com/pkg/errors"
)

type node struct {
	lastupdate time.Time
	alive      bool
}

func (n *node) isQuarantineExpired(quarantine float64) bool {
	return math.Abs(time.Now().Sub(n.lastupdate).Minutes()) > quarantine
	//	d := time.Now().Sub(n.lastupdate).Minutes()
	//	log.Println("Mins from last =", d, " quarantine =", quarantine)
	//	return d > quarantine
}

type pool struct {
	data map[string]node
}

// GlobalPool is global pool
var GlobalPool = pool{data: make(map[string]node)}

func (p *pool) UpdateStatus(server string, alive bool) {
	p.data[server] = node{lastupdate: time.Now(), alive: alive}
	//	log.Println("Status of server", server, "set alive =", alive)
}

func (p *pool) GetNextServer(server string, servers []string, quarantine float64) (string, error) {
	if len(servers) == 0 {
		return "", errors.New("Your server list is empty")
	}

	start := -1
	size := len(servers)
	count := size
	if len(server) > 0 {
		for i, s := range servers {
			if server == s {
				start = i + 1
				break
			}
		}
		if start == -1 {
			return "", errors.New(strings.Join([]string{"Server not found in your server list: ", server}, ""))
		}
		count--
	} else {
		start = 0
	}

	// try to find node with no status, alive nodes:
	for i := 0; i < count; i++ {
		index := (i + start) % size
		s := servers[index]
		node, found := p.data[s]
		if !found {
			return s, nil
		}
		if node.alive {
			return s, nil
		}
	}

	// try to find node where quarantine is expired:
	for i := 0; i < count; i++ {
		index := (i + start) % size
		s := servers[index]
		node, found := p.data[s]
		if !found {
			return s, nil
		}
		if node.alive {
			return s, nil
		}
		if node.isQuarantineExpired(quarantine) {
			return s, nil
		}
	}
	return "", io.EOF
}
<|endoftext|>


<commit_before>package idx

import (
	"fmt"
	"reflect"
	"testing"
	"unsafe"

	goi "github.com/robert-milan/go-object-interning"
	. "github.com/smartystreets/goconvey/convey"
)

var defs []*MetricDefinition
var tags []string
var metName = "metric.names.can.be.a.bit.longer.than.normal.sometimes"
var metNameRepeating = "metric.metric.metric.metric.metric.metric.metric.metric"
var testSz string

func genTags(num int) []string {
	szs := make([]string, num)
	for i := 0; i < len(szs); i++ {
		szs[i] = fmt.Sprintf("key%d=val%d", i, i)
	}
	return szs
}

func genMetricDefinitionsWithSameName(num int, defs []*MetricDefinition) {
	for i := 0; i < len(defs); i++ {
		defs[i] = &MetricDefinition{
			OrgId:      1,
			Interval:   10,
			LastUpdate: 0,
			Partition:  1,
		}
		defs[i].SetTags(genTags(5))
		defs[i].SetMType("rate")
		defs[i].SetMetricName("anotheryetlonger.short.metric.name")
		defs[i].SetUnit("test")
		defs[i].SetId()
	}
}

func genMetricDefinitionsWithoutTags(num int, defs []*MetricDefinition) {
	for i := 0; i < len(defs); i++ {
		defs[i] = &MetricDefinition{
			OrgId:      1,
			Interval:   10,
			LastUpdate: 0,
			Partition:  1,
		}
		defs[i].SetMType("rate")
		defs[i].SetMetricName(fmt.Sprintf("metric%d", i))
		defs[i].SetUnit("test")
		defs[i].SetId()
	}
}

func TestCreateDeleteMetricDefinition10(t *testing.T) {
	testCreateDeleteMetricDefinition(t, 10)
}

func TestCreateDeleteMetricDefinition1000(t *testing.T) {
	testCreateDeleteMetricDefinition(t, 1000)
}

func testCreateDeleteMetricDefinition(t *testing.T, num int) {
	IdxIntern = goi.NewObjectIntern(nil)
	defs = make([]*MetricDefinition, num)
	name := "anotheryetlonger.short.metric.name"

	genMetricDefinitionsWithSameName(num, defs)

	originalNameAddress := defs[0].Name.Nodes()[0]

	Convey("When creating MetricDefinitions", t, func() {
		Convey(fmt.Sprintf("number of definitions should be %d", num), func() {
			So(defs, ShouldHaveLength, num)
		})
		Convey(fmt.Sprintf("reference counts should be at %d", num), func() {
			for _, md := range defs {
				for _, ptr := range md.Name.Nodes() {
					cnt, err := IdxIntern.RefCnt(ptr)
					So(err, ShouldBeNil)
					So(cnt, ShouldEqual, num)
				}

			}
		})
		Convey(fmt.Sprintf("After deleting half of the metricdefinitions reference count should be %d", num/2), func() {
			for i := 0; i < num/2; i++ {
				InternReleaseMetricDefinition(*defs[i])
				defs[i] = nil
			}
			defs = defs[num/2:]
			for _, md := range defs {
				for _, ptr := range md.Name.Nodes() {
					cnt, err := IdxIntern.RefCnt(ptr)
					So(err, ShouldBeNil)
					So(cnt, ShouldEqual, num/2)
				}

			}
		})
		Convey("Name should still be the same when retrieved", func() {
			for _, md := range defs {
				So(md.Name.String(), ShouldEqual, name)
			}
		})
		Convey("After deleting all of the metricdefinitions the name should be removed from the object store", func() {
			nameAddr := defs[0].Name.Nodes()[0]
			for i := 0; i < len(defs); i++ {
				InternReleaseMetricDefinition(*defs[i])
				defs[i] = nil
			}
			cnt, err := IdxIntern.RefCnt(nameAddr)
			So(err, ShouldNotBeNil)
			So(cnt, ShouldEqual, 0)
		})
		Convey("After adding more metricdefinitions with the same name as before we should have a new object address for their names", func() {
			defs = nil
			defs = make([]*MetricDefinition, num)
			// create this to use the first memory offset of the new slab in a fresh slabPool in case
			// MMap decides to use the same memory chunk. The string is the same length as what should be in slot 0.
			IdxIntern.AddOrGetSzNoCprsn([]byte("bopuifszfumpohfs"))
			genMetricDefinitionsWithSameName(num, defs)
			So(originalNameAddress, ShouldNotEqual, defs[0].Name.Nodes()[0])
		})
	})
}

func TestMetricNameAndTagAddresses(t *testing.T) {
	IdxIntern = goi.NewObjectIntern(nil)
	defs = make([]*MetricDefinition, 5)
	for i := 0; i < len(defs); i++ {
		defs[i] = &MetricDefinition{
			OrgId:      uint32(i),
			Interval:   10,
			LastUpdate: 0,
			Partition:  1,
		}
		defs[i].SetTags(genTags(5))
		defs[i].SetMType("rate")
		defs[i].SetMetricName("some.short.metric.name")
		defs[i].SetUnit("test")
		defs[i].SetId()
	}

	Convey("When creating MetricDefinitions with the same name and tags", t, func() {
		Convey("names should be using the same object addresses", func() {
			for _, md := range defs {
				for _, mdd := range defs {
					So(md.Name.Nodes(), ShouldResemble, mdd.Name.Nodes())
				}
			}
		})
		Convey("reference counts should be at 5", func() {
			for _, md := range defs {
				for _, ptr := range md.Name.Nodes() {
					cnt, err := IdxIntern.RefCnt(ptr)
					So(err, ShouldBeNil)
					So(cnt, ShouldEqual, 5)
				}

			}
		})
		Convey("tags should be using the same object addresses", func() {
			for idx, tag := range defs[0].Tags {
				keyData := (*reflect.StringHeader)(unsafe.Pointer(&tag.Key)).Data
				valData := (*reflect.StringHeader)(unsafe.Pointer(&tag.Value)).Data
				for i := 1; i < len(defs); i++ {
					So(keyData, ShouldEqual, (*reflect.StringHeader)(unsafe.Pointer(&defs[i].Tags[idx].Key)).Data)
					So(valData, ShouldEqual, (*reflect.StringHeader)(unsafe.Pointer(&defs[i].Tags[idx].Value)).Data)
				}
			}
		})
	})
}

func TestTagKeyValuesAndNameWithTags(t *testing.T) {
	IdxIntern = goi.NewObjectIntern(nil)
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithSameName(1, defs)
	tags = genTags(5)
	defs[0].SetTags(tags)

	Convey("After adding tags to a MetricDefinition", t, func() {
		Convey("the Strings() function of TagKeyValues should resemble the original []string", func() {
			So(tags, ShouldResemble, defs[0].Tags.Strings())
		})
		Convey("and NameWithTags should be predictable", func() {
			nwt := defs[0].NameWithTags()
			expected := "anotheryetlonger.short.metric.name;key0=val0;key1=val1;key2=val2;key3=val3;key4=val4"
			So(nwt, ShouldEqual, expected)
		})
	})
}

func BenchmarkSetTags10(b *testing.B) {
	benchmarkSetTags(b, 10)
}

func BenchmarkSetTags100(b *testing.B) {
	benchmarkSetTags(b, 100)
}

func BenchmarkSetTags1000(b *testing.B) {
	benchmarkSetTags(b, 1000)
}

func benchmarkSetTags(b *testing.B, num int) {
	tags = genTags(num)
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithoutTags(1, defs)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		defs[0].SetTags(tags)
	}
}

func BenchmarkSetMetricName(b *testing.B) {
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithoutTags(1, defs)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		defs[0].SetMetricName(metName)
	}
}

func BenchmarkSetMetricNameRepeatingWords(b *testing.B) {
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithoutTags(1, defs)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		defs[0].SetMetricName(metNameRepeating)
	}
}

func BenchmarkGetMetricName(b *testing.B) {
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithoutTags(1, defs)
	defs[0].SetMetricName(metName)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		testSz = defs[0].Name.String()
	}
}

func BenchmarkGetNameWithTags(b *testing.B) {
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithSameName(1, defs)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		testSz = defs[0].NameWithTags()
	}
}

func BenchmarkGetTags(b *testing.B) {
	defs = make([]*MetricDefinition, 1)
	genMetricDefinitionsWithSameName(1, defs)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tags = defs[0].Tags.Strings()
	}
}
<commit_msg>simplify getting of tags and []*MetricDefinition<commit_after>package idx

import (
	"fmt"
	"reflect"
	"testing"
	"unsafe"

	goi "github.com/robert-milan/go-object-interning"
	. "github.com/smartystreets/goconvey/convey"
)

var metName = "metric.names.can.be.a.bit.longer.than.normal.sometimes"
var metNameRepeating = "metric.metric.metric.metric.metric.metric.metric.metric"
var testSz string

func genTags(num int) []string {
	szs := make([]string, num)
	for i := 0; i < num; i++ {
		szs[i] = fmt.Sprintf("key%d=val%d", i, i)
	}
	return szs
}

func genMetricDefinitionsWithSameName(num int) []*MetricDefinition {
	defs := make([]*MetricDefinition, num)
	for i := 0; i < num; i++ {
		defs[i] = &MetricDefinition{
			OrgId:      1,
			Interval:   10,
			LastUpdate: 0,
			Partition:  1,
		}
		defs[i].SetTags(genTags(5))
		defs[i].SetMType("rate")
		defs[i].SetMetricName("anotheryetlonger.short.metric.name")
		defs[i].SetUnit("test")
		defs[i].SetId()
	}
	return defs
}

func genMetricDefinitionsWithoutTags(num int) []*MetricDefinition {
	defs := make([]*MetricDefinition, num)
	for i := 0; i < num; i++ {
		defs[i] = &MetricDefinition{
			OrgId:      1,
			Interval:   10,
			LastUpdate: 0,
			Partition:  1,
		}
		defs[i].SetMType("rate")
		defs[i].SetMetricName(fmt.Sprintf("metric%d", i))
		defs[i].SetUnit("test")
		defs[i].SetId()
	}
	return defs
}

func TestCreateDeleteMetricDefinition10(t *testing.T) {
	testCreateDeleteMetricDefinition(t, 10)
}

func TestCreateDeleteMetricDefinition1000(t *testing.T) {
	testCreateDeleteMetricDefinition(t, 1000)
}

func testCreateDeleteMetricDefinition(t *testing.T, num int) {
	IdxIntern = goi.NewObjectIntern(nil)
	defs := genMetricDefinitionsWithSameName(num)
	name := "anotheryetlonger.short.metric.name"

	originalNameAddress := defs[0].Name.Nodes()[0]

	Convey("When creating MetricDefinitions", t, func() {
		Convey(fmt.Sprintf("number of definitions should be %d", num), func() {
			So(defs, ShouldHaveLength, num)
		})
		Convey(fmt.Sprintf("reference counts should be at %d", num), func() {
			for _, md := range defs {
				for _, ptr := range md.Name.Nodes() {
					cnt, err := IdxIntern.RefCnt(ptr)
					So(err, ShouldBeNil)
					So(cnt, ShouldEqual, num)
				}

			}
		})
		Convey(fmt.Sprintf("After deleting half of the metricdefinitions reference count should be %d", num/2), func() {
			for i := 0; i < num/2; i++ {
				InternReleaseMetricDefinition(*defs[i])
				defs[i] = nil
			}
			defs = defs[num/2:]
			for _, md := range defs {
				for _, ptr := range md.Name.Nodes() {
					cnt, err := IdxIntern.RefCnt(ptr)
					So(err, ShouldBeNil)
					So(cnt, ShouldEqual, num/2)
				}

			}
		})
		Convey("Name should still be the same when retrieved", func() {
			for _, md := range defs {
				So(md.Name.String(), ShouldEqual, name)
			}
		})
		Convey("After deleting all of the metricdefinitions the name should be removed from the object store", func() {
			nameAddr := defs[0].Name.Nodes()[0]
			for i := 0; i < len(defs); i++ {
				InternReleaseMetricDefinition(*defs[i])
				defs[i] = nil
			}
			cnt, err := IdxIntern.RefCnt(nameAddr)
			So(err, ShouldNotBeNil)
			So(cnt, ShouldEqual, 0)
		})
		Convey("After adding more metricdefinitions with the same name as before we should have a new object address for their names", func() {
			// create this to use the first memory offset of the new slab in a fresh slabPool in case
			// MMap decides to use the same memory chunk. The string is the same length as what should be in slot 0.
			IdxIntern.AddOrGetSzNoCprsn([]byte("bopuifszfumpohfs"))
			defs := genMetricDefinitionsWithSameName(num)
			So(originalNameAddress, ShouldNotEqual, defs[0].Name.Nodes()[0])
		})
	})
}

func TestMetricNameAndTagAddresses(t *testing.T) {
	IdxIntern = goi.NewObjectIntern(nil)
	defs := make([]*MetricDefinition, 5)
	for i := 0; i < len(defs); i++ {
		defs[i] = &MetricDefinition{
			OrgId:      uint32(i),
			Interval:   10,
			LastUpdate: 0,
			Partition:  1,
		}
		defs[i].SetTags(genTags(5))
		defs[i].SetMType("rate")
		defs[i].SetMetricName("some.short.metric.name")
		defs[i].SetUnit("test")
		defs[i].SetId()
	}

	Convey("When creating MetricDefinitions with the same name and tags", t, func() {
		Convey("names should be using the same object addresses", func() {
			for _, md := range defs {
				for _, mdd := range defs {
					So(md.Name.Nodes(), ShouldResemble, mdd.Name.Nodes())
				}
			}
		})
		Convey("reference counts should be at 5", func() {
			for _, md := range defs {
				for _, ptr := range md.Name.Nodes() {
					cnt, err := IdxIntern.RefCnt(ptr)
					So(err, ShouldBeNil)
					So(cnt, ShouldEqual, 5)
				}

			}
		})
		Convey("tags should be using the same object addresses", func() {
			for idx, tag := range defs[0].Tags {
				keyData := (*reflect.StringHeader)(unsafe.Pointer(&tag.Key)).Data
				valData := (*reflect.StringHeader)(unsafe.Pointer(&tag.Value)).Data
				for i := 1; i < len(defs); i++ {
					So(keyData, ShouldEqual, (*reflect.StringHeader)(unsafe.Pointer(&defs[i].Tags[idx].Key)).Data)
					So(valData, ShouldEqual, (*reflect.StringHeader)(unsafe.Pointer(&defs[i].Tags[idx].Value)).Data)
				}
			}
		})
	})
}

func TestTagKeyValuesAndNameWithTags(t *testing.T) {
	IdxIntern = goi.NewObjectIntern(nil)
	defs := genMetricDefinitionsWithSameName(1)
	tags := genTags(5)
	defs[0].SetTags(tags)

	Convey("After adding tags to a MetricDefinition", t, func() {
		Convey("the Strings() function of TagKeyValues should resemble the original []string", func() {
			So(tags, ShouldResemble, defs[0].Tags.Strings())
		})
		Convey("and NameWithTags should be predictable", func() {
			nwt := defs[0].NameWithTags()
			expected := "anotheryetlonger.short.metric.name;key0=val0;key1=val1;key2=val2;key3=val3;key4=val4"
			So(nwt, ShouldEqual, expected)
		})
	})
}

func BenchmarkSetTags10(b *testing.B) {
	benchmarkSetTags(b, 10)
}

func BenchmarkSetTags100(b *testing.B) {
	benchmarkSetTags(b, 100)
}

func BenchmarkSetTags1000(b *testing.B) {
	benchmarkSetTags(b, 1000)
}

func benchmarkSetTags(b *testing.B, num int) {
	tags := genTags(num)
	defs := genMetricDefinitionsWithoutTags(1)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		defs[0].SetTags(tags)
	}
}

func BenchmarkSetMetricName(b *testing.B) {
	defs := genMetricDefinitionsWithoutTags(1)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		defs[0].SetMetricName(metName)
	}
}

func BenchmarkSetMetricNameRepeatingWords(b *testing.B) {
	defs := genMetricDefinitionsWithoutTags(1)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		defs[0].SetMetricName(metNameRepeating)
	}
}

func BenchmarkGetMetricName(b *testing.B) {
	defs := genMetricDefinitionsWithoutTags(1)
	defs[0].SetMetricName(metName)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		testSz = defs[0].Name.String()
	}
}

func BenchmarkGetNameWithTags(b *testing.B) {
	defs := genMetricDefinitionsWithSameName(1)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		testSz = defs[0].NameWithTags()
	}
}

func BenchmarkGetTags(b *testing.B) {
	defs := genMetricDefinitionsWithSameName(1)
	var tags []string

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tags = defs[0].Tags.Strings()
	}
	if len(tags) != 5 {
		panic("incorrect number of tags returned")
	}
}
<|endoftext|>