<commit_before>/* 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 http://mozilla.org/MPL/2.0/. */

package main

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

	"github.com/labstack/echo"
)

// Datacenter holds the datacenter response from datacenter-store
type Datacenter struct {
	ID              int    `json:"id"`
	GroupID         int    `json:"group_id"`
	Name            string `json:"name"`
	Type            string `json:"type"`
	Region          string `json:"region"`
	Username        string `json:"username"`
	Password        string `json:"password"`
	VCloudURL       string `json:"vcloud_url"`
	VseURL          string `json:"vse_url"`
	ExternalNetwork string `json:"external_network"`
}

// Validate the datacenter
func (d *Datacenter) Validate() error {
	if d.Name == "" {
		return errors.New("Datacenter name is empty")
	}

	if d.GroupID == 0 {
		return errors.New("Datacenter group is empty")
	}

	if d.Type == "" {
		return errors.New("Datacenter type is empty")
	}

	if d.Username == "" {
		return errors.New("Datacenter username is empty")
	}

	if d.Password == "" {
		return errors.New("Datacenter password is empty")
	}

	if d.Type == "vcloud" && d.VCloudURL == "" {
		return errors.New("Datacenter vcloud url is empty")
	}

	return nil
}

// Map : maps a datacenter from a request's body and validates the input
func (d *Datacenter) Map(c echo.Context) *echo.HTTPError {
	body := c.Request().Body()
	data, err := ioutil.ReadAll(body)
	if err != nil {
		return ErrBadReqBody
	}

	err = json.Unmarshal(data, &d)
	if err != nil {
		return ErrBadReqBody
	}

	err = d.Validate()
	if err != nil {
		return ErrBadReqBody
	}

	return nil
}

