<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location        string
	ChecksUrl       string
	MeasurementsUrl string
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func MeasureLoop(config Config, checks chan Check, measurements chan Measurement) {
	for {
		c := <-checks
		m := c.Measure(config)

		measurements <- m
	}
}

func RecordLoop(config Config, measurements chan Measurement) {
	payload := make([]Measurement, 0, 100)
	for {
		m := <-measurements
		payload = append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			panic(err)
		}
		resp.Body.Close()
		payload = make([]Measurement, 0, 100)

		fmt.Println(resp)
	}
}

func GetChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")

	fmt.Printf("%s\n", config.MeasurementsUrl)

	check_list := GetChecks(config)

	checks := make(chan Check)
	measurements := make(chan Measurement)

	go MeasureLoop(config, checks, measurements)
	go RecordLoop(config, measurements)

	for {
		for _, c := range check_list {
			checks <- c
		}

		time.Sleep(1000 * time.Millisecond)
	}
}
<commit_msg>record millis rather than secs<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location        string
	ChecksUrl       string
	MeasurementsUrl string
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	nanos := now.UnixNano()
	millis := nanos / 1000000

	m.T = int(millis)

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func MeasureLoop(config Config, checks chan Check, measurements chan Measurement) {
	for {
		c := <-checks
		m := c.Measure(config)

		measurements <- m
	}
}

func RecordLoop(config Config, measurements chan Measurement) {
	payload := make([]Measurement, 0, 100)
	for {
		m := <-measurements
		payload = append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			panic(err)
		}
		resp.Body.Close()
		payload = make([]Measurement, 0, 100)

		fmt.Println(resp)
	}
}

func GetChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")

	fmt.Printf("%s\n", config.MeasurementsUrl)

	check_list := GetChecks(config)

	checks := make(chan Check)
	measurements := make(chan Measurement)

	go MeasureLoop(config, checks, measurements)
	go RecordLoop(config, measurements)

	for {
		for _, c := range check_list {
			checks <- c
		}

		time.Sleep(1000 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/nightlyone/lockfile"
	"github.com/tylerb/graceful"

	"github.com/stripe/sequins/backend"
	"github.com/stripe/sequins/sharding"
	"github.com/stripe/sequins/workqueue"
	"github.com/stripe/sequins/zk"
)

const defaultMaxLoads = 10

var errDirLocked = errors.New("failed to acquire lock")

type sequins struct {
	config  sequinsConfig
	http    *http.Server
	backend backend.Backend

	dbs     map[string]*db
	dbsLock sync.RWMutex

	address   string
	peers     *sharding.Peers
	zkWatcher *zk.Watcher

	refreshLock   sync.Mutex
	workQueue     *workqueue.WorkQueue
	refreshTicker *time.Ticker
	sighups       chan os.Signal

	stats *statsd.Client

	storeLock lockfile.Lockfile
}

func newSequins(backend backend.Backend, config sequinsConfig) *sequins {
	return &sequins{
		config:      config,
		backend:     backend,
		refreshLock: sync.Mutex{},
	}
}

func (s *sequins) init() error {
	// Start Datadog client if configured
	if s.config.Datadog.Url != "" {
		statsdClient, err := statsd.New(s.config.Datadog.Url)
		if err != nil {
			log.Fatalf("Error connecting to statsd: %s", err)
		}
		statsdClient.Namespace = "sequins."
		statsdClient.Tags = append(statsdClient.Tags, fmt.Sprintf("sequins:%s", s.config.Sharding.ClusterName))
		s.stats = statsdClient
	}

	if s.config.Sharding.Enabled {
		err := s.initCluster()
		if err != nil {
			return err
		}
	}

	// Create local directories, and load any cached versions we have.
	err := s.initLocalStore()
	if err != nil {
		return fmt.Errorf("error initializing local store: %s", err)
	}

	// This lock limits load parallelism across all dbs.
	maxLoads := s.config.MaxParallelLoads
	if maxLoads == 0 {
		maxLoads = defaultMaxLoads
	}
	s.workQueue = workqueue.NewWorkQueue(maxLoads)

	// Trigger loads before we start up.
	s.refreshAll()
	s.refreshLock.Lock()
	defer s.refreshLock.Unlock()

	// Automatically refresh, if configured to do so.
	refresh := s.config.RefreshPeriod.Duration
	if refresh != 0 {
		s.refreshTicker = time.NewTicker(refresh)
		go func() {
			log.Println("Automatically checking for new versions every", refresh.String())
			for range s.refreshTicker.C {
				s.refreshAll()
			}
		}()
	}

	// Refresh on SIGHUP.
	sighups := make(chan os.Signal)
	signal.Notify(sighups, syscall.SIGHUP)
	go func() {
		for range sighups {
			s.refreshAll()
		}
	}()

	s.sighups = sighups
	return nil
}

func (s *sequins) joinCluster(zkWatcher *zk.Watcher, routableIpAddress string) (*sharding.Peers, error) {
	lock := zkWatcher.CreateIDAssignmentLock()
	lock.Lock()
	defer lock.Unlock()

	shardID := s.config.Sharding.ShardID
	if shardID == "" {
		fp := filepath.Join(s.config.LocalStore, "self_assigned_id")
		_, err := os.Stat(fp)
		if err != nil {
			if os.IsNotExist(err) {
				log.Print("No self-assigned ID file found")

				peersNotJoined := sharding.WatchPeersNoJoin(zkWatcher)
				peersNotJoined.WaitToConverge(s.config.Sharding.TimeToConverge.Duration)

				shardID, err = peersNotJoined.SmallestAvailableShardID()
				if err != nil {
					return nil, fmt.Errorf("Error running SmallestAvailableShardID: %q", err)
				}
				log.Printf("Identified %q as smallest missing shardID", shardID)

				err := ioutil.WriteFile(fp, []byte(shardID), 0600)
				if err != nil {
					return nil, fmt.Errorf("Could not write shardID to file: %q", err)
				}
			} else {
				return nil, fmt.Errorf("Could not check %q exists: %q", fp, err)
			}
		} else {
			log.Print("Self-assigned ID file found. Attemtping to read.")
			bytes, err := ioutil.ReadFile(fp)
			if err != nil {
				return nil, fmt.Errorf("Could not read %q: %q", fp, err)
			}
			shardID = string(bytes)
			log.Printf("Extracted shardID %q from file", shardID)
		}
	}
	log.Printf("Using %q as shardID", shardID)

	peers := sharding.WatchPeers(zkWatcher, shardID, routableIpAddress)
	peers.WaitToConverge(s.config.Sharding.TimeToConverge.Duration)

	return peers, nil
}

func (s *sequins) initCluster() error {
	// This config property is calculated if not set.
	if s.config.Sharding.ProxyStageTimeout.Duration == 0 {
		stageTimeout := s.config.Sharding.ProxyTimeout.Duration / time.Duration(s.config.Sharding.Replication)
		s.config.Sharding.ProxyStageTimeout = duration{stageTimeout}
	}

	prefix := path.Join("/", s.config.Sharding.ClusterName)
	zkWatcher, err := zk.Connect(s.config.ZK.Servers, prefix,
		s.config.ZK.ConnectTimeout.Duration, s.config.ZK.SessionTimeout.Duration)
	if err != nil {
		return err
	}

	go zkWatcher.TriggerCleanup()

	hostname := s.config.Sharding.AdvertisedHostname
	if hostname == "" {
		hostname, err = os.Hostname()
		if err != nil {
			return err
		}
	}

	_, port, err := net.SplitHostPort(s.config.Bind)
	if err != nil {
		return err
	}

	ip := hostname
	ipAddresses, err := net.LookupHost(hostname)
	if err != nil {
		return err
	}
	if len(ipAddresses) == 1 {
		ip = ipAddresses[0]
	}

	routableIpAddress := net.JoinHostPort(ip, port)
	routableAddress := net.JoinHostPort(hostname, port)

	peers, err := s.joinCluster(zkWatcher, routableIpAddress)
	if err != nil {
		log.Fatalf("Error joining cluster: %q", err)
	}

	s.address = routableAddress
	s.zkWatcher = zkWatcher
	s.peers = peers
	return nil
}

func (s *sequins) initLocalStore() error {
	dataPath := filepath.Join(s.config.LocalStore, "data")
	err := os.MkdirAll(dataPath, 0755|os.ModeDir)
	if err != nil {
		return err
	}

	lock, _ := lockfile.New(filepath.Join(s.config.LocalStore, "sequins.lock"))
	s.storeLock = lock
	err = s.storeLock.TryLock()
	if err != nil {
		p, err := s.storeLock.GetOwner()
		if err == nil {
			log.Printf("The local store at %s is locked by process %d", s.config.LocalStore, p.Pid)
		}

		return errDirLocked
	}

	return nil
}

func (s *sequins) start() {
	defer s.shutdown()

	var h http.Handler = s
	if s.config.Debug.Bind != "" && s.config.Debug.Expvars {
		h = trackQueries(s)
	}

	log.Println("Listening on", s.config.Bind)
	graceful.Run(s.config.Bind, time.Second, h)
}

func (s *sequins) shutdown() {
	log.Println("Shutting down...")
	signal.Stop(s.sighups)

	if s.refreshTicker != nil {
		s.refreshTicker.Stop()
	}

	zk := s.zkWatcher
	if zk != nil {
		zk.Close()
	}

	// TODO: figure out how to cancel in-progress downloads
	// s.dbsLock.Lock()
	// defer s.dbsLock.Unlock()

	// for _, db := range s.dbs {
	// 	db.close()
	// }

	s.storeLock.Unlock()
}

func (s *sequins) listDBs() ([]string, error) {
	dbs, err := s.backend.ListDBs()
	if err != nil {
		return nil, err
	}
	return filterPaths(dbs), nil
}

func (s *sequins) refreshAll() {
	s.refreshLock.Lock()
	defer s.refreshLock.Unlock()

	dbs, err := s.listDBs()
	if err != nil {
		log.Printf("Error listing DBs from %s: %s", s.backend.DisplayPath(""), err)
		return
	}

	// Add any new DBS, and trigger refreshes for existing ones. This
	// Lock pattern is a bit goofy, since the lock is on the request
	// path and we want to hold it for as short a time as possible.
	s.dbsLock.RLock()

	newDBs := make(map[string]*db)
	var backfills sync.WaitGroup
	for _, name := range dbs {
		db := s.dbs[name]
		if db == nil {
			db = newDB(s, name)

			backfills.Add(1)
			go func() {
				db.backfillVersions()
				backfills.Done()
			}()
		} else {
			go func() {
				err := db.refresh()
				if err != nil {
					log.Printf("Error refreshing %s: %s", db.name, err)
				}
			}()
		}

		newDBs[name] = db
	}

	backfills.Wait()

	// Now, grab the full lock to switch the new map in.
	s.dbsLock.RUnlock()
	s.dbsLock.Lock()

	oldDBs := s.dbs
	s.dbs = newDBs

	// Finally, close any dbs that we're removing.
	s.dbsLock.Unlock()
	s.dbsLock.RLock()

	for name, db := range oldDBs {
		if s.dbs[name] == nil {
			log.Println("Removing and clearing database", name)
			db.close()
			db.delete()
		}
	}

	s.dbsLock.RUnlock()

	// Cleanup any zkNodes for deleted versions and dbs.
	if s.zkWatcher != nil {
		s.zkWatcher.TriggerCleanup()
	}
}

func (s *sequins) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if r.URL.Path == "/" {
		s.serveStatus(w, r)
		return
	}

	if r.URL.Path == "/healthcheck" || r.URL.Path == "/healthz" {
		s.serveHealth(w, r)
		return
	}

	var dbName, key string
	path := strings.TrimPrefix(r.URL.Path, "/")
	split := strings.Index(path, "/")
	if split == -1 {
		dbName = path
		key = ""
	} else {
		dbName = path[:split]
		key = path[split+1:]
	}

	if dbName == "" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	s.dbsLock.RLock()
	db := s.dbs[dbName]
	s.dbsLock.RUnlock()

	// If this is a proxy request, we don't want to confuse "we don't have this
	// db" with "we don't have this key"; if this is a proxied request, then the
	// peer apparently does have the db, and thinks we do too (which should never
	// happen). So we use 501 Not Implemented to indicate the former. To users,
	// we present a uniform 404.
	if db == nil {
		if r.URL.Query().Get("proxy") != "" {
			w.WriteHeader(http.StatusNotImplemented)
		} else {
			w.WriteHeader(http.StatusNotFound)
		}

		return
	}

	db.serveKey(w, r, key)
}
<commit_msg>Shard IDs: Don't wait for convergence too much<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/nightlyone/lockfile"
	"github.com/tylerb/graceful"

	"github.com/stripe/sequins/backend"
	"github.com/stripe/sequins/sharding"
	"github.com/stripe/sequins/workqueue"
	"github.com/stripe/sequins/zk"
)