// getDatacentersHandler : get all datacenters
func getDatacentersHandler(c echo.Context) error {
	au := authenticatedUser(c)
	query := fmt.Sprintf(`{"group_id": %d}`, au.GroupID)
	msg, err := n.Request("datacenter.find", []byte(query), 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// getDatacenterHandler : get a datancenter by id
func getDatacenterHandler(c echo.Context) error {
	var query string
	au := authenticatedUser(c)

	if au.Admin {
		query = fmt.Sprintf(`{"id": %s}`, c.Param("datacenter"))
	} else {
		query = fmt.Sprintf(`{"id": %s, "group_id": %d}`, c.Param("datacenter"), au.GroupID)
	}

	msg, err := n.Request("datacenter.get", []byte(query), 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// createDatacenterHandler : Endpoint to create a datacenter
func createDatacenterHandler(c echo.Context) error {
	var d Datacenter
	if d.Map(c) != nil {
		return ErrBadReqBody
	}

	au := authenticatedUser(c)
	if au.Admin != true && d.GroupID != au.GroupID {
		return ErrUnauthorized
	}

	data, err := json.Marshal(d)
	if err != nil {
		return ErrInternal
	}

	msg, err := n.Request("datacenter.set", data, 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// updateDatacenterHandler : Updates a datacenter through its store
func updateDatacenterHandler(c echo.Context) error {
	var d Datacenter
	if d.Map(c) != nil {
		return ErrBadReqBody
	}

	au := authenticatedUser(c)
	if au.Admin != true || d.GroupID != au.GroupID {
		return ErrUnauthorized
	}

	data, err := json.Marshal(d)
	if err != nil {
		return ErrInternal
	}

	msg, err := n.Request("datacenter.set", data, 5*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// deleteDatacenterHandler : Deletes a datancenter though its store
func deleteDatacenterHandler(c echo.Context) error {
	au := authenticatedUser(c)

	query := fmt.Sprintf(`{"id": %s, "group_id": %d}`, c.Param("datacenter"), au.GroupID)
	msg, err := n.Request("datacenter.del", []byte(query), 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.String(http.StatusOK, "")
}
<commit_msg>Removed validation on GroupID to datacenters<commit_after>/* 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 http://mozilla.org/MPL/2.0/. */

package main

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

	"github.com/labstack/echo"
)

// Datacenter holds the datacenter response from datacenter-store
type Datacenter struct {
	ID              int    `json:"id"`
	GroupID         int    `json:"group_id"`
	Name            string `json:"name"`
	Type            string `json:"type"`
	Region          string `json:"region"`
	Username        string `json:"username"`
	Password        string `json:"password"`
	VCloudURL       string `json:"vcloud_url"`
	VseURL          string `json:"vse_url"`
	ExternalNetwork string `json:"external_network"`
}

// Validate the datacenter
func (d *Datacenter) Validate() error {
	if d.Name == "" {
		return errors.New("Datacenter name is empty")
	}

	if d.Type == "" {
		return errors.New("Datacenter type is empty")
	}

	if d.Username == "" {
		return errors.New("Datacenter username is empty")
	}

	if d.Password == "" {
		return errors.New("Datacenter password is empty")
	}

	if d.Type == "vcloud" && d.VCloudURL == "" {
		return errors.New("Datacenter vcloud url is empty")
	}

	return nil
}

// Map : maps a datacenter from a request's body and validates the input
func (d *Datacenter) Map(c echo.Context) *echo.HTTPError {
	body := c.Request().Body()
	data, err := ioutil.ReadAll(body)
	if err != nil {
		return ErrBadReqBody
	}

	err = json.Unmarshal(data, &d)
	if err != nil {
		return ErrBadReqBody
	}

	err = d.Validate()
	if err != nil {
		return ErrBadReqBody
	}

	return nil
}

// getDatacentersHandler : get all datacenters
func getDatacentersHandler(c echo.Context) error {
	au := authenticatedUser(c)
	query := fmt.Sprintf(`{"group_id": %d}`, au.GroupID)
	msg, err := n.Request("datacenter.find", []byte(query), 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// getDatacenterHandler : get a datancenter by id
func getDatacenterHandler(c echo.Context) error {
	var query string
	au := authenticatedUser(c)

	if au.Admin {
		query = fmt.Sprintf(`{"id": %s}`, c.Param("datacenter"))
	} else {
		query = fmt.Sprintf(`{"id": %s, "group_id": %d}`, c.Param("datacenter"), au.GroupID)
	}

	msg, err := n.Request("datacenter.get", []byte(query), 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// createDatacenterHandler : Endpoint to create a datacenter
func createDatacenterHandler(c echo.Context) error {
	var d Datacenter
	if d.Map(c) != nil {
		return ErrBadReqBody
	}

	au := authenticatedUser(c)
	if au.Admin != true && d.GroupID != au.GroupID {
		return ErrUnauthorized
	}

	data, err := json.Marshal(d)
	if err != nil {
		return ErrInternal
	}

	msg, err := n.Request("datacenter.set", data, 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// updateDatacenterHandler : Updates a datacenter through its store
func updateDatacenterHandler(c echo.Context) error {
	var d Datacenter
	if d.Map(c) != nil {
		return ErrBadReqBody
	}

	au := authenticatedUser(c)
	if au.Admin != true || d.GroupID != au.GroupID {
		return ErrUnauthorized
	}

	data, err := json.Marshal(d)
	if err != nil {
		return ErrInternal
	}

	msg, err := n.Request("datacenter.set", data, 5*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.JSONBlob(http.StatusOK, msg.Data)
}

// deleteDatacenterHandler : Deletes a datancenter though its store
func deleteDatacenterHandler(c echo.Context) error {
	au := authenticatedUser(c)

	query := fmt.Sprintf(`{"id": %s, "group_id": %d}`, c.Param("datacenter"), au.GroupID)
	msg, err := n.Request("datacenter.del", []byte(query), 1*time.Second)
	if err != nil {
		return ErrGatewayTimeout
	}

	if re := responseErr(msg); re != nil {
		return re.HTTPError
	}

	return c.String(http.StatusOK, "")
}
<|endoftext|>


<commit_before>package taskrunner

import (
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
	ti "github.com/hashicorp/nomad/client/allocrunner/taskrunner/interfaces"
	"github.com/hashicorp/nomad/client/consul"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/pkg/errors"
)

const (
	// the name of this hook, used in logs
	sidsHookName = "consul_si_token"

	// sidsBackoffBaseline is the baseline time for exponential backoff when
	// attempting to retrieve a Consul SI token
	sidsBackoffBaseline = 5 * time.Second

	// sidsBackoffLimit is the limit of the exponential backoff when attempting
	// to retrieve a Consul SI token
	sidsBackoffLimit = 3 * time.Minute

	// sidsDerivationTimeout limits the amount of time we may spend trying to
	// derive a SI token. If the hook does not get a token within this amount of
	// time, the result is a failure.
	sidsDerivationTimeout = 5 * time.Minute

	// sidsTokenFile is the name of the file holding the Consul SI token inside
	// the task's secret directory
	sidsTokenFile = "si_token"

	// sidsTokenFilePerms is the level of file permissions granted on the file
	// in the secrets directory for the task
	sidsTokenFilePerms = 0440
)

type sidsHookConfig struct {
	alloc      *structs.Allocation
	task       *structs.Task
	sidsClient consul.ServiceIdentityAPI
	lifecycle  ti.TaskLifecycle
	logger     hclog.Logger
}

// Service Identities hook for managing SI tokens of connect enabled tasks.
type sidsHook struct {
	// alloc is the allocation
	alloc *structs.Allocation

	// taskName is the name of the task
	task *structs.Task

	// sidsClient is the Consul client [proxy] for requesting SI tokens
	sidsClient consul.ServiceIdentityAPI

	// lifecycle is used to signal, restart, and kill a task
	lifecycle ti.TaskLifecycle

	// derivationTimeout is the amount of time we may wait for Consul to successfully
	// provide a SI token. Making this configurable for testing, otherwise
	// default to sidsDerivationTimeout
	derivationTimeout time.Duration

	// logger is used to log
	logger hclog.Logger

	// lock variables that can be manipulated after hook creation
	lock sync.Mutex
	// firstRun keeps track of whether the hook is being called for the first
	// time (for this task) during the lifespan of the Nomad Client process.
	firstRun bool
}

func newSIDSHook(c sidsHookConfig) *sidsHook {
	return &sidsHook{
		alloc:             c.alloc,
		task:              c.task,
		sidsClient:        c.sidsClient,
		lifecycle:         c.lifecycle,
		derivationTimeout: sidsDerivationTimeout,
		logger:            c.logger.Named(sidsHookName),
		firstRun:          true,
	}
}

func (h *sidsHook) Name() string {
	return sidsHookName
}

func (h *sidsHook) Prestart(
	ctx context.Context,
	req *interfaces.TaskPrestartRequest,
	resp *interfaces.TaskPrestartResponse) error {

	h.lock.Lock()
	defer h.lock.Unlock()

	// do nothing if we have already done things
	if h.earlyExit() {
		resp.Done = true
		return nil
	}

	// optimistically try to recover token from disk
	token, err := h.recoverToken(req.TaskDir.SecretsDir)
	if err != nil {
		return err
	}

	// need to ask for a new SI token & persist it to disk
	if token == "" {
		if token, err = h.deriveSIToken(ctx); err != nil {
			return err
		}
		if err := h.writeToken(req.TaskDir.SecretsDir, token); err != nil {
			return err
		}
	}

	h.logger.Info("derived SI token", "task", h.task.Name, "si_task", h.task.Kind.Value())

	resp.Done = true
	return nil
}

// earlyExit returns true if the Prestart hook has already been executed during
// the instantiation of this task runner.
//
// assumes h is locked
func (h *sidsHook) earlyExit() bool {
	if h.firstRun {
		h.firstRun = false
		return false
	}
	return true
}

// writeToken writes token into the secrets directory for the task.
func (h *sidsHook) writeToken(dir string, token string) error {
	tokenPath := filepath.Join(dir, sidsTokenFile)
	if err := ioutil.WriteFile(tokenPath, []byte(token), sidsTokenFilePerms); err != nil {
		return errors.Wrap(err, "failed to write SI token")
	}
	return nil
}

// recoverToken returns the token saved to disk in the secrets directory for the
// task if it exists, or the empty string if the file does not exist. an error
// is returned only for some other (e.g. disk IO) error.
func (h *sidsHook) recoverToken(dir string) (string, error) {
	tokenPath := filepath.Join(dir, sidsTokenFile)
	token, err := ioutil.ReadFile(tokenPath)
	if err != nil {
		if !os.IsNotExist(err) {
			h.logger.Error("failed to recover SI token", "error", err)
			return "", errors.Wrap(err, "failed to recover SI token")
		}
		h.logger.Trace("no pre-existing SI token to recover", "task", h.task.Name)
		return "", nil // token file does not exist yet
	}
	h.logger.Trace("recovered pre-existing SI token", "task", h.task.Name)
	return string(token), nil
}

// siDerivationResult is used to pass along the result of attempting to derive
// an SI token between the goroutine doing the derivation and its caller
type siDerivationResult struct {
	token string
	err   error
}

// deriveSIToken spawns and waits on a goroutine which will make attempts to
// derive an SI token until a token is successfully created, or ctx is signaled
// done.
func (h *sidsHook) deriveSIToken(ctx context.Context) (string, error) {
	ctx2, cancel := context.WithTimeout(ctx, h.derivationTimeout)
	defer cancel()

	resultCh := make(chan siDerivationResult)

	// keep trying to get the token in the background
	go h.tryDerive(ctx2, resultCh)

	// wait until we get a token, or we get a signal to quit
	for {
		select {
		case result := <-resultCh:
			if result.err != nil {
				h.kill(ctx, errors.Wrap(result.err, "consul: failed to derive SI token"))
				return "", result.err
			}
			return result.token, nil
		case <-ctx2.Done():
			return "", ctx2.Err()
		}
	}
}

func (h *sidsHook) kill(ctx context.Context, err error) {
	_ = h.lifecycle.Kill(
		ctx,
		structs.NewTaskEvent(structs.TaskKilling).
			SetFailsTask().
			SetDisplayMessage(err.Error()),
	)
}

// tryDerive loops forever until a token is created, or ctx is done.
func (h *sidsHook) tryDerive(ctx context.Context, ch chan<- siDerivationResult) {
	for attempt := 0; backoff(ctx, attempt); attempt++ {

		tokens, err := h.sidsClient.DeriveSITokens(h.alloc, []string{h.task.Name})

		switch {
		case err == nil:
			token, exists := tokens[h.task.Name]
			if !exists {
				err := errors.New("response does not include token for task")
				h.logger.Error("derive SI token is missing token for task", "error", err, "task", h.task.Name)
				ch <- siDerivationResult{token: "", err: err}
				return
			}
			ch <- siDerivationResult{token: token, err: nil}
			return
		case structs.IsServerSide(err):
			// the error is known to be a server problem, just die
			h.logger.Error("failed to derive SI token", "error", err, "task", h.task.Name, "server_side", true)
			ch <- siDerivationResult{token: "", err: err}
			return
		case !structs.IsRecoverable(err):
			// the error is known not to be recoverable, just die
			h.logger.Error("failed to derive SI token", "error", err, "task", h.task.Name, "recoverable", false)
			ch <- siDerivationResult{token: "", err: err}
			return

		default:
			// the error is marked recoverable, retry after some backoff
			h.logger.Error("failed attempt to derive SI token", "error", err, "recoverable", true)
		}
	}
}

func backoff(ctx context.Context, attempt int) bool {
	next := computeBackoff(attempt)
	select {
	case <-ctx.Done():
		return false
	case <-time.After(next):
		return true
	}
}

func computeBackoff(attempt int) time.Duration {
	switch attempt {
	case 0:
		return 0
	case 1:
		// go fast on first retry, because a unit test should be fast
		return 100 * time.Millisecond
	default:
		wait := time.Duration(attempt) * sidsBackoffBaseline
		if wait > sidsBackoffLimit {
			wait = sidsBackoffLimit
		}
		return wait
	}
}
<commit_msg>client: PR cleanup - shadow context variable<commit_after>package taskrunner

import (
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
	ti "github.com/hashicorp/nomad/client/allocrunner/taskrunner/interfaces"
	"github.com/hashicorp/nomad/client/consul"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/pkg/errors"
)

const (
	// the name of this hook, used in logs
	sidsHookName = "consul_si_token"

	// sidsBackoffBaseline is the baseline time for exponential backoff when
	// attempting to retrieve a Consul SI token
	sidsBackoffBaseline = 5 * time.Second

	// sidsBackoffLimit is the limit of the exponential backoff when attempting
	// to retrieve a Consul SI token
	sidsBackoffLimit = 3 * time.Minute

	// sidsDerivationTimeout limits the amount of time we may spend trying to
	// derive a SI token. If the hook does not get a token within this amount of
	// time, the result is a failure.
	sidsDerivationTimeout = 5 * time.Minute

	// sidsTokenFile is the name of the file holding the Consul SI token inside
	// the task's secret directory
	sidsTokenFile = "si_token"

	// sidsTokenFilePerms is the level of file permissions granted on the file
	// in the secrets directory for the task
	sidsTokenFilePerms = 0440
)

type sidsHookConfig struct {
	alloc      *structs.Allocation
	task       *structs.Task
	sidsClient consul.ServiceIdentityAPI
	lifecycle  ti.TaskLifecycle
	logger     hclog.Logger
}

// Service Identities hook for managing SI tokens of connect enabled tasks.
type sidsHook struct {
	// alloc is the allocation
	alloc *structs.Allocation

	// taskName is the name of the task
	task *structs.Task

	// sidsClient is the Consul client [proxy] for requesting SI tokens
	sidsClient consul.ServiceIdentityAPI

	// lifecycle is used to signal, restart, and kill a task
	lifecycle ti.TaskLifecycle

	// derivationTimeout is the amount of time we may wait for Consul to successfully
	// provide a SI token. Making this configurable for testing, otherwise
	// default to sidsDerivationTimeout
	derivationTimeout time.Duration

	// logger is used to log
	logger hclog.Logger

	// lock variables that can be manipulated after hook creation
	lock sync.Mutex
	// firstRun keeps track of whether the hook is being called for the first
	// time (for this task) during the lifespan of the Nomad Client process.
	firstRun bool
}

func newSIDSHook(c sidsHookConfig) *sidsHook {
	return &sidsHook{
		alloc:             c.alloc,
		task:              c.task,
		sidsClient:        c.sidsClient,
		lifecycle:         c.lifecycle,
		derivationTimeout: sidsDerivationTimeout,
		logger:            c.logger.Named(sidsHookName),
		firstRun:          true,
	}
}

func (h *sidsHook) Name() string {
	return sidsHookName
}

func (h *sidsHook) Prestart(
	ctx context.Context,
	req *interfaces.TaskPrestartRequest,
	resp *interfaces.TaskPrestartResponse) error {

	h.lock.Lock()
	defer h.lock.Unlock()

	// do nothing if we have already done things
	if h.earlyExit() {
		resp.Done = true
		return nil
	}

	// optimistically try to recover token from disk
	token, err := h.recoverToken(req.TaskDir.SecretsDir)
	if err != nil {
		return err
	}

	// need to ask for a new SI token & persist it to disk
	if token == "" {
		if token, err = h.deriveSIToken(ctx); err != nil {
			return err
		}
		if err := h.writeToken(req.TaskDir.SecretsDir, token); err != nil {
			return err
		}
	}

	h.logger.Info("derived SI token", "task", h.task.Name, "si_task", h.task.Kind.Value())

	resp.Done = true
	return nil
}

// earlyExit returns true if the Prestart hook has already been executed during
// the instantiation of this task runner.
//
// assumes h is locked
func (h *sidsHook) earlyExit() bool {
	if h.firstRun {
		h.firstRun = false
		return false
	}
	return true
}

// writeToken writes token into the secrets directory for the task.
func (h *sidsHook) writeToken(dir string, token string) error {
	tokenPath := filepath.Join(dir, sidsTokenFile)
	if err := ioutil.WriteFile(tokenPath, []byte(token), sidsTokenFilePerms); err != nil {
		return errors.Wrap(err, "failed to write SI token")
	}
	return nil
}

// recoverToken returns the token saved to disk in the secrets directory for the
// task if it exists, or the empty string if the file does not exist. an error
// is returned only for some other (e.g. disk IO) error.
func (h *sidsHook) recoverToken(dir string) (string, error) {
	tokenPath := filepath.Join(dir, sidsTokenFile)
	token, err := ioutil.ReadFile(tokenPath)
	if err != nil {
		if !os.IsNotExist(err) {
			h.logger.Error("failed to recover SI token", "error", err)
			return "", errors.Wrap(err, "failed to recover SI token")
		}
		h.logger.Trace("no pre-existing SI token to recover", "task", h.task.Name)
		return "", nil // token file does not exist yet
	}
	h.logger.Trace("recovered pre-existing SI token", "task", h.task.Name)
	return string(token), nil
}

// siDerivationResult is used to pass along the result of attempting to derive
// an SI token between the goroutine doing the derivation and its caller
type siDerivationResult struct {
	token string
	err   error
}

// deriveSIToken spawns and waits on a goroutine which will make attempts to
// derive an SI token until a token is successfully created, or ctx is signaled
// done.
func (h *sidsHook) deriveSIToken(ctx context.Context) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, h.derivationTimeout)
	defer cancel()

	resultCh := make(chan siDerivationResult)

	// keep trying to get the token in the background
	go h.tryDerive(ctx, resultCh)

	// wait until we get a token, or we get a signal to quit
	for {
		select {
		case result := <-resultCh:
			if result.err != nil {
				h.kill(ctx, errors.Wrap(result.err, "consul: failed to derive SI token"))
				return "", result.err
			}
			return result.token, nil
		case <-ctx.Done():
			return "", ctx.Err()
		}
	}
}

func (h *sidsHook) kill(ctx context.Context, err error) {
	_ = h.lifecycle.Kill(
		ctx,
		structs.NewTaskEvent(structs.TaskKilling).
			SetFailsTask().
			SetDisplayMessage(err.Error()),
	)
}

// tryDerive loops forever until a token is created, or ctx is done.
func (h *sidsHook) tryDerive(ctx context.Context, ch chan<- siDerivationResult) {
	for attempt := 0; backoff(ctx, attempt); attempt++ {

		tokens, err := h.sidsClient.DeriveSITokens(h.alloc, []string{h.task.Name})

		switch {
		case err == nil:
			token, exists := tokens[h.task.Name]
			if !exists {
				err := errors.New("response does not include token for task")
				h.logger.Error("derive SI token is missing token for task", "error", err, "task", h.task.Name)
				ch <- siDerivationResult{token: "", err: err}
				return
			}
			ch <- siDerivationResult{token: token, err: nil}
			return
		case structs.IsServerSide(err):
			// the error is known to be a server problem, just die
			h.logger.Error("failed to derive SI token", "error", err, "task", h.task.Name, "server_side", true)
			ch <- siDerivationResult{token: "", err: err}
			return
		case !structs.IsRecoverable(err):
			// the error is known not to be recoverable, just die
			h.logger.Error("failed to derive SI token", "error", err, "task", h.task.Name, "recoverable", false)
			ch <- siDerivationResult{token: "", err: err}
			return

		default:
			// the error is marked recoverable, retry after some backoff
			h.logger.Error("failed attempt to derive SI token", "error", err, "recoverable", true)
		}
	}
}

func backoff(ctx context.Context, attempt int) bool {
	next := computeBackoff(attempt)
	select {
	case <-ctx.Done():
		return false
	case <-time.After(next):
		return true
	}
}

func computeBackoff(attempt int) time.Duration {
	switch attempt {
	case 0:
		return 0
	case 1:
		// go fast on first retry, because a unit test should be fast
		return 100 * time.Millisecond
	default:
		wait := time.Duration(attempt) * sidsBackoffBaseline
		if wait > sidsBackoffLimit {
			wait = sidsBackoffLimit
		}
		return wait
	}
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"
	"sync"

	usbdrivedetecter "github.com/deepakjois/gousbdrivedetector"
	moocfetcher "github.com/moocfetcher/moocfetcher-appliance/backend/lib"
)

var copyStatusPath = regexp.MustCompile("^/api/copy-status/([a-zA-Z0-9\\-]+)$")
var copyCancelPath = regexp.MustCompile("^/api/copy/([a-zA-Z0-9\\-]+)$")

// MOOCFetcherApplianceServer take request from the MOOCFetcher Appliance frontend app
// and performs copy operations, provides status updates etc.
type MOOCFetcherApplianceServer struct {
	*http.ServeMux
	courseFoldersPath string
	courseMetadata    moocfetcher.CourseData
	stats             map[string]int
	jobs              map[string]*CopyJob
	currJobId         string
	Copier            CourseCopier
	sync.Mutex
}

func NewServer(courseFoldersPath string, courseMetadata moocfetcher.CourseData) *MOOCFetcherApplianceServer {
	m := &MOOCFetcherApplianceServer{
		ServeMux:          http.NewServeMux(),
		courseFoldersPath: courseFoldersPath,
		courseMetadata:    courseMetadata,
		stats:             make(map[string]int),
		jobs:              make(map[string]*CopyJob),
	}

	m.readStats()

	m.Handle("/api/copy", http.HandlerFunc(m.copyHandler))
	m.Handle("/api/copy/", http.HandlerFunc(m.copyHandler))
	m.Handle("/api/copy-status/", http.HandlerFunc(m.progressHandler))
	m.Handle("/api/stats", http.HandlerFunc(m.statsHandler))

	return m
}

func (s *MOOCFetcherApplianceServer) copyHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "POST":
		s.copyStartHandler(w, r)
		return
	case "PUT":
		s.copyCancelHandler(w, r)
		return
	default:
		http.NotFound(w, r)
		return
	}

}

func (s *MOOCFetcherApplianceServer) copyStartHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()
	// Get request JSON and parse
	var courseData moocfetcher.CourseData

	defer r.Body.Close()
	err := json.NewDecoder(r.Body).Decode(&courseData)
	if err != nil {
		httpJSONError(w, err.Error(), http.StatusBadRequest)
		return
	}

	if s.currJobId != "" {
		httpJSONError(w, fmt.Sprintf("Job %s already running.", s.currJobId), http.StatusForbidden)
		return
	}

	for _, c := range courseData.Courses {
		s.stats[c.Slug]++
	}

	// Detect USB media
	drives, err := usbdrivedetecter.Detect()
	if err != nil {
		httpJSONError(w, fmt.Sprintf("Error trying to detect USB: %s", err.Error()), http.StatusInternalServerError)
		return
	}

	if len(drives) == 0 {
		httpJSONError(w, "Could not locate USB Pen Drive", http.StatusInternalServerError)
		return
	}

	// TODO write code to select most likely candidate
	drivePath := drives[len(drives)-1]

	if !checkFreeSpace(drivePath, calcSpaceRequired(courseData)) {
		httpJSONError(w, fmt.Sprintf("Not enough free space on USB media %s", drivePath), http.StatusInternalServerError)
		return
	}

	copier := s.Copier
	if copier == nil {
		copier = NewFileSystemCopier(s.courseFoldersPath, drivePath)
	}

	job := NewCopyJob(courseData, copier)
	s.currJobId = job.ID
	s.jobs[job.ID] = job

	resp := fmt.Sprintf("{ \"id\": \"%s\"}", job.ID)
	w.Write([]byte(resp))
	w.Header().Set("Content-Type", "application/json")

	// Track job
	go func() {
		done := job.Done
		go job.Run()
		<-done
		s.writeStats()
		s.Lock()
		defer s.Unlock()
		s.currJobId = ""
	}()
}

func (s *MOOCFetcherApplianceServer) copyCancelHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	var job *CopyJob
	// Check the job ID
	if m := copyCancelPath.FindStringSubmatch(r.URL.Path); m == nil || m[1] != s.currJobId {
		http.Error(w, "404 page not found", http.StatusNotFound)
		return
	} else {
		job = s.jobs[m[1]]
	}
	if job.status == "running" {
		go job.Cancel()
	}
}

func (s *MOOCFetcherApplianceServer) progressHandler(w http.ResponseWriter, r *http.Request) {
	// Get the job ID
	m := copyStatusPath.FindStringSubmatch(r.URL.Path)
	if m == nil {
		http.Error(w, "404 page not found", http.StatusNotFound)
		return
	}

	id := m[1]
	job, ok := s.jobs[id]
	if ok {
		progress := job.Progress()
		js, err := json.Marshal(progress)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(js)
		return
	}
	http.NotFound(w, r)
}

func (s *MOOCFetcherApplianceServer) statsHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Under Construction", http.StatusNotImplemented)
}

func (s *MOOCFetcherApplianceServer) readStats() {
	file, err := ioutil.ReadFile("./stats.json")
	if err != nil {
		log.Printf("Unable to read stats file: %s\n", err)
		return
	}

	// Read stats data from file
	var stats map[string]int
	err = json.Unmarshal(file, &stats)
	if err != nil {
		log.Printf("Error reading stats file %s\n", err)
		return
	}
	s.stats = stats
}

func (s *MOOCFetcherApplianceServer) writeStats() {
	stats, err := json.Marshal(s.stats)
	if err != nil {
		log.Printf("Error writing stats file: %s\n", err)
		return
	}
	ioutil.WriteFile("./stats.json", stats, 0644)
}
<commit_msg>Add hacks to improve USB drive detection<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"
	"sync"

	usbdrivedetecter "github.com/deepakjois/gousbdrivedetector"
	moocfetcher "github.com/moocfetcher/moocfetcher-appliance/backend/lib"
)

var copyStatusPath = regexp.MustCompile("^/api/copy-status/([a-zA-Z0-9\\-]+)$")
var copyCancelPath = regexp.MustCompile("^/api/copy/([a-zA-Z0-9\\-]+)$")

// MOOCFetcherApplianceServer take request from the MOOCFetcher Appliance frontend app
// and performs copy operations, provides status updates etc.
type MOOCFetcherApplianceServer struct {
	*http.ServeMux
	courseFoldersPath string
	courseMetadata    moocfetcher.CourseData
	stats             map[string]int
	jobs              map[string]*CopyJob
	currJobId         string
	Copier            CourseCopier
	sync.Mutex
}

func NewServer(courseFoldersPath string, courseMetadata moocfetcher.CourseData) *MOOCFetcherApplianceServer {
	m := &MOOCFetcherApplianceServer{
		ServeMux:          http.NewServeMux(),
		courseFoldersPath: courseFoldersPath,
		courseMetadata:    courseMetadata,
		stats:             make(map[string]int),
		jobs:              make(map[string]*CopyJob),
	}

	m.readStats()

	m.Handle("/api/copy", http.HandlerFunc(m.copyHandler))
	m.Handle("/api/copy/", http.HandlerFunc(m.copyHandler))
	m.Handle("/api/copy-status/", http.HandlerFunc(m.progressHandler))
	m.Handle("/api/stats", http.HandlerFunc(m.statsHandler))

	return m
}

func (s *MOOCFetcherApplianceServer) copyHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "POST":
		s.copyStartHandler(w, r)
		return
	case "PUT":
		s.copyCancelHandler(w, r)
		return
	default:
		http.NotFound(w, r)
		return
	}

}

func (s *MOOCFetcherApplianceServer) copyStartHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()
	// Get request JSON and parse
	var courseData moocfetcher.CourseData

	defer r.Body.Close()
	err := json.NewDecoder(r.Body).Decode(&courseData)
	if err != nil {
		httpJSONError(w, err.Error(), http.StatusBadRequest)
		return
	}

	if s.currJobId != "" {
		httpJSONError(w, fmt.Sprintf("Job %s already running.", s.currJobId), http.StatusForbidden)
		return
	}

	for _, c := range courseData.Courses {
		s.stats[c.Slug]++
	}

	// Detect USB media
	drives, err := usbdrivedetecter.Detect()
	if err != nil {
		httpJSONError(w, fmt.Sprintf("Error trying to detect USB: %s", err.Error()), http.StatusInternalServerError)
		return
	}

	// Filter out drive containing courses (if required)
	for i := len(drives) - 1; i >= 0; i-- {
		drive := drives[i]
		if strings.Contains(s.courseFoldersPath, drive) {
			drives = append(drives[:i], drives[i+1:]...)
		}
	}

	// FIXME Ugly hack for Windows, hardcoded for our setup.
	// Try E:\
	if len(drives) == 0 && runtime.GOOS == "windows" {
		log.Println("Last ditch effort on Windows to detect USB Drive")
		_, err := os.Open("E:\\")
		if err != nil {
			drives = []string{"E:\\"}
		}
	}

	if len(drives) == 0 {
		httpJSONError(w, "Could not locate USB Pen Drive", http.StatusInternalServerError)
		return
	}

	// TODO write code to select most likely candidate
	drivePath := drives[len(drives)-1]

	if !checkFreeSpace(drivePath, calcSpaceRequired(courseData)) {
		httpJSONError(w, fmt.Sprintf("Not enough free space on USB media %s", drivePath), http.StatusInternalServerError)
		return
	}

	copier := s.Copier
	if copier == nil {
		copier = NewFileSystemCopier(s.courseFoldersPath, drivePath)
	}

	job := NewCopyJob(courseData, copier)
	s.currJobId = job.ID
	s.jobs[job.ID] = job

	resp := fmt.Sprintf("{ \"id\": \"%s\"}", job.ID)
	w.Write([]byte(resp))
	w.Header().Set("Content-Type", "application/json")

	// Track job
	go func() {
		done := job.Done
		go job.Run()
		<-done
		s.writeStats()
		s.Lock()
		defer s.Unlock()
		s.currJobId = ""
	}()
}

func (s *MOOCFetcherApplianceServer) copyCancelHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	var job *CopyJob
	// Check the job ID
	if m := copyCancelPath.FindStringSubmatch(r.URL.Path); m == nil || m[1] != s.currJobId {
		http.Error(w, "404 page not found", http.StatusNotFound)
		return
	} else {
		job = s.jobs[m[1]]
	}
	if job.status == "running" {
		go job.Cancel()
	}
}

func (s *MOOCFetcherApplianceServer) progressHandler(w http.ResponseWriter, r *http.Request) {
	// Get the job ID
	m := copyStatusPath.FindStringSubmatch(r.URL.Path)
	if m == nil {
		http.Error(w, "404 page not found", http.StatusNotFound)
		return
	}

	id := m[1]
	job, ok := s.jobs[id]
	if ok {
		progress := job.Progress()
		js, err := json.Marshal(progress)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(js)
		return
	}
	http.NotFound(w, r)
}

func (s *MOOCFetcherApplianceServer) statsHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Under Construction", http.StatusNotImplemented)
}

func (s *MOOCFetcherApplianceServer) readStats() {
	file, err := ioutil.ReadFile("./stats.json")
	if err != nil {
		log.Printf("Unable to read stats file: %s\n", err)
		return
	}

	// Read stats data from file
	var stats map[string]int
	err = json.Unmarshal(file, &stats)
	if err != nil {
		log.Printf("Error reading stats file %s\n", err)
		return
	}
	s.stats = stats
}

func (s *MOOCFetcherApplianceServer) writeStats() {
	stats, err := json.Marshal(s.stats)
	if err != nil {
		log.Printf("Error writing stats file: %s\n", err)
		return
	}
	ioutil.WriteFile("./stats.json", stats, 0644)
}
<|endoftext|>


<commit_before>package viewer

import (
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
)

var (
	pingWait  = 70 * time.Second
	writeWait = 1 * time.Second
	// Maximum message size allowed from peer.
	maxMessageSize int64 = 512
)

type WsClient struct {
	out       io.Writer
	conn      *websocket.Conn
	interrupt chan struct{}
	done      chan struct{}
}

func WsConnect(addr string, out io.Writer, interrupt chan struct{}) {
	u := url.URL{Scheme: "ws", Host: addr, Path: "/"}

	conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Printf("WsConnect: %v", err)
		return
	}
	client := &WsClient{
		out:       out,
		conn:      conn,
		done:      make(chan struct{}),
		interrupt: interrupt,
	}

	go client.pull()
	client.push()
}

func (c *WsClient) push() {
	defer func() {
		c.conn.Close()
	}()

	select {
	case <-c.done:
	case <-c.interrupt:
		log.Println("interrupted")
		// To cleanly close a connection, a client should send a close
		// frame and wait for the server to close the connection.
		err := c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err != nil {
			log.Println("write close:", err)
			return
		}
		select {
		case <-c.done:
		case <-time.After(time.Second):
		}
		return
	}
}

func (c *WsClient) pull() {
	defer func() {
		c.conn.Close()
		close(c.done)
	}()

	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetPingHandler(func(message string) error {
		c.conn.SetReadDeadline(time.Now().Add(pingWait))
		err := c.conn.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(writeWait))
		if err == websocket.ErrCloseSent {
			return nil
		} else if e, ok := err.(net.Error); ok && e.Temporary() {
			return nil
		}
		return err
	})

	for {
		msgType, r, err := c.conn.NextReader()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("error: %v", err)
			}
			return
		}

		if msgType == websocket.TextMessage {
			msg, err := ioutil.ReadAll(r)
			if err != nil {
				log.Printf("error: %v", err)
				return
			}
			c.out.Write([]byte("\r"))
			c.out.Write(msg)
			//if _, err := io.Copy(c.out, r); err != nil {
			//	log.Printf("error: %v", err)
			//	return
			//}
		}
		//if err := c.out.Close(); err != nil {
		//	log.Printf("error: %v", err)
		//	return
		//}
	}
}
<commit_msg>viewer: fix CR output for ticker<commit_after>package viewer