const defaultMaxLoads = 10

// While we wait to pick a shard ID, other nodes are prohibited from joining the cluster because
// we're holding a lock. This timeout can therefore be pretty short.
//
// Indeed, it needs to be shorter than the regular cluster convergence timeout, or
// the cluster will be seen to "converge" when it's really just waiting for someone to release
// the lock.
const shardIDConvergenceMax = 3 * time.Second

var errDirLocked = errors.New("failed to acquire lock")

type sequins struct {
	config  sequinsConfig
	http    *http.Server
	backend backend.Backend

	dbs     map[string]*db
	dbsLock sync.RWMutex

	address   string
	peers     *sharding.Peers
	zkWatcher *zk.Watcher

	refreshLock   sync.Mutex
	workQueue     *workqueue.WorkQueue
	refreshTicker *time.Ticker
	sighups       chan os.Signal

	stats *statsd.Client

	storeLock lockfile.Lockfile
}

func newSequins(backend backend.Backend, config sequinsConfig) *sequins {
	return &sequins{
		config:      config,
		backend:     backend,
		refreshLock: sync.Mutex{},
	}
}

func (s *sequins) init() error {
	// Start Datadog client if configured
	if s.config.Datadog.Url != "" {
		statsdClient, err := statsd.New(s.config.Datadog.Url)
		if err != nil {
			log.Fatalf("Error connecting to statsd: %s", err)
		}
		statsdClient.Namespace = "sequins."
		statsdClient.Tags = append(statsdClient.Tags, fmt.Sprintf("sequins:%s", s.config.Sharding.ClusterName))
		s.stats = statsdClient
	}

	if s.config.Sharding.Enabled {
		err := s.initCluster()
		if err != nil {
			return err
		}
	}

	// Create local directories, and load any cached versions we have.
	err := s.initLocalStore()
	if err != nil {
		return fmt.Errorf("error initializing local store: %s", err)
	}

	// This lock limits load parallelism across all dbs.
	maxLoads := s.config.MaxParallelLoads
	if maxLoads == 0 {
		maxLoads = defaultMaxLoads
	}
	s.workQueue = workqueue.NewWorkQueue(maxLoads)

	// Trigger loads before we start up.
	s.refreshAll()
	s.refreshLock.Lock()
	defer s.refreshLock.Unlock()

	// Automatically refresh, if configured to do so.
	refresh := s.config.RefreshPeriod.Duration
	if refresh != 0 {
		s.refreshTicker = time.NewTicker(refresh)
		go func() {
			log.Println("Automatically checking for new versions every", refresh.String())
			for range s.refreshTicker.C {
				s.refreshAll()
			}
		}()
	}

	// Refresh on SIGHUP.
	sighups := make(chan os.Signal)
	signal.Notify(sighups, syscall.SIGHUP)
	go func() {
		for range sighups {
			s.refreshAll()
		}
	}()

	s.sighups = sighups
	return nil
}

func (s *sequins) ensureShardID(zkWatcher *zk.Watcher, routableIpAddress string) (*sharding.Peers, error) {
	lock := zkWatcher.CreateIDAssignmentLock()
	lock.Lock()
	defer lock.Unlock()

	convergenceTime := s.config.Sharding.TimeToConverge.Duration
	if convergenceTime > shardIDConvergenceMax {
		convergenceTime = shardIDConvergenceMax
	}

	shardID := s.config.Sharding.ShardID
	if shardID == "" {
		fp := filepath.Join(s.config.LocalStore, "self_assigned_id")
		_, err := os.Stat(fp)
		if err != nil {
			if os.IsNotExist(err) {
				log.Print("No self-assigned ID file found")

				peersNotJoined := sharding.WatchPeersNoJoin(zkWatcher)
				peersNotJoined.WaitToConverge(convergenceTime)

				shardID, err = peersNotJoined.SmallestAvailableShardID()
				if err != nil {
					return nil, fmt.Errorf("Error running SmallestAvailableShardID: %q", err)
				}
				log.Printf("Identified %q as smallest missing shardID", shardID)

				err := ioutil.WriteFile(fp, []byte(shardID), 0600)
				if err != nil {
					return nil, fmt.Errorf("Could not write shardID to file: %q", err)
				}
			} else {
				return nil, fmt.Errorf("Could not check %q exists: %q", fp, err)
			}
		} else {
			log.Print("Self-assigned ID file found. Attempting to read.")
			bytes, err := ioutil.ReadFile(fp)
			if err != nil {
				return nil, fmt.Errorf("Could not read %q: %q", fp, err)
			}
			shardID = string(bytes)
			log.Printf("Extracted shardID %q from file", shardID)
		}
	}
	log.Printf("Using %q as shardID", shardID)
	peers := sharding.WatchPeers(zkWatcher, shardID, routableIpAddress)
	return peers, nil
}

func (s *sequins) joinCluster(zkWatcher *zk.Watcher, routableIpAddress string) (*sharding.Peers, error) {
	// The shard-ID lock is released inside here--we don't need that lock to be held while we wait to converge.
	peers, err := s.ensureShardID(zkWatcher, routableIpAddress)
	if err != nil {
		return nil, err
	}
	peers.WaitToConverge(s.config.Sharding.TimeToConverge.Duration)
	return peers, nil
}