import (
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
)

var (
	pingWait  = 70 * time.Second
	writeWait = 1 * time.Second
	// Maximum message size allowed from peer.
	maxMessageSize int64 = 512
)

type WsClient struct {
	out       io.Writer
	conn      *websocket.Conn
	interrupt chan struct{}
	done      chan struct{}
}

func WsConnect(addr string, out io.Writer, interrupt chan struct{}) {
	u := url.URL{Scheme: "ws", Host: addr, Path: "/"}

	conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Printf("WsConnect: %v", err)
		return
	}
	client := &WsClient{
		out:       out,
		conn:      conn,
		done:      make(chan struct{}),
		interrupt: interrupt,
	}

	go client.pull()
	client.push()
}

func (c *WsClient) push() {
	defer func() {
		c.conn.Close()
	}()

	select {
	case <-c.done:
	case <-c.interrupt:
		log.Println("interrupted")
		// To cleanly close a connection, a client should send a close
		// frame and wait for the server to close the connection.
		err := c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err != nil {
			log.Println("write close:", err)
			return
		}
		select {
		case <-c.done:
		case <-time.After(time.Second):
		}
		return
	}
}

func (c *WsClient) pull() {
	defer func() {
		c.conn.Close()
		close(c.done)
	}()

	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetPingHandler(func(message string) error {
		c.conn.SetReadDeadline(time.Now().Add(pingWait))
		err := c.conn.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(writeWait))
		if err == websocket.ErrCloseSent {
			return nil
		} else if e, ok := err.(net.Error); ok && e.Temporary() {
			return nil
		}
		return err
	})

	for {
		msgType, r, err := c.conn.NextReader()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("error: %v", err)
			}
			return
		}

		if msgType == websocket.TextMessage {
			msg, err := ioutil.ReadAll(r)
			if err != nil {
				log.Printf("error: %v", err)
				return
			}
			c.out.Write(msg)
			c.out.Write([]byte("\r"))
			//if _, err := io.Copy(c.out, r); err != nil {
			//	log.Printf("error: %v", err)
			//	return
			//}
		}
		//if err := c.out.Close(); err != nil {
		//	log.Printf("error: %v", err)
		//	return
		//}
	}
}
<|endoftext|>


<commit_before>package main

import(
	"fmt"
	"strings"
	"io/ioutil"
	"log"
	"strconv"
)

func main(){
	// Load the file in to memory.
	input, err := ioutil.ReadFile("input")
	handleError(err)
	// Convert the file into a string.
	inputString := string(input)

	// Split the string into a bunch of strings at each newline
	newLineSplit := strings.Split(inputString, "\n")
	// Now we should have an array of strings that look like this: "LLxWWxHH"

	// We need an array of arrays that can store our strings once we split them below.
	// This is our actual data. Hence the name.
	// The minus one is to deal with the last line of input, which will almost always be empty
	data := make([][]int, len(newLineSplit) - 1)

	// Split each of the resulting strings into length, width, and height by using x as a delimiter
	// temp will hold the split strings until we go through and convert them into ints
	// side is an iterator that is declared outside of the loop to save a bit of time in large datasets
	var side int
	for line := 0; line < len(newLineSplit) - 1; line++{
		temp := strings.Split(newLineSplit[line], "x")
		// First we have to make the line inside of our data
		data[line] = make([]int, 3)
		// This loop turns each of the strings into an int
		for side = 0; side < len(temp); side++{
			// Now we just store the values of temp in each element of data[line]
			data[line][side], err = strconv.Atoi(temp[side])
			handleError(err)
		}
	}
	// Now we have an nx3 array, where n is the number of lines in our input

	// Determine the total area needed for each package
	// This should be a for loop that performs a function
	// on each of the elements in the array created above
	// Add up the areas and store them in total
	totalFeet := 0
	for line := 0; line < len(data); line++{
		totalFeet += calcArea(data[line][0], data[line][1], data[line][2])
	}

	fmt.Printf("The elves need a total of %d square feet of wrapping paper!\n\n", totalFeet)
}

// This takes three inputs, and finds the total square footage needed for a given length, width, and height returns an int
// First it will take the three measurments needed (lxw, lxh, wxh)
// To find the slack needed, it will search the three values (lxw, lxh, wxh) for the smallest such value, and add it to the total square footage
// Finally, multiply each of the measurements by 2 and add them to teh total, then return total
func calcArea(length int, width int, height int) int{
	var i int
	measurements := [3]int {length * width, length * height, width * height}
	smallest := -1
  for i = 0; i < len(measurements); i++{
		if measurements[i] < smallest || smallest < 0{
			smallest = measurements[i]
		}
	}
	total := smallest
	for i = 0; i < len(measurements); i++{
		total += 2 * measurements[i]
	}

	return total
}

// Poor man's error handling.
func handleError(err error){
  if err != nil{
    log.Fatal(err)
  }
}
<commit_msg>Finished part 1 of day 2<commit_after>package main

import(
	"fmt"
	"strings"
	"io/ioutil"
	"log"
	"strconv"
)