func (s *sequins) initCluster() error {
	// This config property is calculated if not set.
	if s.config.Sharding.ProxyStageTimeout.Duration == 0 {
		stageTimeout := s.config.Sharding.ProxyTimeout.Duration / time.Duration(s.config.Sharding.Replication)
		s.config.Sharding.ProxyStageTimeout = duration{stageTimeout}
	}

	prefix := path.Join("/", s.config.Sharding.ClusterName)
	zkWatcher, err := zk.Connect(s.config.ZK.Servers, prefix,
		s.config.ZK.ConnectTimeout.Duration, s.config.ZK.SessionTimeout.Duration)
	if err != nil {
		return err
	}

	go zkWatcher.TriggerCleanup()

	hostname := s.config.Sharding.AdvertisedHostname
	if hostname == "" {
		hostname, err = os.Hostname()
		if err != nil {
			return err
		}
	}

	_, port, err := net.SplitHostPort(s.config.Bind)
	if err != nil {
		return err
	}

	ip := hostname
	ipAddresses, err := net.LookupHost(hostname)
	if err != nil {
		return err
	}
	if len(ipAddresses) == 1 {
		ip = ipAddresses[0]
	}

	routableIpAddress := net.JoinHostPort(ip, port)
	routableAddress := net.JoinHostPort(hostname, port)

	peers, err := s.joinCluster(zkWatcher, routableIpAddress)
	if err != nil {
		log.Fatalf("Error joining cluster: %q", err)
	}

	s.address = routableAddress
	s.zkWatcher = zkWatcher
	s.peers = peers
	return nil
}

func (s *sequins) initLocalStore() error {
	dataPath := filepath.Join(s.config.LocalStore, "data")
	err := os.MkdirAll(dataPath, 0755|os.ModeDir)
	if err != nil {
		return err
	}

	lock, _ := lockfile.New(filepath.Join(s.config.LocalStore, "sequins.lock"))
	s.storeLock = lock
	err = s.storeLock.TryLock()
	if err != nil {
		p, err := s.storeLock.GetOwner()
		if err == nil {
			log.Printf("The local store at %s is locked by process %d", s.config.LocalStore, p.Pid)
		}

		return errDirLocked
	}

	return nil
}

func (s *sequins) start() {
	defer s.shutdown()

	var h http.Handler = s
	if s.config.Debug.Bind != "" && s.config.Debug.Expvars {
		h = trackQueries(s)
	}

	log.Println("Listening on", s.config.Bind)
	graceful.Run(s.config.Bind, time.Second, h)
}

func (s *sequins) shutdown() {
	log.Println("Shutting down...")
	signal.Stop(s.sighups)

	if s.refreshTicker != nil {
		s.refreshTicker.Stop()
	}

	zk := s.zkWatcher
	if zk != nil {
		zk.Close()
	}

	// TODO: figure out how to cancel in-progress downloads
	// s.dbsLock.Lock()
	// defer s.dbsLock.Unlock()

	// for _, db := range s.dbs {
	// 	db.close()
	// }

	s.storeLock.Unlock()
}

func (s *sequins) listDBs() ([]string, error) {
	dbs, err := s.backend.ListDBs()
	if err != nil {
		return nil, err
	}
	return filterPaths(dbs), nil
}

func (s *sequins) refreshAll() {
	s.refreshLock.Lock()
	defer s.refreshLock.Unlock()

	dbs, err := s.listDBs()
	if err != nil {
		log.Printf("Error listing DBs from %s: %s", s.backend.DisplayPath(""), err)
		return
	}

	// Add any new DBS, and trigger refreshes for existing ones. This
	// Lock pattern is a bit goofy, since the lock is on the request
	// path and we want to hold it for as short a time as possible.
	s.dbsLock.RLock()

	newDBs := make(map[string]*db)
	var backfills sync.WaitGroup
	for _, name := range dbs {
		db := s.dbs[name]
		if db == nil {
			db = newDB(s, name)

			backfills.Add(1)
			go func() {
				db.backfillVersions()
				backfills.Done()
			}()
		} else {
			go func() {
				err := db.refresh()
				if err != nil {
					log.Printf("Error refreshing %s: %s", db.name, err)
				}
			}()
		}

		newDBs[name] = db
	}

	backfills.Wait()

	// Now, grab the full lock to switch the new map in.
	s.dbsLock.RUnlock()
	s.dbsLock.Lock()

	oldDBs := s.dbs
	s.dbs = newDBs

	// Finally, close any dbs that we're removing.
	s.dbsLock.Unlock()
	s.dbsLock.RLock()

	for name, db := range oldDBs {
		if s.dbs[name] == nil {
			log.Println("Removing and clearing database", name)
			db.close()
			db.delete()
		}
	}

	s.dbsLock.RUnlock()

	// Cleanup any zkNodes for deleted versions and dbs.
	if s.zkWatcher != nil {
		s.zkWatcher.TriggerCleanup()
	}
}

func (s *sequins) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if r.URL.Path == "/" {
		s.serveStatus(w, r)
		return
	}

	if r.URL.Path == "/healthcheck" || r.URL.Path == "/healthz" {
		s.serveHealth(w, r)
		return
	}

	var dbName, key string
	path := strings.TrimPrefix(r.URL.Path, "/")
	split := strings.Index(path, "/")
	if split == -1 {
		dbName = path
		key = ""
	} else {
		dbName = path[:split]
		key = path[split+1:]
	}

	if dbName == "" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	s.dbsLock.RLock()
	db := s.dbs[dbName]
	s.dbsLock.RUnlock()

	// If this is a proxy request, we don't want to confuse "we don't have this
	// db" with "we don't have this key"; if this is a proxied request, then the
	// peer apparently does have the db, and thinks we do too (which should never
	// happen). So we use 501 Not Implemented to indicate the former. To users,
	// we present a uniform 404.
	if db == nil {
		if r.URL.Query().Get("proxy") != "" {
			w.WriteHeader(http.StatusNotImplemented)
		} else {
			w.WriteHeader(http.StatusNotFound)
		}

		return
	}

	db.serveKey(w, r, key)
}
<|endoftext|>


<commit_before>/**
 * Copyright 2013-2016 Seagate Technology LLC.
 *
 * This Source Code Form is subject to the terms of the Mozilla
 * Public License, v. 2.0. If a copy of the MPL was not
 * distributed with this file, You can obtain one at
 * https://mozilla.org/MP:/2.0/.
 *
 * This program is distributed in the hope that it will be useful,
 * but is provided AS-IS, WITHOUT ANY WARRANTY; including without
 * the implied warranty of MERCHANTABILITY, NON-INFRINGEMENT or
 * FITNESS FOR A PARTICULAR PURPOSE. See the Mozilla Public
 * License for more details.
 *
 * See www.openkinetic.org for more project information
 */

package kinetic

import (
	"crypto/tls"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	kproto "github.com/Kinetic/kinetic-go/proto"
	proto "github.com/golang/protobuf/proto"
)

var (
	networkTimeout = 20 * time.Second
)

func newMessage(t kproto.Message_AuthType) *kproto.Message {
	msg := &kproto.Message{
		AuthType: t.Enum(),
	}
	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.HmacAuth = &kproto.Message_HMACauth{}
	}

	return msg
}

func newCommand(t kproto.Command_MessageType) *kproto.Command {
	return &kproto.Command{
		Header: &kproto.Command_Header{
			MessageType: t.Enum(),
		},
	}
}

type networkService struct {
	rxMu           sync.Mutex
	txMu           sync.Mutex
	mapMu          sync.Mutex
	conn           net.Conn
	clusterVersion int64                      // Cluster version
	seq            int64                      // Operation sequence ID
	connID         int64                      // current connection ID
	option         ClientOptions              // current connection operation
	hmap           map[int64]*ResponseHandler // Message handler map
	fatal          bool                       // Network has fatal failure
	fatalError     error                      // Network fatal error details
	device         Log                        // Store device information from handshake package
}

func newNetworkService(op ClientOptions) (*networkService, error) {
	var conn net.Conn
	var err error
	if op.Timeout > 0 {
		networkTimeout = time.Duration(op.Timeout) * time.Millisecond
	}

	target := fmt.Sprintf("%s:%d", op.Host, op.Port)
	if op.UseSSL {
		// TODO: Need to enable verify certification later
		config := tls.Config{InsecureSkipVerify: true}
		d := &net.Dialer{Timeout: networkTimeout}
		conn, err = tls.DialWithDialer(d, "tcp", target, &config)
	} else {
		conn, err = net.DialTimeout("tcp", target, networkTimeout)
	}

	if err != nil {
		klog.Error("Can't establish connection to ", op.Host, err)
		return nil, err
	}

	ns := &networkService{
		conn:           conn,
		clusterVersion: 0,
		seq:            0,
		connID:         -1,
		option:         op,
		hmap:           make(map[int64]*ResponseHandler),
		fatal:          false,
		fatalError:     nil,
	}

	ns.rxMu.Lock()
	// Do the handshake.
	// Device Configuration and Limits from handshake will be stored in networkService.device
	_, _, _, err = ns.receive()
	ns.rxMu.Unlock()

	if err != nil {
		klog.Error("Can't establish connection to %s", op.Host)
		return nil, err
	}

	klog.Debugf("Connected to %s:%d", op.Host, op.Port)
	klog.Debugf("\tVendor: %s", ns.device.Configuration.Vendor)
	klog.Debugf("\tModel: %s", ns.device.Configuration.Model)
	klog.Debugf("\tWorldWideName: %s", ns.device.Configuration.WorldWideName)
	klog.Debugf("\tSerial Number: %s", ns.device.Configuration.SerialNumber)
	klog.Debugf("\tFirmware Version: %s", ns.device.Configuration.Version)
	klog.Debugf("\tKinetic Protocol Version: %s", ns.device.Configuration.ProtocolVersion)
	klog.Debugf("\tPort: %d", ns.device.Configuration.Port)
	klog.Debugf("\tTlsPort: %d", ns.device.Configuration.TLSPort)
	klog.Debugf("\tCurrentPowerLevel : %s", ns.device.Configuration.CurrentPowerLevel.String())

	return ns, nil
}