func main(){
	// Load the file in to memory.
	input, err := ioutil.ReadFile("input")
	handleError(err)
	// Convert the file into a string.
	inputString := string(input)

	// Split the string into a bunch of strings at each newline
	newLineSplit := strings.Split(inputString, "\n")
	// Now we should have an array of strings that look like this: "LLxWWxHH"

	// We need an array of arrays that can store our strings once we split them below.
	// This is our actual data. Hence the name.
	// The minus one is to deal with the last line of input, which will almost always be empty
	data := make([][]int, len(newLineSplit) - 1)

	// Split each of the resulting strings into length, width, and height by using x as a delimiter
	// temp will hold the split strings until we go through and convert them into ints
	// side is an iterator that is declared outside of the loop to save a bit of time in large datasets
	// TODO: Break this out into a function
	var side int
	for line := 0; line < len(newLineSplit) - 1; line++{
		temp := strings.Split(newLineSplit[line], "x")
		// First we have to make the line inside of our data
		data[line] = make([]int, 3)
		// This loop turns each of the strings into an int
		for side = 0; side < len(temp); side++{
			// Now we just store the values of temp in each element of data[line]
			data[line][side], err = strconv.Atoi(temp[side])
			handleError(err)
		}
	}
	// Now we have an nx3 array, where n is the number of lines in our input

	// Determine the total area needed for each package
	// This should be a for loop that performs a function
	// on each of the elements in the array created above
	// Add up the areas and store them in total
	totalFeet := 0
	for line := 0; line < len(data); line++{
		totalFeet += calcArea(data[line][0], data[line][1], data[line][2])
	}

	fmt.Printf("The elves need a total of %d square feet of wrapping paper!\n\n", totalFeet)
}

// This takes three inputs, and finds the total square footage needed for a given length, width, and height returns an int
// First it will take the three measurments needed (lxw, lxh, wxh)
// To find the slack needed, it will search the three values (lxw, lxh, wxh) for the smallest such value, and add it to the total square footage
// Finally, multiply each of the measurements by 2 and add them to teh total, then return total
func calcArea(length int, width int, height int) int{
	var i int
	measurements := [3]int {length * width, length * height, width * height}
	smallest := -1
  for i = 0; i < len(measurements); i++{
		if measurements[i] < smallest || smallest < 0{
			smallest = measurements[i]
		}
	}
	total := smallest
	for i = 0; i < len(measurements); i++{
		total += 2 * measurements[i]
	}

	return total
}

// Poor man's error handling.
func handleError(err error){
  if err != nil{
    log.Fatal(err)
  }
}
<|endoftext|>


<commit_before>package manager

import (
	"fmt"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/cache"
	ds "github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/deliveryservice"
	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/enum"
	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/health"
	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/peer"
	todata "github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/trafficopsdata"
)

type DurationMapThreadsafe struct {
	durationMap map[enum.CacheName]time.Duration // TODO change string -> CacheName
	m           *sync.Mutex
}

func copyDurationMap(a map[enum.CacheName]time.Duration) map[enum.CacheName]time.Duration {
	b := map[enum.CacheName]time.Duration{}
	for k, v := range a {
		b[k] = v
	}
	return b
}

func NewDurationMapThreadsafe() DurationMapThreadsafe {
	return DurationMapThreadsafe{m: &sync.Mutex{}, durationMap: map[enum.CacheName]time.Duration{}}
}

func (o *DurationMapThreadsafe) Get() map[enum.CacheName]time.Duration {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	return copyDurationMap(o.durationMap)
}

func (o *DurationMapThreadsafe) GetDuration(cacheName enum.CacheName) (time.Duration, bool) {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	duration, ok := o.durationMap[cacheName]
	return duration, ok
}

func (o *DurationMapThreadsafe) Set(cacheName enum.CacheName, d time.Duration) {
	o.m.Lock()
	o.durationMap[cacheName] = d
	o.m.Unlock()
}

type TimeMapThreadsafe struct {
	timeMap map[enum.CacheName]time.Time // TODO change string -> CacheName
	m       *sync.Mutex
}

func copyTimeMap(a map[enum.CacheName]time.Time) map[enum.CacheName]time.Time {
	b := map[enum.CacheName]time.Time{}
	for k, v := range a {
		b[k] = v
	}
	return b
}

func NewTimeMapThreadsafe() TimeMapThreadsafe {
	return TimeMapThreadsafe{m: &sync.Mutex{}, timeMap: map[enum.CacheName]time.Time{}}
}

func (o TimeMapThreadsafe) Get() map[enum.CacheName]time.Time {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	return copyTimeMap(o.timeMap)
}

func (o *TimeMapThreadsafe) GetTime(cacheName enum.CacheName) (time.Time, bool) {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	time, ok := o.timeMap[cacheName]
	return time, ok
}

func (o *TimeMapThreadsafe) Set(cacheName enum.CacheName, d time.Time) {
	o.m.Lock()
	o.timeMap[cacheName] = d
	o.m.Unlock()
}

type ResultsThreadsafe struct {
	r map[enum.CacheName][]cache.Result // TODO change string -> CacheName
	m *sync.Mutex
}

func NewResultsThreadsafe() ResultsThreadsafe {
	return ResultsThreadsafe{m: &sync.Mutex{}, r: map[enum.CacheName][]cache.Result{}}
}

func (o *ResultsThreadsafe) Get(cacheName enum.CacheName) []cache.Result {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	return o.r[cacheName]
}

func (o *ResultsThreadsafe) Set(cacheName enum.CacheName, newR []cache.Result) {
	o.m.Lock()
	o.r[cacheName] = newR
	o.m.Unlock()
}

// StartHealthResultManager starts the goroutine which listens for health results.
// Returns the last health durations, events, and the local cache statuses.
func StartHealthResultManager(cacheHealthChan <-chan cache.Result, toData todata.TODataThreadsafe, localStates peer.CRStatesThreadsafe, statHistory StatHistoryThreadsafe, monitorConfig TrafficMonitorConfigMapThreadsafe, peerStates peer.CRStatesPeersThreadsafe, combinedStates peer.CRStatesThreadsafe, fetchCount UintThreadsafe, errorCount UintThreadsafe) (DurationMapThreadsafe, EventsThreadsafe, CacheAvailableStatusThreadsafe, DSStatsThreadsafe, StatsLastKbpsThreadsafe) {
	lastHealthDurations := NewDurationMapThreadsafe()
	events := NewEventsThreadsafe()
	localCacheStatus := NewCacheAvailableStatusThreadsafe()
	dsStats := NewDSStatsThreadsafe()
	lastKbpsStats := NewStatsLastKbpsThreadsafe()
	go healthResultManagerListen(cacheHealthChan, toData, localStates, lastHealthDurations, statHistory, monitorConfig, peerStates, combinedStates, fetchCount, errorCount, events, localCacheStatus, dsStats, lastKbpsStats)
	return lastHealthDurations, events, localCacheStatus, dsStats, lastKbpsStats
}

func healthResultManagerListen(cacheHealthChan <-chan cache.Result, toData todata.TODataThreadsafe, localStates peer.CRStatesThreadsafe, lastHealthDurations DurationMapThreadsafe, statHistory StatHistoryThreadsafe, monitorConfig TrafficMonitorConfigMapThreadsafe, peerStates peer.CRStatesPeersThreadsafe, combinedStates peer.CRStatesThreadsafe, fetchCount UintThreadsafe, errorCount UintThreadsafe, events EventsThreadsafe, localCacheStatus CacheAvailableStatusThreadsafe, dsStats DSStatsThreadsafe, lastKbpsStats StatsLastKbpsThreadsafe) {
	lastHealthEndTimes := map[enum.CacheName]time.Time{}
	healthHistory := map[enum.CacheName][]cache.Result{}
	eventIndex := uint64(0) // TODO move to EventsThreadsafe.Add() ?
	// lastHealthEndTimes := NewTimeMapThreadsafe()
	// healthHistory := NewResultsThreadsafe()
	// eventIndex := NewUintThreadsafe()

	for {
		select {
		case healthResult := <-cacheHealthChan:
			fmt.Printf("DEBUG poll %v %v healthresultman start\n", healthResult.PollID, time.Now())
			// go func() {
			{
				fetchCount.Inc()
				toDataCopy := toData.Get() // create a copy, so the same data used for all processing of this cache health result
				var prevResult cache.Result
				healthResultHistory := healthHistory[enum.CacheName(healthResult.Id)]
				// healthResultHistory := healthHistory.Get(enum.CacheName(healthResult.Id))
				if len(healthResultHistory) != 0 {
					prevResult = healthResultHistory[len(healthResultHistory)-1]
				}

				monitorConfigCopy := monitorConfig.Get() // copy now, so all calculations are on the same data
				health.GetVitals(&healthResult, &prevResult, &monitorConfigCopy)
				// healthHistory.Set(enum.CacheName(healthResult.Id), pruneHistory(append(healthHistory.Get(enum.CacheName(healthResult.Id)), healthResult), defaultMaxHistory))
				healthHistory[enum.CacheName(healthResult.Id)] = pruneHistory(append(healthHistory[enum.CacheName(healthResult.Id)], healthResult), defaultMaxHistory)
				isAvailable, whyAvailable := health.EvalCache(healthResult, &monitorConfigCopy)
				if localStates.Get().Caches[healthResult.Id].IsAvailable != isAvailable {
					fmt.Println("Changing state for", healthResult.Id, " was:", prevResult.Available, " is now:", isAvailable, " because:", whyAvailable, " errors:", healthResult.Errors)
					e := Event{Index: eventIndex, Time: time.Now().Unix(), Description: whyAvailable, Name: healthResult.Id, Hostname: healthResult.Id, Type: toDataCopy.ServerTypes[enum.CacheName(healthResult.Id)].String(), Available: isAvailable}
					// e := Event{Index: eventIndex.Get(), Time: time.Now().Unix(), Description: whyAvailable, Name: healthResult.Id, Hostname: healthResult.Id, Type: toDataCopy.ServerTypes[enum.CacheName(healthResult.Id)].String(), Available: isAvailable}
					eventIndex++
					// eventIndex.Inc()
					events.Add(e)
				}

				localCacheStatus.Set(enum.CacheName(healthResult.Id), CacheAvailableStatus{Available: isAvailable, Status: monitorConfigCopy.TrafficServer[healthResult.Id].Status}) // TODO move within localStates
				localStates.SetCache(healthResult.Id, peer.IsAvailable{IsAvailable: isAvailable})
				fmt.Printf("DEBUG poll %v %v calculateDeliveryServiceState start\n", healthResult.PollID, time.Now())
				calculateDeliveryServiceState(toDataCopy.DeliveryServiceServers, localStates)
				fmt.Printf("DEBUG poll %v %v calculateDeliveryServiceState end\n", healthResult.PollID, time.Now())

				// TODO determine if we should combineCrStates() here

				now := time.Now()

				var err error
				createStatsCopyStatHistory := statHistory.Get()
				createStatsCopyCombinedStates := combinedStates.Get()
				createStatsCopyLastKbpsStats := lastKbpsStats.Get()
				fmt.Printf("DEBUG poll %v %v CreateStats start\n", healthResult.PollID, time.Now())
				newDsStats, newLastKbpsStats, err := ds.CreateStats(createStatsCopyStatHistory, toDataCopy, createStatsCopyCombinedStates, createStatsCopyLastKbpsStats, now)
				fmt.Printf("DEBUG poll %v %v CreateStats end\n", healthResult.PollID, time.Now())
				if err != nil {
					errorCount.Inc()
					log.Printf("ERROR getting deliveryservice: %v\n", err)
				} else {
					dsStats.Set(newDsStats)
					lastKbpsStats.Set(newLastKbpsStats)
				}

				//				if lastHealthStart, ok := lastHealthEndTimes.GetTime(enum.CacheName(healthResult.Id)); ok {
				if lastHealthStart, ok := lastHealthEndTimes[enum.CacheName(healthResult.Id)]; ok {
					d := time.Since(lastHealthStart)
					lastHealthDurations.Set(enum.CacheName(healthResult.Id), d)
				}
				lastHealthEndTimes[enum.CacheName(healthResult.Id)] = now
				// lastHealthEndTimes.Set(enum.CacheName(healthResult.Id), now)

				fmt.Printf("DEBUG poll %v %v finish\n", healthResult.PollID, time.Now())
				// }()
			}
		}
	}
}

// calculateDeliveryServiceState calculates the state of delivery services from the new cache state data `cacheState` and the CRConfig data `deliveryServiceServers` and puts the calculated state in the outparam `deliveryServiceStates`
func calculateDeliveryServiceState(deliveryServiceServers map[string][]string, states peer.CRStatesThreadsafe) {
	deliveryServices := states.GetDeliveryServices()
	for deliveryServiceName, deliveryServiceState := range deliveryServices {
		if _, ok := deliveryServiceServers[deliveryServiceName]; !ok {
			// log.Printf("ERROR CRConfig does not have delivery service %s, but traffic monitor poller does; skipping\n", deliveryServiceName)
			continue
		}
		deliveryServiceState.IsAvailable = false
		deliveryServiceState.DisabledLocations = nil
		for _, server := range deliveryServiceServers[deliveryServiceName] {
			if states.GetCache(server).IsAvailable {
				deliveryServiceState.IsAvailable = true
			} else {
				deliveryServiceState.DisabledLocations = append(deliveryServiceState.DisabledLocations, server)
			}
		}
		deliveryServices[deliveryServiceName] = deliveryServiceState
	}
	states.SetDeliveryServices(deliveryServices)
}

// intersection returns strings in both a and b.
// Note this modifies a and b. Specifically, it sorts them. If that isn't acceptable, pass copies of your real data.
func intersection(a []string, b []string) []string {
	sort.Strings(a)
	sort.Strings(b)
	var c []string
	for _, s := range a {
		i := sort.SearchStrings(b, s)
		if i < len(b) && b[i] == s {
			c = append(c, s)
		}
	}
	return c
}
<commit_msg>Change TM2 cache result man to bulk aggregate<commit_after>package manager

import (
	"fmt"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/cache"
	ds "github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/deliveryservice"
	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/enum"
	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/health"
	"github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/peer"
	todata "github.com/Comcast/traffic_control/traffic_monitor/experimental/traffic_monitor/trafficopsdata"
)

type DurationMapThreadsafe struct {
	durationMap map[enum.CacheName]time.Duration // TODO change string -> CacheName
	m           *sync.Mutex
}

func copyDurationMap(a map[enum.CacheName]time.Duration) map[enum.CacheName]time.Duration {
	b := map[enum.CacheName]time.Duration{}
	for k, v := range a {
		b[k] = v
	}
	return b
}

func NewDurationMapThreadsafe() DurationMapThreadsafe {
	return DurationMapThreadsafe{m: &sync.Mutex{}, durationMap: map[enum.CacheName]time.Duration{}}
}

func (o *DurationMapThreadsafe) Get() map[enum.CacheName]time.Duration {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	return copyDurationMap(o.durationMap)
}

func (o *DurationMapThreadsafe) GetDuration(cacheName enum.CacheName) (time.Duration, bool) {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	duration, ok := o.durationMap[cacheName]
	return duration, ok
}

func (o *DurationMapThreadsafe) Set(cacheName enum.CacheName, d time.Duration) {
	o.m.Lock()
	o.durationMap[cacheName] = d
	o.m.Unlock()
}

type TimeMapThreadsafe struct {
	timeMap map[enum.CacheName]time.Time // TODO change string -> CacheName
	m       *sync.Mutex
}

func copyTimeMap(a map[enum.CacheName]time.Time) map[enum.CacheName]time.Time {
	b := map[enum.CacheName]time.Time{}
	for k, v := range a {
		b[k] = v
	}
	return b
}

func NewTimeMapThreadsafe() TimeMapThreadsafe {
	return TimeMapThreadsafe{m: &sync.Mutex{}, timeMap: map[enum.CacheName]time.Time{}}
}

func (o TimeMapThreadsafe) Get() map[enum.CacheName]time.Time {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	return copyTimeMap(o.timeMap)
}

func (o *TimeMapThreadsafe) GetTime(cacheName enum.CacheName) (time.Time, bool) {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	time, ok := o.timeMap[cacheName]
	return time, ok
}

func (o *TimeMapThreadsafe) Set(cacheName enum.CacheName, d time.Time) {
	o.m.Lock()
	o.timeMap[cacheName] = d
	o.m.Unlock()
}

type ResultsThreadsafe struct {
	r map[enum.CacheName][]cache.Result // TODO change string -> CacheName
	m *sync.Mutex
}

func NewResultsThreadsafe() ResultsThreadsafe {
	return ResultsThreadsafe{m: &sync.Mutex{}, r: map[enum.CacheName][]cache.Result{}}
}

func (o *ResultsThreadsafe) Get(cacheName enum.CacheName) []cache.Result {
	o.m.Lock()
	defer func() {
		o.m.Unlock()
	}()
	return o.r[cacheName]
}

func (o *ResultsThreadsafe) Set(cacheName enum.CacheName, newR []cache.Result) {
	o.m.Lock()
	o.r[cacheName] = newR
	o.m.Unlock()
}

// StartHealthResultManager starts the goroutine which listens for health results.
// Returns the last health durations, events, and the local cache statuses.
func StartHealthResultManager(cacheHealthChan <-chan cache.Result, toData todata.TODataThreadsafe, localStates peer.CRStatesThreadsafe, statHistory StatHistoryThreadsafe, monitorConfig TrafficMonitorConfigMapThreadsafe, peerStates peer.CRStatesPeersThreadsafe, combinedStates peer.CRStatesThreadsafe, fetchCount UintThreadsafe, errorCount UintThreadsafe) (DurationMapThreadsafe, EventsThreadsafe, CacheAvailableStatusThreadsafe, DSStatsThreadsafe, StatsLastKbpsThreadsafe) {
	lastHealthDurations := NewDurationMapThreadsafe()
	events := NewEventsThreadsafe()
	localCacheStatus := NewCacheAvailableStatusThreadsafe()
	dsStats := NewDSStatsThreadsafe()
	lastKbpsStats := NewStatsLastKbpsThreadsafe()
	go healthResultManagerListen(cacheHealthChan, toData, localStates, lastHealthDurations, statHistory, monitorConfig, peerStates, combinedStates, fetchCount, errorCount, events, localCacheStatus, dsStats, lastKbpsStats)
	return lastHealthDurations, events, localCacheStatus, dsStats, lastKbpsStats
}

// cacheAggregateSeconds is how often to aggregate stats, if the health chan is never empty. (Otherwise, we read from the chan until it's empty, then aggregate, continuously)
const cacheAggregateSeconds = 1

func healthResultManagerListen(cacheHealthChan <-chan cache.Result, toData todata.TODataThreadsafe, localStates peer.CRStatesThreadsafe, lastHealthDurations DurationMapThreadsafe, statHistory StatHistoryThreadsafe, monitorConfig TrafficMonitorConfigMapThreadsafe, peerStates peer.CRStatesPeersThreadsafe, combinedStates peer.CRStatesThreadsafe, fetchCount UintThreadsafe, errorCount UintThreadsafe, events EventsThreadsafe, localCacheStatus CacheAvailableStatusThreadsafe, dsStats DSStatsThreadsafe, lastKbpsStats StatsLastKbpsThreadsafe) {
	lastHealthEndTimes := map[enum.CacheName]time.Time{}
	healthHistory := map[enum.CacheName][]cache.Result{}
	eventIndex := uint64(0) // TODO move to EventsThreadsafe.Add() ?
	// lastHealthEndTimes := NewTimeMapThreadsafe()
	// healthHistory := NewResultsThreadsafe()
	// eventIndex := NewUintThreadsafe()

	var results []cache.Result
	for {
		select {
		// TODO add a tick, in case cacheHealthChan is never empty (starvation)
		case r := <-cacheHealthChan:
			results = append(results, r)
		default:
			// NOTE this busy waits. TODO figure out if there's a way in Go to listen to a chan until it's written to, then drain the channel, without defaulting and busy-waiting or sleeping
			if len(results) == 0 {
				continue
			}
			toDataCopy := toData.Get()               // create a copy, so the same data used for all processing of this cache health result
			monitorConfigCopy := monitorConfig.Get() // copy now, so all calculations are on the same data
			for _, healthResult := range results {
				fmt.Printf("DEBUG poll %v %v healthresultman start\n", healthResult.PollID, time.Now())
				fetchCount.Inc()
				var prevResult cache.Result
				healthResultHistory := healthHistory[enum.CacheName(healthResult.Id)]
				// healthResultHistory := healthHistory.Get(enum.CacheName(healthResult.Id))
				if len(healthResultHistory) != 0 {
					prevResult = healthResultHistory[len(healthResultHistory)-1]
				}

				health.GetVitals(&healthResult, &prevResult, &monitorConfigCopy)
				// healthHistory.Set(enum.CacheName(healthResult.Id), pruneHistory(append(healthHistory.Get(enum.CacheName(healthResult.Id)), healthResult), defaultMaxHistory))
				healthHistory[enum.CacheName(healthResult.Id)] = pruneHistory(append(healthHistory[enum.CacheName(healthResult.Id)], healthResult), defaultMaxHistory)
				isAvailable, whyAvailable := health.EvalCache(healthResult, &monitorConfigCopy)
				if localStates.Get().Caches[healthResult.Id].IsAvailable != isAvailable {
					fmt.Println("Changing state for", healthResult.Id, " was:", prevResult.Available, " is now:", isAvailable, " because:", whyAvailable, " errors:", healthResult.Errors)
					e := Event{Index: eventIndex, Time: time.Now().Unix(), Description: whyAvailable, Name: healthResult.Id, Hostname: healthResult.Id, Type: toDataCopy.ServerTypes[enum.CacheName(healthResult.Id)].String(), Available: isAvailable}
					// e := Event{Index: eventIndex.Get(), Time: time.Now().Unix(), Description: whyAvailable, Name: healthResult.Id, Hostname: healthResult.Id, Type: toDataCopy.ServerTypes[enum.CacheName(healthResult.Id)].String(), Available: isAvailable}
					eventIndex++
					// eventIndex.Inc()
					events.Add(e)
				}

				localCacheStatus.Set(enum.CacheName(healthResult.Id), CacheAvailableStatus{Available: isAvailable, Status: monitorConfigCopy.TrafficServer[healthResult.Id].Status}) // TODO move within localStates
				localStates.SetCache(healthResult.Id, peer.IsAvailable{IsAvailable: isAvailable})
				fmt.Printf("DEBUG poll %v %v calculateDeliveryServiceState start\n", healthResult.PollID, time.Now())
				calculateDeliveryServiceState(toDataCopy.DeliveryServiceServers, localStates)
				fmt.Printf("DEBUG poll %v %v calculateDeliveryServiceState end\n", healthResult.PollID, time.Now())
			}
			// TODO determine if we should combineCrStates() here

			now := time.Now()

			var err error
			createStatsCopyStatHistory := statHistory.Get()
			createStatsCopyCombinedStates := combinedStates.Get()
			createStatsCopyLastKbpsStats := lastKbpsStats.Get()

			for _, healthResult := range results {
				fmt.Printf("DEBUG poll %v %v CreateStats start\n", healthResult.PollID, time.Now())
			}

			newDsStats, newLastKbpsStats, err := ds.CreateStats(createStatsCopyStatHistory, toDataCopy, createStatsCopyCombinedStates, createStatsCopyLastKbpsStats, now)

			for _, healthResult := range results {
				fmt.Printf("DEBUG poll %v %v CreateStats end\n", healthResult.PollID, time.Now())
			}

			if err != nil {
				errorCount.Inc()
				log.Printf("ERROR getting deliveryservice: %v\n", err)
			} else {
				dsStats.Set(newDsStats)
				lastKbpsStats.Set(newLastKbpsStats)
			}

			for _, healthResult := range results {
				//				if lastHealthStart, ok := lastHealthEndTimes.GetTime(enum.CacheName(healthResult.Id)); ok {
				if lastHealthStart, ok := lastHealthEndTimes[enum.CacheName(healthResult.Id)]; ok {
					d := time.Since(lastHealthStart)
					lastHealthDurations.Set(enum.CacheName(healthResult.Id), d)
				}
				lastHealthEndTimes[enum.CacheName(healthResult.Id)] = now
				// lastHealthEndTimes.Set(enum.CacheName(healthResult.Id), now)

				fmt.Printf("DEBUG poll %v %v finish\n", healthResult.PollID, time.Now())
			}

			results = nil
		}
	}
}