// When client network service has error, call error handling
// from all Messagehandler current in Queue.
func (ns *networkService) clientError(s Status, mh *ResponseHandler) {
	ns.mapMu.Lock()
	for ack, h := range ns.hmap {
		h.fail(s)
		delete(ns.hmap, ack)
	}
	ns.mapMu.Unlock()

	if mh != nil {
		mh.fail(s)
	}
}

func (ns *networkService) listen() error {
	if ns.fatal {
		return errors.New("Can't listen, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.mapMu.Lock()
	if len(ns.hmap) == 0 {
		ns.mapMu.Unlock()
		return nil
	}
	ns.mapMu.Unlock()

	ns.rxMu.Lock()
	msg, cmd, value, err := ns.receive()
	ns.rxMu.Unlock()
	if err != nil {
		klog.Error("Network Service listen error")
		return err
	}

	if cmd.GetHeader() != nil {
		klog.Debug("Kinetic response received ", cmd.GetHeader().GetMessageType().String(),
			", AckSeq = ", cmd.GetHeader().GetAckSequence(),
			", Code = ", cmd.GetStatus().GetCode())
	} else if msg.GetAuthType() == kproto.Message_UNSOLICITEDSTATUS {
		klog.Debug("Kinetic UNSOLICITEDSTATUS : ",
			"Code = ", cmd.GetStatus().GetCode(),
			", StatusMessage = ", cmd.GetStatus().GetStatusMessage())
	}

	// For UNSOLICITEDSTATUS, command may not have Header or AckSequence, set the ack to -1 so
	// no ResponseHandler will be found from hmap table.
	var ack int64 = -1
	if cmd.Header != nil && cmd.Header.AckSequence != nil {
		ack = cmd.GetHeader().GetAckSequence()
	}

	ns.mapMu.Lock()
	h, ok := ns.hmap[ack]
	ns.mapMu.Unlock()
	if ok == false {
		// It's high chance this is an UNSOLICITEDSTATUS message, display the Status.
		klog.Errorf("Couldn't find a handler for acksequence %d, status=%s", ack, getStatusFromProto(cmd).String())
		// This is an unexpected packet. Each listen() call expect remove one ResponseHandler from hmap.
		// So need to fire another listen() to make sure ResponseHandler in hmap got chance to exit.
		// Either by receive correct packet, or network read failure.
		go ns.listen()
		return nil
	}

	h.handle(cmd, value)

	ns.mapMu.Lock()
	delete(ns.hmap, ack)
	ns.mapMu.Unlock()

	return nil
}

// submit will send the message to kinetic device, insert ResponseHandler for this message sequence number.
// ResponseHandler can be nil if the message no require for Ack, eg batch PUT / DELETE.
func (ns *networkService) submit(msg *kproto.Message, cmd *kproto.Command, value []byte, h *ResponseHandler) error {
	if ns.fatal {
		return errors.New("Can't submit, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.txMu.Lock()

	cmd.GetHeader().ConnectionID = &ns.connID
	cmd.GetHeader().Sequence = &ns.seq
	cmd.GetHeader().ClusterVersion = &ns.clusterVersion

	cmdBytes, err := proto.Marshal(cmd)
	if err != nil {
		klog.Error("Error marshl Kinetic Command")
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Command"}
		ns.clientError(s, h)
		return err
	}
	msg.CommandBytes = cmdBytes[:]

	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.GetHmacAuth().Identity = &ns.option.User
		msg.GetHmacAuth().Hmac = computeHmac(msg.CommandBytes, ns.option.Hmac)
	}

	klog.Debug("Kinetic message send ", cmd.GetHeader().GetMessageType().String(), " Seq = ", ns.seq)

	err = ns.send(msg, value)

	if err != nil {
		return err
	}

	if h != nil {
		ns.mapMu.Lock()
		ns.hmap[ns.seq] = h
		ns.mapMu.Unlock()
	}

	ns.seq++
	ns.txMu.Unlock()

	return nil
}

func (ns *networkService) send(msg *kproto.Message, value []byte) error {
	msgBytes, err := proto.Marshal(msg)
	if err != nil {
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Message"}
		ns.clientError(s, nil)
		return err
	}

	// Set timeout for send packet
	ns.conn.SetWriteDeadline(time.Now().Add(networkTimeout))

	// Construct message header 9 bytes
	header := make([]byte, 9)
	header[0] = 'F' // Magic number
	binary.BigEndian.PutUint32(header[1:5], uint32(len(msgBytes)))
	binary.BigEndian.PutUint32(header[5:9], uint32(len(value)))

	packet := append(header, msgBytes...)
	if value != nil && len(value) > 0 {
		packet = append(packet, value...)
	}

	_, err = ns.conn.Write(packet)
	if err != nil {
		klog.Error("Network I/O write error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O write error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return err
	}

	return nil
}

func (ns *networkService) receive() (*kproto.Message, *kproto.Command, []byte, error) {
	// Set timeout for receive packet
	ns.conn.SetReadDeadline(time.Now().Add(networkTimeout))

	header := make([]byte, 9)

	_, err := io.ReadFull(ns.conn, header[0:])
	if err != nil {
		klog.Error("Network I/O read error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	magic := header[0]
	if magic != 'F' {
		klog.Error("Network I/O read error Header wrong magic")
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error Header wrong magic"}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = errors.New("Wrong magic number")
		return nil, nil, nil, errors.New("Network I/O read error Header wrong magic")
	}

	protoLen := int(binary.BigEndian.Uint32(header[1:5]))
	valueLen := int(binary.BigEndian.Uint32(header[5:9]))

	protoBuf := make([]byte, protoLen)
	_, err = io.ReadFull(ns.conn, protoBuf)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error receive Kinetic Header, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	msg := &kproto.Message{}
	err = proto.Unmarshal(protoBuf, msg)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error reaceive Kinetic Message, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if msg.GetAuthType() == kproto.Message_HMACAUTH && validateHmac(msg, ns.option.Hmac) == false {
		klog.Error("Response HMAC mismatch")
		s := Status{Code: ClientResponseHMACError, ErrorMsg: "Response HMAC mismatch"}
		ns.clientError(s, nil)
		return nil, nil, nil, err
	}

	cmd := &kproto.Command{}
	err = proto.Unmarshal(msg.CommandBytes, cmd)
	if err != nil {
		klog.Error("Network I/O read error parsing Kinetic Command, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Command, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if cmd.Header != nil && cmd.Header.ConnectionID != nil {
		if ns.connID < 0 {
			// This is handshake packet
			ns.device = getLogFromProto(cmd)

			// Only update client cluster version during Handshake
			if cmd.Header.ClusterVersion != nil {
				ns.clusterVersion = cmd.GetHeader().GetClusterVersion()
			}
		}
		ns.connID = cmd.GetHeader().GetConnectionID()
	}

	if valueLen > 0 {
		valueBuf := make([]byte, valueLen)
		_, err = io.ReadFull(ns.conn, valueBuf)
		if err != nil {
			klog.Error("Network I/O read error parsing Kinetic Value, " + err.Error())
			s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Value, " + err.Error()}
			ns.clientError(s, nil)
			ns.fatal = true
			ns.fatalError = err
			return nil, nil, nil, err
		}

		return msg, cmd, valueBuf, nil
	}

	return msg, cmd, nil, nil
}

func (ns *networkService) close() {
	ns.conn.Close()
	klog.Debugf("Connection to %s closed", ns.option.Host)
}
<commit_msg>Fix panic caused by access nil data structure.<commit_after>/**
 * Copyright 2013-2016 Seagate Technology LLC.
 *
 * This Source Code Form is subject to the terms of the Mozilla
 * Public License, v. 2.0. If a copy of the MPL was not
 * distributed with this file, You can obtain one at
 * https://mozilla.org/MP:/2.0/.
 *
 * This program is distributed in the hope that it will be useful,
 * but is provided AS-IS, WITHOUT ANY WARRANTY; including without
 * the implied warranty of MERCHANTABILITY, NON-INFRINGEMENT or
 * FITNESS FOR A PARTICULAR PURPOSE. See the Mozilla Public
 * License for more details.
 *
 * See www.openkinetic.org for more project information
 */

package kinetic

import (
	"crypto/tls"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	kproto "github.com/Kinetic/kinetic-go/proto"
	proto "github.com/golang/protobuf/proto"
)

var (
	networkTimeout = 20 * time.Second
)

func newMessage(t kproto.Message_AuthType) *kproto.Message {
	msg := &kproto.Message{
		AuthType: t.Enum(),
	}
	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.HmacAuth = &kproto.Message_HMACauth{}
	}

	return msg
}

func newCommand(t kproto.Command_MessageType) *kproto.Command {
	return &kproto.Command{
		Header: &kproto.Command_Header{
			MessageType: t.Enum(),
		},
	}
}

type networkService struct {
	rxMu           sync.Mutex
	txMu           sync.Mutex
	mapMu          sync.Mutex
	conn           net.Conn
	clusterVersion int64                      // Cluster version
	seq            int64                      // Operation sequence ID
	connID         int64                      // current connection ID
	option         ClientOptions              // current connection operation
	hmap           map[int64]*ResponseHandler // Message handler map
	fatal          bool                       // Network has fatal failure
	fatalError     error                      // Network fatal error details
	device         Log                        // Store device information from handshake package
}

func newNetworkService(op ClientOptions) (*networkService, error) {
	var conn net.Conn
	var err error
	if op.Timeout > 0 {
		networkTimeout = time.Duration(op.Timeout) * time.Millisecond
	}

	target := fmt.Sprintf("%s:%d", op.Host, op.Port)
	if op.UseSSL {
		// TODO: Need to enable verify certification later
		config := tls.Config{InsecureSkipVerify: true}
		d := &net.Dialer{Timeout: networkTimeout}
		conn, err = tls.DialWithDialer(d, "tcp", target, &config)
	} else {
		conn, err = net.DialTimeout("tcp", target, networkTimeout)
	}

	if err != nil {
		klog.Error("Can't establish connection to ", op.Host, err)
		return nil, err
	}

	ns := &networkService{
		conn:           conn,
		clusterVersion: 0,
		seq:            0,
		connID:         -1,
		option:         op,
		hmap:           make(map[int64]*ResponseHandler),
		fatal:          false,
		fatalError:     nil,
	}

	ns.rxMu.Lock()
	// Do the handshake.
	// Device Configuration and Limits from handshake will be stored in networkService.device
	_, _, _, err = ns.receive()
	ns.rxMu.Unlock()

	if err != nil {
		klog.Error("Can't establish connection to %s", op.Host)
		return nil, err
	}

	klog.Debugf("Connected to %s:%d", op.Host, op.Port)
	klog.Debugf("\tVendor: %s", ns.device.Configuration.Vendor)
	klog.Debugf("\tModel: %s", ns.device.Configuration.Model)
	klog.Debugf("\tWorldWideName: %s", ns.device.Configuration.WorldWideName)
	klog.Debugf("\tSerial Number: %s", ns.device.Configuration.SerialNumber)
	klog.Debugf("\tFirmware Version: %s", ns.device.Configuration.Version)
	klog.Debugf("\tKinetic Protocol Version: %s", ns.device.Configuration.ProtocolVersion)
	klog.Debugf("\tPort: %d", ns.device.Configuration.Port)
	klog.Debugf("\tTlsPort: %d", ns.device.Configuration.TLSPort)
	klog.Debugf("\tCurrentPowerLevel : %s", ns.device.Configuration.CurrentPowerLevel.String())

	return ns, nil
}

// When client network service has error, call error handling
// from all Messagehandler current in Queue.
func (ns *networkService) clientError(s Status, mh *ResponseHandler) {
	ns.mapMu.Lock()
	for ack, h := range ns.hmap {
		h.fail(s)
		delete(ns.hmap, ack)
	}
	ns.mapMu.Unlock()

	if mh != nil {
		mh.fail(s)
	}
}

func (ns *networkService) listen() error {
	if ns.fatal {
		return errors.New("Can't listen, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.mapMu.Lock()
	if len(ns.hmap) == 0 {
		ns.mapMu.Unlock()
		return nil
	}
	ns.mapMu.Unlock()

	ns.rxMu.Lock()
	msg, cmd, value, err := ns.receive()
	ns.rxMu.Unlock()
	if err != nil {
		klog.Error("Network Service listen error")
		return err
	}

	if cmd.GetHeader() != nil {
		klog.Debug("Kinetic response received ", cmd.GetHeader().GetMessageType().String(),
			", AckSeq = ", cmd.GetHeader().GetAckSequence(),
			", Code = ", cmd.GetStatus().GetCode())
	} else if msg.GetAuthType() == kproto.Message_UNSOLICITEDSTATUS {
		klog.Debug("Kinetic UNSOLICITEDSTATUS : ",
			"Code = ", cmd.GetStatus().GetCode(),
			", StatusMessage = ", cmd.GetStatus().GetStatusMessage())
	}

	// For UNSOLICITEDSTATUS, command may not have Header or AckSequence, set the ack to -1 so
	// no ResponseHandler will be found from hmap table.
	var ack int64 = -1
	if cmd.GetHeader() != nil {
		ack = cmd.GetHeader().GetAckSequence()
	}

	ns.mapMu.Lock()
	h, ok := ns.hmap[ack]
	ns.mapMu.Unlock()
	if ok == false {
		// It's high chance this is an UNSOLICITEDSTATUS message, display the Status.
		klog.Errorf("Couldn't find a handler for acksequence %d, status=%s", ack, getStatusFromProto(cmd).String())
		// This is an unexpected packet. Each listen() call expect remove one ResponseHandler from hmap.
		// So need to fire another listen() to make sure ResponseHandler in hmap got chance to exit.
		// Either by receive correct packet, or network read failure.
		go ns.listen()
		return nil
	}

	h.handle(cmd, value)

	ns.mapMu.Lock()
	delete(ns.hmap, ack)
	ns.mapMu.Unlock()

	return nil
}

// submit will send the message to kinetic device, insert ResponseHandler for this message sequence number.
// ResponseHandler can be nil if the message no require for Ack, eg batch PUT / DELETE.
func (ns *networkService) submit(msg *kproto.Message, cmd *kproto.Command, value []byte, h *ResponseHandler) error {
	if ns.fatal {
		return errors.New("Can't submit, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.txMu.Lock()

	cmd.GetHeader().ConnectionID = &ns.connID
	cmd.GetHeader().Sequence = &ns.seq
	cmd.GetHeader().ClusterVersion = &ns.clusterVersion

	cmdBytes, err := proto.Marshal(cmd)
	if err != nil {
		klog.Error("Error marshl Kinetic Command")
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Command"}
		ns.clientError(s, h)
		return err
	}
	msg.CommandBytes = cmdBytes[:]

	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.GetHmacAuth().Identity = &ns.option.User
		msg.GetHmacAuth().Hmac = computeHmac(msg.CommandBytes, ns.option.Hmac)
	}

	klog.Debug("Kinetic message send ", cmd.GetHeader().GetMessageType().String(), " Seq = ", ns.seq)

	err = ns.send(msg, value)

	if err != nil {
		return err
	}

	if h != nil {
		ns.mapMu.Lock()
		ns.hmap[ns.seq] = h
		ns.mapMu.Unlock()
	}

	ns.seq++
	ns.txMu.Unlock()

	return nil
}

func (ns *networkService) send(msg *kproto.Message, value []byte) error {
	msgBytes, err := proto.Marshal(msg)
	if err != nil {
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Message"}
		ns.clientError(s, nil)
		return err
	}

	// Set timeout for send packet
	ns.conn.SetWriteDeadline(time.Now().Add(networkTimeout))

	// Construct message header 9 bytes
	header := make([]byte, 9)
	header[0] = 'F' // Magic number
	binary.BigEndian.PutUint32(header[1:5], uint32(len(msgBytes)))
	binary.BigEndian.PutUint32(header[5:9], uint32(len(value)))

	packet := append(header, msgBytes...)
	if value != nil && len(value) > 0 {
		packet = append(packet, value...)
	}

	_, err = ns.conn.Write(packet)
	if err != nil {
		klog.Error("Network I/O write error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O write error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return err
	}

	return nil
}

func (ns *networkService) receive() (*kproto.Message, *kproto.Command, []byte, error) {
	// Set timeout for receive packet
	ns.conn.SetReadDeadline(time.Now().Add(networkTimeout))

	header := make([]byte, 9)

	_, err := io.ReadFull(ns.conn, header[0:])
	if err != nil {
		klog.Error("Network I/O read error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	magic := header[0]
	if magic != 'F' {
		klog.Error("Network I/O read error Header wrong magic")
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error Header wrong magic"}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = errors.New("Wrong magic number")
		return nil, nil, nil, errors.New("Network I/O read error Header wrong magic")
	}

	protoLen := int(binary.BigEndian.Uint32(header[1:5]))
	valueLen := int(binary.BigEndian.Uint32(header[5:9]))

	protoBuf := make([]byte, protoLen)
	_, err = io.ReadFull(ns.conn, protoBuf)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error receive Kinetic Header, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	msg := &kproto.Message{}
	err = proto.Unmarshal(protoBuf, msg)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error reaceive Kinetic Message, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if msg.GetAuthType() == kproto.Message_HMACAUTH && validateHmac(msg, ns.option.Hmac) == false {
		klog.Error("Response HMAC mismatch")
		s := Status{Code: ClientResponseHMACError, ErrorMsg: "Response HMAC mismatch"}
		ns.clientError(s, nil)
		return nil, nil, nil, err
	}

	cmd := &kproto.Command{}
	err = proto.Unmarshal(msg.CommandBytes, cmd)
	if err != nil {
		klog.Error("Network I/O read error parsing Kinetic Command, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Command, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if cmd.Header != nil && cmd.Header.ConnectionID != nil {
		if ns.connID < 0 {
			// This is handshake packet
			ns.device = getLogFromProto(cmd)

			// Only update client cluster version during Handshake
			if cmd.Header.ClusterVersion != nil {
				ns.clusterVersion = cmd.GetHeader().GetClusterVersion()
			}
		}
		ns.connID = cmd.GetHeader().GetConnectionID()
	}

	if valueLen > 0 {
		valueBuf := make([]byte, valueLen)
		_, err = io.ReadFull(ns.conn, valueBuf)
		if err != nil {
			klog.Error("Network I/O read error parsing Kinetic Value, " + err.Error())
			s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Value, " + err.Error()}
			ns.clientError(s, nil)
			ns.fatal = true
			ns.fatalError = err
			return nil, nil, nil, err
		}

		return msg, cmd, valueBuf, nil
	}

	return msg, cmd, nil, nil
}

func (ns *networkService) close() {
	ns.conn.Close()
	klog.Debugf("Connection to %s closed", ns.option.Host)
}
<|endoftext|>


<commit_before>// Copyright 2016 go-dockerclient authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"encoding/json"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/docker/docker/api/types/swarm"
	"golang.org/x/net/context"
)

// NoSuchService is the error returned when a given service does not exist.
type NoSuchService struct {
	ID  string
	Err error
}

func (err *NoSuchService) Error() string {
	if err.Err != nil {
		return err.Err.Error()
	}
	return "No such service: " + err.ID
}

// CreateServiceOptions specify parameters to the CreateService function.
//
// See https://goo.gl/KrVjHz for more details.
type CreateServiceOptions struct {
	Auth AuthConfiguration `qs:"-"`
	swarm.ServiceSpec
	Context context.Context
}

// CreateService creates a new service, returning the service instance
// or an error in case of failure.
//
// See https://goo.gl/KrVjHz for more details.
func (c *Client) CreateService(opts CreateServiceOptions) (*swarm.Service, error) {
	headers, err := headersWithAuth(opts.Auth)
	if err != nil {
		return nil, err
	}
	path := "/services/create?" + queryString(opts)
	resp, err := c.do("POST", path, doOptions{
		headers:   headers,
		data:      opts.ServiceSpec,
		forceJSON: true,
		context:   opts.Context,
	})
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var service swarm.Service
	if err := json.NewDecoder(resp.Body).Decode(&service); err != nil {
		return nil, err
	}
	return &service, nil
}

// RemoveServiceOptions encapsulates options to remove a service.
//
// See https://goo.gl/Tqrtya for more details.
type RemoveServiceOptions struct {
	ID      string `qs:"-"`
	Context context.Context
}

// RemoveService removes a service, returning an error in case of failure.
//
// See https://goo.gl/Tqrtya for more details.
func (c *Client) RemoveService(opts RemoveServiceOptions) error {
	path := "/services/" + opts.ID
	resp, err := c.do("DELETE", path, doOptions{context: opts.Context})
	if err != nil {
		if e, ok := err.(*Error); ok && e.Status == http.StatusNotFound {
			return &NoSuchService{ID: opts.ID}
		}
		return err
	}
	resp.Body.Close()
	return nil
}

// UpdateServiceOptions specify parameters to the UpdateService function.
//
// See https://goo.gl/wu3MmS for more details.
type UpdateServiceOptions struct {
	Auth AuthConfiguration `qs:"-"`
	swarm.ServiceSpec
	Context context.Context
	Version uint64
}

// UpdateService updates the service at ID with the options
//
// See https://goo.gl/wu3MmS for more details.
func (c *Client) UpdateService(id string, opts UpdateServiceOptions) error {
	headers, err := headersWithAuth(opts.Auth)
	if err != nil {
		return err
	}
	params := make(url.Values)
	params.Set("version", strconv.FormatUint(opts.Version, 10))
	resp, err := c.do("POST", "/services/"+id+"/update?"+params.Encode(), doOptions{
		headers:   headers,
		data:      opts.ServiceSpec,
		forceJSON: true,
		context:   opts.Context,
	})
	if err != nil {
		if e, ok := err.(*Error); ok && e.Status == http.StatusNotFound {
			return &NoSuchService{ID: id}
		}
		return err
	}
	defer resp.Body.Close()
	return nil
}

// InspectService returns information about a service by its ID.
//
// See https://goo.gl/dHmr75 for more details.
func (c *Client) InspectService(id string) (*swarm.Service, error) {
	path := "/services/" + id
	resp, err := c.do("GET", path, doOptions{})
	if err != nil {
		if e, ok := err.(*Error); ok && e.Status == http.StatusNotFound {
			return nil, &NoSuchService{ID: id}
		}
		return nil, err
	}
	defer resp.Body.Close()
	var service swarm.Service
	if err := json.NewDecoder(resp.Body).Decode(&service); err != nil {
		return nil, err
	}
	return &service, nil
}

// ListServicesOptions specify parameters to the ListServices function.
//
// See https://goo.gl/DwvNMd for more details.
type ListServicesOptions struct {
	Filters map[string][]string
	Context context.Context
}

// ListServices returns a slice of services matching the given criteria.
//
// See https://goo.gl/DwvNMd for more details.
func (c *Client) ListServices(opts ListServicesOptions) ([]swarm.Service, error) {
	path := "/services?" + queryString(opts)
	resp, err := c.do("GET", path, doOptions{context: opts.Context})
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var services []swarm.Service
	if err := json.NewDecoder(resp.Body).Decode(&services); err != nil {
		return nil, err
	}
	return services, nil
}

// LogsOptions represents the set of options used when getting logs from a
// container.
//
// See https://goo.gl/krK0ZH for more details.
type LogsServiceOptions struct {
	Context           context.Context
	Service           string        `qs:"-"`
	OutputStream      io.Writer     `qs:"-"`
	ErrorStream       io.Writer     `qs:"-"`
	InactivityTimeout time.Duration `qs:"-"`
	Tail              string

	// Use raw terminal? Usually true when the container contains a TTY.
	RawTerminal bool `qs:"-"`
	Since       int64
	Follow      bool
	Stdout      bool
	Stderr      bool
	Timestamps  bool
	Details     bool
}

// Logs gets stdout and stderr logs from the specified service.
//
// When LogsOptions.RawTerminal is set to false, go-dockerclient will multiplex
// the streams and send the containers stdout to LogsOptions.OutputStream, and
// stderr to LogsOptions.ErrorStream.
//
// When LogsOptions.RawTerminal is true, callers will get the raw stream on
// LogOptions.OutputStream. The caller can use libraries such as dlog
// (github.com/ahmetalpbalkan/dlog).
//
// See https://goo.gl/krK0ZH for more details.
func (c *Client) LogsService(opts LogsServiceOptions) error {
	if opts.Service == "" {
		return &NoSuchService{ID: opts.Service}
	}
	if opts.Tail == "" {
		opts.Tail = "all"
	}
	path := "/services/" + opts.Service + "/logs?" + queryString(opts)
	return c.stream("GET", path, streamOptions{
		setRawTerminal:    opts.RawTerminal,
		stdout:            opts.OutputStream,
		stderr:            opts.ErrorStream,
		inactivityTimeout: opts.InactivityTimeout,
		context:           opts.Context,
	})
}
<commit_msg>add service logs comments<commit_after>// Copyright 2016 go-dockerclient authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"encoding/json"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/docker/docker/api/types/swarm"
	"golang.org/x/net/context"
)

// NoSuchService is the error returned when a given service does not exist.
type NoSuchService struct {
	ID  string
	Err error
}

func (err *NoSuchService) Error() string {
	if err.Err != nil {
		return err.Err.Error()
	}
	return "No such service: " + err.ID
}

// CreateServiceOptions specify parameters to the CreateService function.
//
// See https://goo.gl/KrVjHz for more details.
type CreateServiceOptions struct {
	Auth AuthConfiguration `qs:"-"`
	swarm.ServiceSpec
	Context context.Context
}

// CreateService creates a new service, returning the service instance
// or an error in case of failure.
//
// See https://goo.gl/KrVjHz for more details.
func (c *Client) CreateService(opts CreateServiceOptions) (*swarm.Service, error) {
	headers, err := headersWithAuth(opts.Auth)
	if err != nil {
		return nil, err
	}
	path := "/services/create?" + queryString(opts)
	resp, err := c.do("POST", path, doOptions{
		headers:   headers,
		data:      opts.ServiceSpec,
		forceJSON: true,
		context:   opts.Context,
	})
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var service swarm.Service
	if err := json.NewDecoder(resp.Body).Decode(&service); err != nil {
		return nil, err
	}
	return &service, nil
}

// RemoveServiceOptions encapsulates options to remove a service.
//
// See https://goo.gl/Tqrtya for more details.
type RemoveServiceOptions struct {
	ID      string `qs:"-"`
	Context context.Context
}

// RemoveService removes a service, returning an error in case of failure.
//
// See https://goo.gl/Tqrtya for more details.
func (c *Client) RemoveService(opts RemoveServiceOptions) error {
	path := "/services/" + opts.ID
	resp, err := c.do("DELETE", path, doOptions{context: opts.Context})
	if err != nil {
		if e, ok := err.(*Error); ok && e.Status == http.StatusNotFound {
			return &NoSuchService{ID: opts.ID}
		}
		return err
	}
	resp.Body.Close()
	return nil
}

// UpdateServiceOptions specify parameters to the UpdateService function.
//
// See https://goo.gl/wu3MmS for more details.
type UpdateServiceOptions struct {
	Auth AuthConfiguration `qs:"-"`
	swarm.ServiceSpec
	Context context.Context
	Version uint64
}

// UpdateService updates the service at ID with the options
//
// See https://goo.gl/wu3MmS for more details.
func (c *Client) UpdateService(id string, opts UpdateServiceOptions) error {
	headers, err := headersWithAuth(opts.Auth)
	if err != nil {
		return err
	}
	params := make(url.Values)
	params.Set("version", strconv.FormatUint(opts.Version, 10))
	resp, err := c.do("POST", "/services/"+id+"/update?"+params.Encode(), doOptions{
		headers:   headers,
		data:      opts.ServiceSpec,
		forceJSON: true,
		context:   opts.Context,
	})
	if err != nil {
		if e, ok := err.(*Error); ok && e.Status == http.StatusNotFound {
			return &NoSuchService{ID: id}
		}
		return err
	}
	defer resp.Body.Close()
	return nil
}

// InspectService returns information about a service by its ID.
//
// See https://goo.gl/dHmr75 for more details.
func (c *Client) InspectService(id string) (*swarm.Service, error) {
	path := "/services/" + id
	resp, err := c.do("GET", path, doOptions{})
	if err != nil {
		if e, ok := err.(*Error); ok && e.Status == http.StatusNotFound {
			return nil, &NoSuchService{ID: id}
		}
		return nil, err
	}
	defer resp.Body.Close()
	var service swarm.Service
	if err := json.NewDecoder(resp.Body).Decode(&service); err != nil {
		return nil, err
	}
	return &service, nil
}

// ListServicesOptions specify parameters to the ListServices function.
//
// See https://goo.gl/DwvNMd for more details.
type ListServicesOptions struct {
	Filters map[string][]string
	Context context.Context
}

// ListServices returns a slice of services matching the given criteria.
//
// See https://goo.gl/DwvNMd for more details.
func (c *Client) ListServices(opts ListServicesOptions) ([]swarm.Service, error) {
	path := "/services?" + queryString(opts)
	resp, err := c.do("GET", path, doOptions{context: opts.Context})
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var services []swarm.Service
	if err := json.NewDecoder(resp.Body).Decode(&services); err != nil {
		return nil, err
	}
	return services, nil
}

// LogsServiceOptions represents the set of options used when getting logs from a
// service.
type LogsServiceOptions struct {
	Context           context.Context
	Service           string        `qs:"-"`
	OutputStream      io.Writer     `qs:"-"`
	ErrorStream       io.Writer     `qs:"-"`
	InactivityTimeout time.Duration `qs:"-"`
	Tail              string

	// Use raw terminal? Usually true when the container contains a TTY.
	RawTerminal bool `qs:"-"`
	Since       int64
	Follow      bool
	Stdout      bool
	Stderr      bool
	Timestamps  bool
	Details     bool
}

// LogsService gets stdout and stderr logs from the specified service.
//
// When LogsServiceOptions.RawTerminal is set to false, go-dockerclient will multiplex
// the streams and send the containers stdout to LogsServiceOptions.OutputStream, and
// stderr to LogsServiceOptions.ErrorStream.
//
// When LogsServiceOptions.RawTerminal is true, callers will get the raw stream on
// LogsServiceOptions.OutputStream.
func (c *Client) LogsService(opts LogsServiceOptions) error {
	if opts.Service == "" {
		return &NoSuchService{ID: opts.Service}
	}
	if opts.Tail == "" {
		opts.Tail = "all"
	}
	path := "/services/" + opts.Service + "/logs?" + queryString(opts)
	return c.stream("GET", path, streamOptions{
		setRawTerminal:    opts.RawTerminal,
		stdout:            opts.OutputStream,
		stderr:            opts.ErrorStream,
		inactivityTimeout: opts.InactivityTimeout,
		context:           opts.Context,
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes 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 task

import (
	"fmt"
	"time"

	"github.com/golang/glog"

	"k8s.io/kubernetes/pkg/client/cache"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/util/workqueue"
)

var (
	keyFunc = cache.DeletionHandlingMetaNamespaceKeyFunc
)

// Queue manages a work queue through an independent worker that
// invokes the given sync function for every work item inserted.
type Queue struct {
	// queue is the work queue the worker polls
	queue workqueue.RateLimitingInterface
	// sync is called for each item in the queue
	sync func(interface{}) error
	// workerDone is closed when the worker exits
	workerDone chan struct{}

	fn func(obj interface{}) (interface{}, error)
}

// Run ...
func (t *Queue) Run(period time.Duration, stopCh <-chan struct{}) {
	wait.Until(t.worker, period, stopCh)
}

// Enqueue enqueues ns/name of the given api object in the task queue.
func (t *Queue) Enqueue(obj interface{}) {
	glog.V(3).Infof("queuing item %v", obj)
	key, err := t.fn(obj)
	if err != nil {
		glog.Errorf("%v", err)
		return
	}
	t.queue.Add(key)
}

func (t *Queue) defaultKeyFunc(obj interface{}) (interface{}, error) {
	key, err := keyFunc(obj)
	if err != nil {
		return "", fmt.Errorf("could not get key for object %+v: %v", obj, err)
	}

	return key, nil
}

// worker processes work in the queue through sync.
func (t *Queue) worker() {
	for {
		key, quit := t.queue.Get()
		if quit {
			close(t.workerDone)
			return
		}
		glog.V(3).Infof("syncing %v", key)
		if err := t.sync(key); err != nil {
			glog.Warningf("requeuing %v, err %v", key, err)
			t.queue.AddRateLimited(key)
		} else {
			t.queue.Forget(key)
		}

		t.queue.Done(key)
	}
}

// Shutdown shuts down the work queue and waits for the worker to ACK
func (t *Queue) Shutdown() {
	t.queue.ShutDown()
	<-t.workerDone
}

// IsShuttingDown returns if the method Shutdown was invoked
func (t *Queue) IsShuttingDown() bool {
	return t.queue.ShuttingDown()
}

// NewTaskQueue creates a new task queue with the given sync function.
// The sync function is called for every element inserted into the queue.
func NewTaskQueue(syncFn func(interface{}) error) *Queue {
	return NewCustomTaskQueue(syncFn, nil)
}

// NewCustomTaskQueue ...
func NewCustomTaskQueue(syncFn func(interface{}) error, fn func(interface{}) (interface{}, error)) *Queue {
	q := &Queue{
		queue:      workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()),
		sync:       syncFn,
		workerDone: make(chan struct{}),
		fn:         fn,
	}

	if fn == nil {
		q.fn = q.defaultKeyFunc
	}

	return q
}
<commit_msg>check queue's status before enqueue<commit_after>/*
Copyright 2015 The Kubernetes 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 task

import (
	"fmt"
	"time"

	"github.com/golang/glog"

	"k8s.io/kubernetes/pkg/client/cache"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/util/workqueue"
)

var (
	keyFunc = cache.DeletionHandlingMetaNamespaceKeyFunc
)

// Queue manages a work queue through an independent worker that
// invokes the given sync function for every work item inserted.
type Queue struct {
	// queue is the work queue the worker polls
	queue workqueue.RateLimitingInterface
	// sync is called for each item in the queue
	sync func(interface{}) error
	// workerDone is closed when the worker exits
	workerDone chan struct{}

	fn func(obj interface{}) (interface{}, error)
}

// Run ...
func (t *Queue) Run(period time.Duration, stopCh <-chan struct{}) {
	wait.Until(t.worker, period, stopCh)
}

// Enqueue enqueues ns/name of the given api object in the task queue.
func (t *Queue) Enqueue(obj interface{}) {
	if t.IsShuttingDown() {
		glog.Errorf("queue has been shutdown, failed to enqueue: %v", obj)
		return
	}

	glog.V(3).Infof("queuing item %v", obj)
	key, err := t.fn(obj)
	if err != nil {
		glog.Errorf("%v", err)
		return
	}
	t.queue.Add(key)
}

func (t *Queue) defaultKeyFunc(obj interface{}) (interface{}, error) {
	key, err := keyFunc(obj)
	if err != nil {
		return "", fmt.Errorf("could not get key for object %+v: %v", obj, err)
	}

	return key, nil
}

// worker processes work in the queue through sync.
func (t *Queue) worker() {
	for {
		key, quit := t.queue.Get()
		if quit {
			close(t.workerDone)
			return
		}
		glog.V(3).Infof("syncing %v", key)
		if err := t.sync(key); err != nil {
			glog.Warningf("requeuing %v, err %v", key, err)
			t.queue.AddRateLimited(key)
		} else {
			t.queue.Forget(key)
		}

		t.queue.Done(key)
	}
}

// Shutdown shuts down the work queue and waits for the worker to ACK
func (t *Queue) Shutdown() {
	t.queue.ShutDown()
	<-t.workerDone
}

// IsShuttingDown returns if the method Shutdown was invoked
func (t *Queue) IsShuttingDown() bool {
	return t.queue.ShuttingDown()
}

// NewTaskQueue creates a new task queue with the given sync function.
// The sync function is called for every element inserted into the queue.
func NewTaskQueue(syncFn func(interface{}) error) *Queue {
	return NewCustomTaskQueue(syncFn, nil)
}

// NewCustomTaskQueue ...
func NewCustomTaskQueue(syncFn func(interface{}) error, fn func(interface{}) (interface{}, error)) *Queue {
	q := &Queue{
		queue:      workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()),
		sync:       syncFn,
		workerDone: make(chan struct{}),
		fn:         fn,
	}

	if fn == nil {
		q.fn = q.defaultKeyFunc
	}

	return q
}
<|endoftext|>


<commit_before>// Code generated by "stringer -type=NodeState"; DO NOT EDIT

package cluster

import "fmt"

const _NodeState_name = "NodeNotReadyNodeReadyNodeUnreachable"

var _NodeState_index = [...]uint8{0, 12, 21, 36}

func (i NodeState) String() string {
	if i < 0 || i >= NodeState(len(_NodeState_index)-1) {
		return fmt.Sprintf("NodeState(%d)", i)
	}
	return _NodeState_name[_NodeState_index[i]:_NodeState_index[i+1]]
}
<commit_msg>new stringer<commit_after>// Code generated by "stringer -type=NodeState"; DO NOT EDIT.

package cluster

import "fmt"

const _NodeState_name = "NodeNotReadyNodeReadyNodeUnreachable"

var _NodeState_index = [...]uint8{0, 12, 21, 36}

func (i NodeState) String() string {
	if i < 0 || i >= NodeState(len(_NodeState_index)-1) {
		return fmt.Sprintf("NodeState(%d)", i)
	}
	return _NodeState_name[_NodeState_index[i]:_NodeState_index[i+1]]
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2022 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"context"
	"fmt"
	"path/filepath"

	humanize "github.com/dustin/go-humanize"
	"github.com/fatih/color"
	"github.com/minio/cli"
	json "github.com/minio/colorjson"
	"github.com/minio/madmin-go"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/pkg/console"
)

func mainAdminSpeedtestDrive(ctx *cli.Context, aliasedURL string) error {
	client, perr := newAdminClient(aliasedURL)
	if perr != nil {
		fatalIf(perr.Trace(aliasedURL), "Unable to initialize admin client.")
		return nil
	}

	ctxt, cancel := context.WithCancel(globalContext)
	defer cancel()

	blocksize, e := humanize.ParseBytes(ctx.String("blocksize"))
	if e != nil {
		fatalIf(probe.NewError(e), "Unable to parse blocksize")
		return nil
	}
	if blocksize < 0 {
		fatalIf(errInvalidArgument(), "blocksize cannot be <= 0")
		return nil
	}

	filesize, e := humanize.ParseBytes(ctx.String("filesize"))
	if e != nil {
		fatalIf(probe.NewError(e), "Unable to parse filesize")
		return nil
	}
	if filesize < 0 {
		fatalIf(errInvalidArgument(), "filesize cannot be <= 0")
		return nil
	}

	serial := ctx.Bool("serial")

	resultCh, err := client.DriveSpeedtest(ctxt, madmin.DriveSpeedTestOpts{
		Serial:    serial,
		BlockSize: uint64(blocksize),
		FileSize:  uint64(filesize),
	})
	fatalIf(probe.NewError(err), "Failed to execute drive speedtest")

	for result := range resultCh {
		if result.Version != "" {
			if globalJSON {
				jsonBytes, e := json.MarshalIndent(result, "", " ")
				fatalIf(probe.NewError(e), "Unable to marshal into JSON.")
				fmt.Println(string(jsonBytes))
				continue
			}
			tbl := console.NewTable(func() []*color.Color {
				colors := []*color.Color{
					color.New(color.FgWhite, color.BgBlack, color.Bold),
				}
				for range result.DrivePerf {
					colors = append(colors, color.New(color.FgGreen))
				}
				return colors
			}(), []bool{false, false, false, false}, 0)
			cellText := make([][]string, len(result.DrivePerf)+1)
			cellText[0] = []string{
				"Node",
				"Path",
				"Read",
				"Write",
			}
			trailerIfGreaterThan := func(in string, max int) string {
				if len(in) < max {
					return ""
				}
				return "..."
			}
			for i, item := range result.DrivePerf {
				cellText[i+1] = []string{
					fmt.Sprintf("%.64s%s", result.Endpoint,
						trailerIfGreaterThan(result.Endpoint, 64)),
					fmt.Sprintf("%.64s%s", filepath.Dir(item.Path),
						trailerIfGreaterThan(result.Endpoint, 64)),
					humanize.IBytes(item.ReadThroughput) + "/s",
					humanize.IBytes(item.WriteThroughput) + "/s",
				}
			}
			if len(result.DrivePerf) > 0 {
				tbl.DisplayTable(cellText)
			}
			if result.Error != "" {
				fmt.Println(color.New(color.FgRed, color.Bold).Sprintf("ERROR"), result.Error)
			}
		}
	}

	return nil
}
<commit_msg>remove trimming of top level path<commit_after>// Copyright (c) 2015-2022 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"context"
	"fmt"

	humanize "github.com/dustin/go-humanize"
	"github.com/fatih/color"
	"github.com/minio/cli"
	json "github.com/minio/colorjson"
	"github.com/minio/madmin-go"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/pkg/console"
)

func mainAdminSpeedtestDrive(ctx *cli.Context, aliasedURL string) error {
	client, perr := newAdminClient(aliasedURL)
	if perr != nil {
		fatalIf(perr.Trace(aliasedURL), "Unable to initialize admin client.")
		return nil
	}

	ctxt, cancel := context.WithCancel(globalContext)
	defer cancel()

	blocksize, e := humanize.ParseBytes(ctx.String("blocksize"))
	if e != nil {
		fatalIf(probe.NewError(e), "Unable to parse blocksize")
		return nil
	}
	if blocksize < 0 {
		fatalIf(errInvalidArgument(), "blocksize cannot be <= 0")
		return nil
	}

	filesize, e := humanize.ParseBytes(ctx.String("filesize"))
	if e != nil {
		fatalIf(probe.NewError(e), "Unable to parse filesize")
		return nil
	}
	if filesize < 0 {
		fatalIf(errInvalidArgument(), "filesize cannot be <= 0")
		return nil
	}

	serial := ctx.Bool("serial")

	resultCh, err := client.DriveSpeedtest(ctxt, madmin.DriveSpeedTestOpts{
		Serial:    serial,
		BlockSize: uint64(blocksize),
		FileSize:  uint64(filesize),
	})
	fatalIf(probe.NewError(err), "Failed to execute drive speedtest")

	for result := range resultCh {
		if result.Version != "" {
			if globalJSON {
				jsonBytes, e := json.MarshalIndent(result, "", " ")
				fatalIf(probe.NewError(e), "Unable to marshal into JSON.")
				fmt.Println(string(jsonBytes))
				continue
			}
			tbl := console.NewTable(func() []*color.Color {
				colors := []*color.Color{
					color.New(color.FgWhite, color.BgBlack, color.Bold),
				}
				for range result.DrivePerf {
					colors = append(colors, color.New(color.FgGreen))
				}
				return colors
			}(), []bool{false, false, false, false}, 0)
			cellText := make([][]string, len(result.DrivePerf)+1)
			cellText[0] = []string{
				"Node",
				"Path",
				"Read",
				"Write",
			}
			trailerIfGreaterThan := func(in string, max int) string {
				if len(in) < max {
					return ""
				}
				return "..."
			}
			for i, item := range result.DrivePerf {
				cellText[i+1] = []string{
					fmt.Sprintf("%.64s%s", result.Endpoint,
						trailerIfGreaterThan(result.Endpoint, 64)),
					fmt.Sprintf("%.64s%s", item.Path,
						trailerIfGreaterThan(result.Endpoint, 64)),
					humanize.IBytes(item.ReadThroughput) + "/s",
					humanize.IBytes(item.WriteThroughput) + "/s",
				}
			}
			if len(result.DrivePerf) > 0 {
				tbl.DisplayTable(cellText)
			}
			if result.Error != "" {
				fmt.Println(color.New(color.FgRed, color.Bold).Sprintf("ERROR"), result.Error)
			}
		}
	}

	return nil
}
<|endoftext|>