// calculateDeliveryServiceState calculates the state of delivery services from the new cache state data `cacheState` and the CRConfig data `deliveryServiceServers` and puts the calculated state in the outparam `deliveryServiceStates`
func calculateDeliveryServiceState(deliveryServiceServers map[string][]string, states peer.CRStatesThreadsafe) {
	deliveryServices := states.GetDeliveryServices()
	for deliveryServiceName, deliveryServiceState := range deliveryServices {
		if _, ok := deliveryServiceServers[deliveryServiceName]; !ok {
			// log.Printf("ERROR CRConfig does not have delivery service %s, but traffic monitor poller does; skipping\n", deliveryServiceName)
			continue
		}
		deliveryServiceState.IsAvailable = false
		deliveryServiceState.DisabledLocations = nil
		for _, server := range deliveryServiceServers[deliveryServiceName] {
			if states.GetCache(server).IsAvailable {
				deliveryServiceState.IsAvailable = true
			} else {
				deliveryServiceState.DisabledLocations = append(deliveryServiceState.DisabledLocations, server)
			}
		}
		deliveryServices[deliveryServiceName] = deliveryServiceState
	}
	states.SetDeliveryServices(deliveryServices)
}

// intersection returns strings in both a and b.
// Note this modifies a and b. Specifically, it sorts them. If that isn't acceptable, pass copies of your real data.
func intersection(a []string, b []string) []string {
	sort.Strings(a)
	sort.Strings(b)
	var c []string
	for _, s := range a {
		i := sort.SearchStrings(b, s)
		if i < len(b) && b[i] == s {
			c = append(c, s)
		}
	}
	return c
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd 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 command

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"unsafe"

	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/services/server"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/mgr"
)

var (
	serviceNameFlag       string
	registerServiceFlag   bool
	unregisterServiceFlag bool
	runServiceFlag        bool
	logFileFlag           string

	kernel32     = windows.NewLazySystemDLL("kernel32.dll")
	setStdHandle = kernel32.NewProc("SetStdHandle")
	allocConsole = kernel32.NewProc("AllocConsole")
	oldStderr    windows.Handle
	panicFile    *os.File
)

const defaultServiceName = "containerd"

// serviceFlags returns an array of flags for configuring containerd to run
// as a Windows service under control of SCM.
func serviceFlags() []cli.Flag {
	return []cli.Flag{
		cli.StringFlag{
			Name:  "service-name",
			Usage: "Set the Windows service name",
			Value: defaultServiceName,
		},
		cli.BoolFlag{
			Name:  "register-service",
			Usage: "Register the service and exit",
		},
		cli.BoolFlag{
			Name:  "unregister-service",
			Usage: "Unregister the service and exit",
		},
		cli.BoolFlag{
			Name:   "run-service",
			Usage:  "",
			Hidden: true,
		},
		cli.StringFlag{
			Name:  "log-file",
			Usage: "Path to the containerd log file",
		},
	}
}

// applyPlatformFlags applies platform-specific flags.
func applyPlatformFlags(context *cli.Context) {
	serviceNameFlag = context.GlobalString("service-name")
	if serviceNameFlag == "" {
		serviceNameFlag = defaultServiceName
	}
	for _, v := range []struct {
		name string
		d    *bool
	}{
		{
			name: "register-service",
			d:    &registerServiceFlag,
		},
		{
			name: "unregister-service",
			d:    &unregisterServiceFlag,
		},
		{
			name: "run-service",
			d:    &runServiceFlag,
		},
	} {
		*v.d = context.GlobalBool(v.name)
	}
	logFileFlag = context.GlobalString("log-file")
}

type handler struct {
	fromsvc chan error
	s       *server.Server
	done    chan struct{} // Indicates back to app main to quit
}

func getServicePath() (string, error) {
	p, err := exec.LookPath(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Abs(p)
}

func registerService() error {
	p, err := getServicePath()
	if err != nil {
		return err
	}
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	c := mgr.Config{
		ServiceType:  windows.SERVICE_WIN32_OWN_PROCESS,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
		DisplayName:  "Containerd",
		Description:  "Container runtime",
	}

	// Configure the service to launch with the arguments that were just passed.
	args := []string{"--run-service"}
	for _, a := range os.Args[1:] {
		if a != "--register-service" && a != "--unregister-service" {
			args = append(args, a)
		}
	}

	s, err := m.CreateService(serviceNameFlag, p, c, args...)
	if err != nil {
		return err
	}
	defer s.Close()

	// See http://stackoverflow.com/questions/35151052/how-do-i-configure-failure-actions-of-a-windows-service-written-in-go
	const (
		scActionNone    = 0
		scActionRestart = 1

		serviceConfigFailureActions = 2
	)

	type serviceFailureActions struct {
		ResetPeriod  uint32
		RebootMsg    *uint16
		Command      *uint16
		ActionsCount uint32
		Actions      uintptr
	}

	type scAction struct {
		Type  uint32
		Delay uint32
	}
	t := []scAction{
		{Type: scActionRestart, Delay: uint32(15 * time.Second / time.Millisecond)},
		{Type: scActionRestart, Delay: uint32(15 * time.Second / time.Millisecond)},
		{Type: scActionNone},
	}
	lpInfo := serviceFailureActions{ResetPeriod: uint32(24 * time.Hour / time.Second), ActionsCount: uint32(3), Actions: uintptr(unsafe.Pointer(&t[0]))}
	err = windows.ChangeServiceConfig2(s.Handle, serviceConfigFailureActions, (*byte)(unsafe.Pointer(&lpInfo)))
	if err != nil {
		return err
	}

	return nil
}

func unregisterService() error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	s, err := m.OpenService(serviceNameFlag)
	if err != nil {
		return err
	}
	defer s.Close()

	err = s.Delete()
	if err != nil {
		return err
	}
	return nil
}

// registerUnregisterService is an entrypoint early in the daemon startup
// to handle (un-)registering against Windows Service Control Manager (SCM).
// It returns an indication to stop on successful SCM operation, and an error.
func registerUnregisterService(root string) (bool, error) {

	if unregisterServiceFlag {
		if registerServiceFlag {
			return true, errors.Wrap(errdefs.ErrInvalidArgument, "--register-service and --unregister-service cannot be used together")
		}
		return true, unregisterService()
	}

	if registerServiceFlag {
		return true, registerService()
	}

	if runServiceFlag {
		// Allocate a conhost for containerd here. We don't actually use this
		// at all in containerd, but it will be inherited by any processes
		// containerd executes, so they won't need to allocate their own
		// conhosts. This is important for two reasons:
		// - Creating a conhost slows down process launch.
		// - We have seen reliability issues when launching many processes.
		//   Sometimes the process invocation will fail due to an error when
		//   creating the conhost.
		//
		// This needs to be done before initializing the panic file, as
		// AllocConsole sets the stdio handles to point to the new conhost,
		// and we want to make sure stderr goes to the panic file.
		r, _, err := allocConsole.Call()
		if r == 0 && err != nil {
			return true, fmt.Errorf("error allocating conhost: %s", err)
		}

		if err := initPanicFile(filepath.Join(root, "panic.log")); err != nil {
			return true, err
		}

		logOutput := ioutil.Discard
		if logFileFlag != "" {
			f, err := os.OpenFile(logFileFlag, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
			if err != nil {
				return true, errors.Wrapf(err, "open log file %q", logFileFlag)
			}
			logOutput = f
		}
		logrus.SetOutput(logOutput)
	}
	return false, nil
}

// launchService is the entry point for running the daemon under SCM.
func launchService(s *server.Server, done chan struct{}) error {

	if !runServiceFlag {
		return nil
	}

	h := &handler{
		fromsvc: make(chan error),
		s:       s,
		done:    done,
	}

	interactive, err := svc.IsAnInteractiveSession() // nolint:staticcheck
	if err != nil {
		return err
	}

	go func() {
		if interactive {
			err = debug.Run(serviceNameFlag, h)
		} else {
			err = svc.Run(serviceNameFlag, h)
		}
		h.fromsvc <- err
	}()

	// Wait for the first signal from the service handler.
	err = <-h.fromsvc
	if err != nil {
		return err
	}
	return nil
}

func (h *handler) Execute(_ []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (bool, uint32) {
	s <- svc.Status{State: svc.StartPending, Accepts: 0}
	// Unblock launchService()
	h.fromsvc <- nil

	s <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown | svc.Accepted(windows.SERVICE_ACCEPT_PARAMCHANGE)}

Loop:
	for c := range r {
		switch c.Cmd {
		case svc.Interrogate:
			s <- c.CurrentStatus
		case svc.Stop, svc.Shutdown:
			s <- svc.Status{State: svc.StopPending, Accepts: 0}
			h.s.Stop()
			break Loop
		}
	}

	removePanicFile()
	close(h.done)
	return false, 0
}

func initPanicFile(path string) error {
	var err error
	panicFile, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0)
	if err != nil {
		return err
	}

	st, err := panicFile.Stat()
	if err != nil {
		return err
	}

	// If there are contents in the file already, move the file out of the way
	// and replace it.
	if st.Size() > 0 {
		panicFile.Close()
		os.Rename(path, path+".old")
		panicFile, err = os.Create(path)
		if err != nil {
			return err
		}
	}

	// Update STD_ERROR_HANDLE to point to the panic file so that Go writes to
	// it when it panics. Remember the old stderr to restore it before removing
	// the panic file.
	sh := uint32(windows.STD_ERROR_HANDLE)
	h, err := windows.GetStdHandle(sh)
	if err != nil {
		return err
	}

	oldStderr = h

	r, _, err := setStdHandle.Call(uintptr(sh), panicFile.Fd())
	if r == 0 && err != nil {
		return err
	}

	// Reset os.Stderr to the panic file (so fmt.Fprintf(os.Stderr,...) actually gets redirected)
	os.Stderr = os.NewFile(panicFile.Fd(), "/dev/stderr")

	// Force threads that panic to write to stderr (the panicFile handle now), otherwise it will go into the ether
	log.SetOutput(os.Stderr)

	return nil
}

func removePanicFile() {
	if st, err := panicFile.Stat(); err == nil {
		if st.Size() == 0 {
			sh := uint32(windows.STD_ERROR_HANDLE)
			setStdHandle.Call(uintptr(sh), uintptr(oldStderr))
			panicFile.Close()
			os.Remove(panicFile.Name())
		}
	}
}
<commit_msg>Fix Windows service panic file to not be read-only<commit_after>/*
   Copyright The containerd 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 command

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"unsafe"

	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/services/server"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/mgr"
)

var (
	serviceNameFlag       string
	registerServiceFlag   bool
	unregisterServiceFlag bool
	runServiceFlag        bool
	logFileFlag           string

	kernel32     = windows.NewLazySystemDLL("kernel32.dll")
	setStdHandle = kernel32.NewProc("SetStdHandle")
	allocConsole = kernel32.NewProc("AllocConsole")
	oldStderr    windows.Handle
	panicFile    *os.File
)

const defaultServiceName = "containerd"

// serviceFlags returns an array of flags for configuring containerd to run
// as a Windows service under control of SCM.
func serviceFlags() []cli.Flag {
	return []cli.Flag{
		cli.StringFlag{
			Name:  "service-name",
			Usage: "Set the Windows service name",
			Value: defaultServiceName,
		},
		cli.BoolFlag{
			Name:  "register-service",
			Usage: "Register the service and exit",
		},
		cli.BoolFlag{
			Name:  "unregister-service",
			Usage: "Unregister the service and exit",
		},
		cli.BoolFlag{
			Name:   "run-service",
			Usage:  "",
			Hidden: true,
		},
		cli.StringFlag{
			Name:  "log-file",
			Usage: "Path to the containerd log file",
		},
	}
}

// applyPlatformFlags applies platform-specific flags.
func applyPlatformFlags(context *cli.Context) {
	serviceNameFlag = context.GlobalString("service-name")
	if serviceNameFlag == "" {
		serviceNameFlag = defaultServiceName
	}
	for _, v := range []struct {
		name string
		d    *bool
	}{
		{
			name: "register-service",
			d:    &registerServiceFlag,
		},
		{
			name: "unregister-service",
			d:    &unregisterServiceFlag,
		},
		{
			name: "run-service",
			d:    &runServiceFlag,
		},
	} {
		*v.d = context.GlobalBool(v.name)
	}
	logFileFlag = context.GlobalString("log-file")
}

type handler struct {
	fromsvc chan error
	s       *server.Server
	done    chan struct{} // Indicates back to app main to quit
}

func getServicePath() (string, error) {
	p, err := exec.LookPath(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Abs(p)
}

func registerService() error {
	p, err := getServicePath()
	if err != nil {
		return err
	}
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	c := mgr.Config{
		ServiceType:  windows.SERVICE_WIN32_OWN_PROCESS,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
		DisplayName:  "Containerd",
		Description:  "Container runtime",
	}

	// Configure the service to launch with the arguments that were just passed.
	args := []string{"--run-service"}
	for _, a := range os.Args[1:] {
		if a != "--register-service" && a != "--unregister-service" {
			args = append(args, a)
		}
	}

	s, err := m.CreateService(serviceNameFlag, p, c, args...)
	if err != nil {
		return err
	}
	defer s.Close()

	// See http://stackoverflow.com/questions/35151052/how-do-i-configure-failure-actions-of-a-windows-service-written-in-go
	const (
		scActionNone    = 0
		scActionRestart = 1

		serviceConfigFailureActions = 2
	)

	type serviceFailureActions struct {
		ResetPeriod  uint32
		RebootMsg    *uint16
		Command      *uint16
		ActionsCount uint32
		Actions      uintptr
	}

	type scAction struct {
		Type  uint32
		Delay uint32
	}
	t := []scAction{
		{Type: scActionRestart, Delay: uint32(15 * time.Second / time.Millisecond)},
		{Type: scActionRestart, Delay: uint32(15 * time.Second / time.Millisecond)},
		{Type: scActionNone},
	}
	lpInfo := serviceFailureActions{ResetPeriod: uint32(24 * time.Hour / time.Second), ActionsCount: uint32(3), Actions: uintptr(unsafe.Pointer(&t[0]))}
	err = windows.ChangeServiceConfig2(s.Handle, serviceConfigFailureActions, (*byte)(unsafe.Pointer(&lpInfo)))
	if err != nil {
		return err
	}

	return nil
}

func unregisterService() error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	s, err := m.OpenService(serviceNameFlag)
	if err != nil {
		return err
	}
	defer s.Close()

	err = s.Delete()
	if err != nil {
		return err
	}
	return nil
}

// registerUnregisterService is an entrypoint early in the daemon startup
// to handle (un-)registering against Windows Service Control Manager (SCM).
// It returns an indication to stop on successful SCM operation, and an error.
func registerUnregisterService(root string) (bool, error) {

	if unregisterServiceFlag {
		if registerServiceFlag {
			return true, errors.Wrap(errdefs.ErrInvalidArgument, "--register-service and --unregister-service cannot be used together")
		}
		return true, unregisterService()
	}

	if registerServiceFlag {
		return true, registerService()
	}

	if runServiceFlag {
		// Allocate a conhost for containerd here. We don't actually use this
		// at all in containerd, but it will be inherited by any processes
		// containerd executes, so they won't need to allocate their own
		// conhosts. This is important for two reasons:
		// - Creating a conhost slows down process launch.
		// - We have seen reliability issues when launching many processes.
		//   Sometimes the process invocation will fail due to an error when
		//   creating the conhost.
		//
		// This needs to be done before initializing the panic file, as
		// AllocConsole sets the stdio handles to point to the new conhost,
		// and we want to make sure stderr goes to the panic file.
		r, _, err := allocConsole.Call()
		if r == 0 && err != nil {
			return true, fmt.Errorf("error allocating conhost: %s", err)
		}

		if err := initPanicFile(filepath.Join(root, "panic.log")); err != nil {
			return true, err
		}

		logOutput := ioutil.Discard
		if logFileFlag != "" {
			f, err := os.OpenFile(logFileFlag, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
			if err != nil {
				return true, errors.Wrapf(err, "open log file %q", logFileFlag)
			}
			logOutput = f
		}
		logrus.SetOutput(logOutput)
	}
	return false, nil
}

// launchService is the entry point for running the daemon under SCM.
func launchService(s *server.Server, done chan struct{}) error {

	if !runServiceFlag {
		return nil
	}

	h := &handler{
		fromsvc: make(chan error),
		s:       s,
		done:    done,
	}

	interactive, err := svc.IsAnInteractiveSession() // nolint:staticcheck
	if err != nil {
		return err
	}

	go func() {
		if interactive {
			err = debug.Run(serviceNameFlag, h)
		} else {
			err = svc.Run(serviceNameFlag, h)
		}
		h.fromsvc <- err
	}()

	// Wait for the first signal from the service handler.
	err = <-h.fromsvc
	if err != nil {
		return err
	}
	return nil
}

func (h *handler) Execute(_ []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (bool, uint32) {
	s <- svc.Status{State: svc.StartPending, Accepts: 0}
	// Unblock launchService()
	h.fromsvc <- nil

	s <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown | svc.Accepted(windows.SERVICE_ACCEPT_PARAMCHANGE)}

Loop:
	for c := range r {
		switch c.Cmd {
		case svc.Interrogate:
			s <- c.CurrentStatus
		case svc.Stop, svc.Shutdown:
			s <- svc.Status{State: svc.StopPending, Accepts: 0}
			h.s.Stop()
			break Loop
		}
	}

	removePanicFile()
	close(h.done)
	return false, 0
}

func initPanicFile(path string) error {
	var err error
	panicFile, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return err
	}

	st, err := panicFile.Stat()
	if err != nil {
		return err
	}

	// If there are contents in the file already, move the file out of the way
	// and replace it.
	if st.Size() > 0 {
		panicFile.Close()
		os.Rename(path, path+".old")
		panicFile, err = os.Create(path)
		if err != nil {
			return err
		}
	}

	// Update STD_ERROR_HANDLE to point to the panic file so that Go writes to
	// it when it panics. Remember the old stderr to restore it before removing
	// the panic file.
	sh := uint32(windows.STD_ERROR_HANDLE)
	h, err := windows.GetStdHandle(sh)
	if err != nil {
		return err
	}

	oldStderr = h

	r, _, err := setStdHandle.Call(uintptr(sh), panicFile.Fd())
	if r == 0 && err != nil {
		return err
	}

	// Reset os.Stderr to the panic file (so fmt.Fprintf(os.Stderr,...) actually gets redirected)
	os.Stderr = os.NewFile(panicFile.Fd(), "/dev/stderr")

	// Force threads that panic to write to stderr (the panicFile handle now), otherwise it will go into the ether
	log.SetOutput(os.Stderr)

	return nil
}

func removePanicFile() {
	if st, err := panicFile.Stat(); err == nil {
		if st.Size() == 0 {
			sh := uint32(windows.STD_ERROR_HANDLE)
			setStdHandle.Call(uintptr(sh), uintptr(oldStderr))
			panicFile.Close()
			os.Remove(panicFile.Name())
		}
	}
}
<|endoftext|>