<commit_before>package syslog

import (
	"os"
	"strconv"
	"testing"
)

func TestSyslogRetryCount(t *testing.T) {
	setRetryCount()
	if retryCount != defaultRetryCount {
		t.Errorf("expected %v got %v", defaultRetryCount, retryCount)
	}

	newRetryCount := uint(20)
	os.Setenv("RETRY_COUNT", strconv.Itoa(int(newRetryCount)))
	defer os.Unsetenv("RETRY_COUNT")
	setRetryCount()
	if retryCount != newRetryCount {
		t.Errorf("expected %v got %v", newRetryCount, retryCount)
	}
}
<commit_msg>add test for reconnect()<commit_after>package syslog

import (
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"testing"
	"text/template"
	"time"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/gliderlabs/logspout/router"

	_ "github.com/gliderlabs/logspout/transports/tcp"
	_ "github.com/gliderlabs/logspout/transports/tls"
	_ "github.com/gliderlabs/logspout/transports/udp"
)

const (
	closeOnMsgIdx = 5
	maxMsgCount   = 10
	testPriority  = "{{.Priority}}"
	testTimestamp = "{{.Timestamp}}"
	testHostname  = "{{.Container.Config.Hostname}}"
	testTag       = "{{.ContainerName}}"
	testPid       = "{{.Container.State.Pid}}"
	testData      = "{{.Data}}"
)

var (
	container = &docker.Container{
		ID:   "8dfafdbc3a40",
		Name: "0michaelshobbs",
		Config: &docker.Config{
			Hostname: "8dfafdbc3a40",
		},
	}
	testTmplStr = fmt.Sprintf("<%s>%s %s %s[%s]: %s\n",
		testPriority, testTimestamp, testHostname, testTag, testPid, testData)
)

type localTCPServer struct {
	lnmu sync.RWMutex
	net.Listener
}

func (ls *localTCPServer) teardown() error {
	ls.lnmu.Lock()
	if ls.Listener != nil {
		ls.Listener.Close()
		ls.Listener = nil
	}
	ls.lnmu.Unlock()
	return nil
}

func TestSyslogRetryCount(t *testing.T) {
	newRetryCount := uint(20)
	os.Setenv("RETRY_COUNT", strconv.Itoa(int(newRetryCount)))
	setRetryCount()
	if retryCount != newRetryCount {
		t.Errorf("expected %v got %v", newRetryCount, retryCount)
	}

	os.Unsetenv("RETRY_COUNT")
	setRetryCount()
	if retryCount != defaultRetryCount {
		t.Errorf("expected %v got %v", defaultRetryCount, retryCount)
	}
}

func TestSyslogReconnectOnClose(t *testing.T) {
	os.Setenv("RETRY_COUNT", strconv.Itoa(int(1)))
	setRetryCount()
	defer func() {
		os.Unsetenv("RETRY_COUNT")
		setRetryCount()
	}()
	tmpl, err := template.New("syslog").Parse(testTmplStr)
	if err != nil {
		t.Fatal(err)
	}
	ls, err := newLocalTCPServer()
	if err != nil {
		t.Fatal(err)
	}
	defer ls.teardown()

	route := &router.Route{
		ID:      "0",
		Adapter: "syslog",
		Address: ls.Listener.Addr().String(),
	}
	transport, found := router.AdapterTransports.Lookup(route.AdapterTransport("tcp"))
	if !found {
		t.Errorf("bad transport: " + route.Adapter)
	}

	datac := make(chan []byte, maxMsgCount)
	errc := make(chan error, 1)
	go acceptAndCloseConn(ls, datac, errc)

	// Dial connection for adapter
	conn, err := net.Dial(ls.Listener.Addr().Network(), ls.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}

	adapter := &Adapter{
		route:     route,
		conn:      conn,
		tmpl:      tmpl,
		transport: transport,
	}

	logstream := make(chan *router.Message)
	done := make(chan bool)
	sentMsgs := [][]byte{}
	// Send msgs to logstream
	go sendLogstream(logstream, done, &sentMsgs, tmpl)

	// Stream logstream to conn
	go adapter.Stream(logstream)

	for err := range errc {
		t.Errorf("%v", err)
	}

	readMsgs := [][]byte{}
	for {
		select {
		case <-done:
			if maxMsgCount-1 != len(datac) {
				t.Errorf("expected %v got %v", maxMsgCount-1, len(datac))
			}
			for msg := range datac {
				readMsgs = append(readMsgs, msg)
			}
			sentMsgs = append(sentMsgs[:closeOnMsgIdx], sentMsgs[closeOnMsgIdx+1:]...)
			for i, v := range sentMsgs {
				sent := strings.Trim(fmt.Sprintf("%s", v), "\n")
				read := strings.Trim(fmt.Sprintf("%s", readMsgs[i]), "\x00\n")
				if sent != read {
					t.Errorf("expected %+q got %+q", sent, read)
				}
			}
		}
		return
	}
}

func newLocalTCPServer() (*localTCPServer, error) {
	ln, err := newLocalListener()
	if err != nil {
		return nil, err
	}
	return &localTCPServer{Listener: ln}, nil
}

func newLocalListener() (net.Listener, error) {
	ln, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	return ln, nil
}

func acceptAndCloseConn(ls *localTCPServer, datac chan []byte, errc chan error) {
	defer func() {
		close(datac)
		close(errc)
	}()
	c, err := ls.Accept()
	if err != nil {
		errc <- err
		return
	}
	count := 0
	for {
		switch count {
		case maxMsgCount - closeOnMsgIdx:
			c.Close()
			c, err = ls.Accept()
			if err != nil {
				errc <- err
				return
			}
			c.SetReadDeadline(time.Now().Add(5 * time.Second))
			readConn(c, datac)
			count++
		case maxMsgCount:
			return
		default:
			readConn(c, datac)
			count++
		}
	}
}

func readConn(c net.Conn, ch chan []byte) error {
	b := make([]byte, 256)
	_, err := c.Read(b)
	if err != nil {
		return err
	}
	ch <- b
	return nil
}

func sendLogstream(logstream chan *router.Message, done chan bool, msgs *[][]byte, tmpl *template.Template) {
	defer func() {
		close(logstream)
		close(done)
	}()
	var count int
	for {
		if count == maxMsgCount {
			done <- true
			return
		}
		msg := &router.Message{
			Container: container,
			Data:      "hellooo",
			Time:      time.Now(),
		}
		m := &Message{msg}
		buf, _ := m.Render(tmpl)
		*msgs = append(*msgs, buf)
		logstream <- msg
		count++
		time.Sleep(250 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>package chug_commands

import (
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/veritas/say"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/chug"
)

func Prettify(relativeTime string, data string, src io.Reader) error {
	out := make(chan chug.Entry)
	go chug.Chug(src, out)

	if data != "none" && data != "short" && data != "long" {
		return fmt.Errorf("invalid data specification: %s", data)
	}

	s := &stenographer{
		Data: data,
	}

	switch relativeTime {
	case "":
		s.Absolute = true
	case "now":
		s.RelativeTime = time.Now()
	case "first":
		break
	default:
		seconds, err := strconv.ParseFloat(relativeTime, 64)
		if err != nil {
			return fmt.Errorf("invalid relative time specification: %s", relativeTime)
		}
		s.RelativeTime = time.Unix(0, int64(seconds*1e9))
	}

	for entry := range out {
		s.PrettyPrint(entry)
	}

	return nil
}

type stenographer struct {
	Absolute     bool
	RelativeTime time.Time
	Data         string
}

func (s *stenographer) PrettyPrint(entry chug.Entry) {
	if !s.Absolute && s.RelativeTime.IsZero() && entry.IsLager {
		s.RelativeTime = entry.Log.Timestamp
	}

	if entry.IsLager {
		s.PrettyPrintLog(entry.Log)
	} else {
		s.PrettyPrintRaw(entry.Raw)
	}
}

func (s *stenographer) PrettyPrintRaw(raw []byte) {
	say.Println(0, say.Gray(string(raw)))
}

func (s *stenographer) PrettyPrintLog(log chug.LogEntry) {
	components := []string{}

	color := say.DefaultStyle
	ornamentColor := say.GrayColor
	level := "I"

	switch log.LogLevel {
	case lager.INFO:
	case lager.DEBUG:
		color = say.GrayColor
		level = "D"
	case lager.ERROR:
		color = say.RedColor
		ornamentColor = say.RedColor
		level = "E"
	case lager.FATAL:
		color = say.RedColor
		ornamentColor = say.RedColor
		level = "F"
	}

	var timestamp string
	components = append(components, say.Colorize(color, level))
	if s.Absolute {
		timestamp = log.Timestamp.Format(time.StampMilli)
		components = append(components, say.Colorize(color, timestamp))
	} else {
		timestamp = log.Timestamp.Sub(s.RelativeTime).String()
		components = append(components, say.Colorize(color, "%12s", timestamp))
	}
	components = append(components, fmt.Sprintf("%-16s", "["+log.Source+"]"))

	session := log.Session()
	components = append(components, say.Colorize(ornamentColor, "%-10s", session))

	components = append(components, say.Colorize(color, log.Message()))

	if log.Error != nil {
		components = append(components, say.Red(" - Error: "+log.Error.Error()))
	}

	if len(log.Data) > 0 && s.Data == "short" {
		dataJSON, _ := json.Marshal(log.Data)
		components = append(components, say.Colorize(color, string(dataJSON)))
	}

	say.Println(0, strings.Join(components, " "))

	if len(log.Data) > 0 && s.Data == "long" {
		dataJSON, _ := json.MarshalIndent(log.Data, "", "  ")
		say.Println(2, say.Colorize(color, string(dataJSON)))
	}
}
<commit_msg>more chug<commit_after>package chug_commands

import (
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/veritas/say"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/chug"
)

func Prettify(relativeTime string, data string, src io.Reader) error {
	out := make(chan chug.Entry)
	go chug.Chug(src, out)

	if data != "none" && data != "short" && data != "long" {
		return fmt.Errorf("invalid data specification: %s", data)
	}

	s := &stenographer{
		Data: data,
	}

	switch relativeTime {
	case "":
		s.Absolute = true
	case "now":
		s.RelativeTime = time.Now()
	case "first":
		break
	default:
		seconds, err := strconv.ParseFloat(relativeTime, 64)
		if err != nil {
			return fmt.Errorf("invalid relative time specification: %s", relativeTime)
		}
		s.RelativeTime = time.Unix(0, int64(seconds*1e9))
	}

	for entry := range out {
		s.PrettyPrint(entry)
	}

	return nil
}

type stenographer struct {
	Absolute     bool
	RelativeTime time.Time
	Data         string
}

func (s *stenographer) PrettyPrint(entry chug.Entry) {
	if !s.Absolute && s.RelativeTime.IsZero() && entry.IsLager {
		s.RelativeTime = entry.Log.Timestamp
	}

	if entry.IsLager {
		s.PrettyPrintLog(entry.Log)
	} else {
		s.PrettyPrintRaw(entry.Raw)
	}
}

func (s *stenographer) PrettyPrintRaw(raw []byte) {
	say.Println(0, say.Gray(string(raw)))
}

func (s *stenographer) PrettyPrintLog(log chug.LogEntry) {
	components := []string{}

	color := say.DefaultStyle
	level := ""

	switch log.LogLevel {
	case lager.INFO:
		level = say.Green("[INFO] ")
	case lager.DEBUG:
		level = say.Gray("[DEBUG]")
	case lager.ERROR:
		color = say.RedColor
		level = say.Red("[ERROR]")
	case lager.FATAL:
		color = say.RedColor
		level = say.Red("[FATAL]")
	}

	var timestamp string
	components = append(components, fmt.Sprintf("%-16s", "["+log.Source+"]"))
	components = append(components, level)
	if s.Absolute {
		timestamp = log.Timestamp.Format(time.StampMilli)
		components = append(components, say.Colorize(color, timestamp))
	} else {
		timestamp = log.Timestamp.Sub(s.RelativeTime).String()
		components = append(components, say.Colorize(color, "%12s", timestamp))
	}

	components = append(components, say.Gray("%-10s", log.Session))

	components = append(components, say.Colorize(color, log.Message))

	if log.Error != nil {
		components = append(components, say.Red(" - Error: "+log.Error.Error()))
	}

	say.Println(0, strings.Join(components, " "))

	if len(log.Data) > 0 && s.Data == "short" {
		dataJSON, _ := json.Marshal(log.Data)
		say.Println(28, say.Colorize(color, string(dataJSON)))
	}

	if len(log.Data) > 0 && s.Data == "long" {
		dataJSON, _ := json.MarshalIndent(log.Data, "", "  ")
		say.Println(28, say.Colorize(color, string(dataJSON)))
	}
}
<|endoftext|>


<commit_before>package main

type extensionHandler func(*ARWServer, *User, ARWObject)

type ExtensionRequest struct {
	cmd     string
	handler func(*ARWServer, *User, ARWObject)
}
<commit_msg>Handler rename<commit_after>package main

type ExtensionHandler func(*ARWServer, *User, ARWObject)

type ExtensionRequest struct {
	cmd     string
	handler func(*ARWServer, *User, ARWObject)
}
<|endoftext|>


<commit_before>package device

import (
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/Comcast/webpa-common/convey/conveyhttp"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/wrp"
	"github.com/Comcast/webpa-common/xhttp"
	"github.com/go-kit/kit/log"
	"github.com/gorilla/websocket"
)

const MaxDevicesHeader = "X-Xmidt-Max-Devices"

var authStatus *websocket.PreparedMessage

func init() {
	var err error
	authStatus, err = websocket.NewPreparedMessage(
		websocket.BinaryMessage,
		wrp.MustEncode(&wrp.AuthorizationStatus{Status: wrp.AuthStatusAuthorized}, wrp.Msgpack),
	)

	if err != nil {
		panic(err)
	}
}

// Connector is a strategy interface for managing device connections to a server.
// Implementations are responsible for upgrading websocket connections and providing
// for explicit disconnection.
type Connector interface {
	// Connect upgrade an HTTP connection to a websocket and begins concurrent
	// management of the device.
	Connect(http.ResponseWriter, *http.Request, http.Header) (Interface, error)

	// Disconnect disconnects the device associated with the given id.
	// If the id was found, this method returns true.
	Disconnect(ID) bool

	// DisconnectIf iterates over all devices known to this manager, applying the
	// given predicate.  For any devices that result in true, this method disconnects them.
	// Note that this method may pause connections and disconnections while it is executing.
	// This method returns the number of devices that were disconnected.
	//
	// Only disconnection by ID is supported, which means that any identifier matching
	// the predicate will result in *all* duplicate devices under that ID being removed.
	//
	// No methods on this Manager should be called from within the predicate function, or
	// a deadlock will likely occur.
	DisconnectIf(func(ID) bool) int

	// DisconnectAll disconnects all devices from this instance, and returns the count of
	// devices disconnected.
	DisconnectAll() int
}

// Router handles dispatching messages to devices.
type Router interface {
	// Route dispatches a WRP request to exactly one device, identified by the ID
	// field of the request.  Route is synchronous, and honors the cancellation semantics
	// of the Request's context.
	Route(*Request) (*Response, error)
}

// Registry is the strategy interface for querying the set of connected devices.  Methods
// in this interface follow the Visitor pattern and are typically executed under a read lock.
type Registry interface {
	// Get returns the device associated with the given ID, if any
	Get(ID) (Interface, bool)

	// VisitAll applies the given visitor function to each device known to this manager.
	//
	// No methods on this Manager should be called from within the visitor function, or
	// a deadlock will likely occur.
	VisitAll(func(Interface)) int
}

// Manager supplies a hub for connecting and disconnecting devices as well as
// an access point for obtaining device metadata.
type Manager interface {
	Connector
	Router
	Registry
}

// NewManager constructs a Manager from a set of options.  A ConnectionFactory will be
// created from the options if one is not supplied.
func NewManager(o *Options) Manager {
	var (
		logger   = o.logger()
		measures = NewMeasures(o.metricsProvider())
	)

	return &manager{
		logger:   logger,
		errorLog: logging.Error(logger),
		debugLog: logging.Debug(logger),

		readDeadline:     NewDeadline(o.idlePeriod(), o.now()),
		writeDeadline:    NewDeadline(o.writeTimeout(), o.now()),
		upgrader:         o.upgrader(),
		conveyTranslator: conveyhttp.NewHeaderTranslator("", nil),
		devices: newRegistry(registryOptions{
			Logger:   logger,
			Limit:    o.maxDevices(),
			Measures: measures,
		}),
		deviceMessageQueueSize: o.deviceMessageQueueSize(),
		pingPeriod:             o.pingPeriod(),
		authDelay:              o.authDelay(),

		listeners: o.listeners(),
		measures:  measures,
	}
}

// manager is the internal Manager implementation.
type manager struct {
	logger   log.Logger
	errorLog log.Logger
	debugLog log.Logger

	readDeadline     func() time.Time
	writeDeadline    func() time.Time
	upgrader         *websocket.Upgrader
	conveyTranslator conveyhttp.HeaderTranslator

	devices *registry

	deviceMessageQueueSize int
	pingPeriod             time.Duration
	authDelay              time.Duration

	listeners []Listener
	measures  Measures
}

func (m *manager) Connect(response http.ResponseWriter, request *http.Request, responseHeader http.Header) (Interface, error) {
	m.debugLog.Log(logging.MessageKey(), "device connect", "url", request.URL)
	id, ok := GetID(request.Context())
	if !ok {
		xhttp.WriteError(
			response,
			http.StatusInternalServerError,
			ErrorMissingDeviceNameContext,
		)

		return nil, ErrorMissingDeviceNameContext
	}

	d := newDevice(deviceOptions{ID: id, QueueSize: m.deviceMessageQueueSize, Logger: m.logger})
	if convey, err := m.conveyTranslator.FromHeader(request.Header); err == nil {
		d.infoLog.Log("convey", convey)
	} else if err != conveyhttp.ErrMissingHeader {
		d.errorLog.Log(logging.MessageKey(), "badly formatted convey data", logging.ErrorKey(), err)
	}

	c, err := m.upgrader.Upgrade(response, request, responseHeader)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "failed websocket upgrade", logging.ErrorKey(), err)
		return nil, err
	}

	d.debugLog.Log(logging.MessageKey(), "websocket upgrade complete", "localAddress", c.LocalAddr().String())

	pinger, err := NewPinger(c, m.measures.Ping, []byte(d.ID()), m.writeDeadline)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to create pinger", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	if err := m.devices.add(d); err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to register device", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	m.dispatch(
		&Event{
			Type:   Connect,
			Device: d,
		},
	)

	SetPongHandler(c, m.measures.Pong, m.readDeadline)
	closeOnce := new(sync.Once)
	go m.readPump(d, InstrumentReader(c, d.statistics), closeOnce)
	go m.writePump(d, InstrumentWriter(c, d.statistics), pinger, closeOnce)

	return d, nil
}

func (m *manager) dispatch(e *Event) {
	for _, listener := range m.listeners {
		listener(e)
	}
}

// pumpClose handles the proper shutdown and logging of a device's pumps.
// This method should be executed within a sync.Once, so that it only executes
// once for a given device.
//
// Note that the write pump does additional cleanup.  In particular, the write pump
// dispatches message failed events for any messages that were waiting to be delivered
// at the time of pump closure.
func (m *manager) pumpClose(d *device, c io.Closer, pumpError error) {
	if pumpError != nil {
		d.errorLog.Log(logging.MessageKey(), "pump close due to error", logging.ErrorKey(), pumpError)
	} else {
		d.debugLog.Log(logging.MessageKey(), "pump close")
	}

	// remove will invoke requestClose()
	m.devices.remove(d.id)

	if closeError := c.Close(); closeError != nil {
		d.errorLog.Log(logging.MessageKey(), "Error closing device connection", logging.ErrorKey(), closeError)
	} else {
		d.debugLog.Log(logging.MessageKey(), "Closed device connection")
	}

	m.dispatch(
		&Event{
			Type:   Disconnect,
			Device: d,
		},
	)
}

// readPump is the goroutine which handles the stream of WRP messages from a device.
// This goroutine exits when any error occurs on the connection.
func (m *manager) readPump(d *device, r ReadCloser, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "readPump exiting")
	d.debugLog.Log(logging.MessageKey(), "readPump starting")

	var (
		readError error
		decoder   = wrp.NewDecoder(nil, wrp.Msgpack)
	)

	// all the read pump has to do is ensure the device and the connection are closed
	// it is the write pump's responsibility to do further cleanup
	defer closeOnce.Do(func() { m.pumpClose(d, r, readError) })

	for {
		messageType, data, readError := r.ReadMessage()
		if readError != nil {
			d.errorLog.Log(logging.MessageKey(), "read error", logging.ErrorKey(), readError)
			return
		}

		if messageType != websocket.BinaryMessage {
			d.errorLog.Log(logging.MessageKey(), "skipping non-binary frame", "messageType", messageType)
			continue
		}

		var (
			message = new(wrp.Message)
			event   = Event{
				Type:     MessageReceived,
				Device:   d,
				Message:  message,
				Format:   wrp.Msgpack,
				Contents: data,
			}
		)

		decoder.ResetBytes(data)
		err := decoder.Decode(message)
		decoder.ResetBytes(nil)
		if err != nil {
			d.errorLog.Log(logging.MessageKey(), "skipping malformed WRP message", logging.ErrorKey(), err)
			continue
		}

		if message.Type == wrp.SimpleRequestResponseMessageType {
			m.measures.RequestResponse.Add(1.0)
		}

		// update any waiting transaction
		if message.IsTransactionPart() {
			err := d.transactions.Complete(
				message.TransactionKey(),
				&Response{
					Device:   d,
					Message:  message,
					Format:   wrp.Msgpack,
					Contents: data,
				},
			)

			if err != nil {
				d.errorLog.Log(logging.MessageKey(), "Error while completing transaction", logging.ErrorKey(), err)
				event.Type = TransactionBroken
				event.Error = err
			} else {
				event.Type = TransactionComplete
			}
		}

		m.dispatch(&event)
	}
}

// writePump is the goroutine which services messages addressed to the device.
// this goroutine exits when either an explicit shutdown is requested or any
// error occurs on the connection.
func (m *manager) writePump(d *device, w WriteCloser, pinger func() error, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "writePump exiting")
	d.debugLog.Log(logging.MessageKey(), "writePump starting")

	var (
		envelope   *envelope
		encoder    = wrp.NewEncoder(nil, wrp.Msgpack)
		writeError error

		pingTicker = time.NewTicker(m.pingPeriod)

		// wait for the delay, then send an auth status request to the device
		authStatusTimer = time.AfterFunc(m.authDelay, func() {
			// TODO: This will keep the device from being garbage collected until the timer
			// triggers.  This is only a problem if a device connects then disconnects faster
			// than the authDelay setting.
			w.WritePreparedMessage(authStatus)
		})
	)

	// cleanup: we not only ensure that the device and connection are closed but also
	// ensure that any messages that were waiting and/or failed are dispatched to
	// the configured listener
	defer func() {
		pingTicker.Stop()
		authStatusTimer.Stop()
		closeOnce.Do(func() { m.pumpClose(d, w, writeError) })

		// notify listener of any message that just now failed
		// any writeError is passed via this event
		if envelope != nil {
			m.dispatch(&Event{
				Type:     MessageFailed,
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			})
		}

		// drain the messages, dispatching them as message failed events.  we never close
		// the message channel, so just drain until a receive would block.
		//
		// Nil is passed explicitly as the error to indicate that these messages failed due
		// to the device disconnecting, not due to an actual I/O error.
		for {
			select {
			case undeliverable := <-d.messages:
				d.errorLog.Log(logging.MessageKey(), "undeliverable message", "deviceMessage", undeliverable)
				m.dispatch(&Event{
					Type:     MessageFailed,
					Device:   d,
					Message:  undeliverable.request.Message,
					Format:   undeliverable.request.Format,
					Contents: undeliverable.request.Contents,
					Error:    writeError,
				})
			default:
				return
			}
		}
	}()

	for writeError == nil {
		envelope = nil

		select {
		case <-d.shutdown:
			d.debugLog.Log(logging.MessageKey(), "explicit shutdown")
			writeError = w.Close()
			return

		case envelope = <-d.messages:
			var frameContents []byte
			if envelope.request.Format == wrp.Msgpack && len(envelope.request.Contents) > 0 {
				frameContents = envelope.request.Contents
			} else {
				// if the request was in a format other than Msgpack, or if the caller did not pass
				// Contents, then do the encoding here.
				encoder.ResetBytes(&frameContents)
				writeError = encoder.Encode(envelope.request.Message)
				encoder.ResetBytes(nil)
			}

			if writeError == nil {
				writeError = w.WriteMessage(websocket.BinaryMessage, frameContents)
			}

			event := Event{
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			}

			if writeError != nil {
				envelope.complete <- writeError
				event.Type = MessageFailed
			} else {
				event.Type = MessageSent
			}

			close(envelope.complete)
			m.dispatch(&event)

		case <-pingTicker.C:
			writeError = pinger()
		}
	}
}

func (m *manager) Disconnect(id ID) bool {
	_, ok := m.devices.remove(id)
	return ok
}

func (m *manager) DisconnectIf(filter func(ID) bool) int {
	return m.devices.removeIf(func(d *device) bool {
		return filter(d.id)
	})
}

func (m *manager) DisconnectAll() int {
	return m.devices.removeAll()
}

func (m *manager) Get(id ID) (Interface, bool) {
	return m.devices.get(id)
}

func (m *manager) VisitAll(visitor func(Interface)) int {
	return m.devices.visit(func(d *device) {
		visitor(d)
	})
}

func (m *manager) Route(request *Request) (*Response, error) {
	if destination, err := request.ID(); err != nil {
		return nil, err
	} else if d, ok := m.devices.get(destination); ok {
		return d.Send(request)
	} else {
		return nil, ErrorDeviceNotFound
	}
}
<commit_msg>Update the final log entry for a device connection.<commit_after>package device

import (
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/Comcast/webpa-common/convey/conveyhttp"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/wrp"
	"github.com/Comcast/webpa-common/xhttp"
	"github.com/go-kit/kit/log"
	"github.com/gorilla/websocket"
)

const MaxDevicesHeader = "X-Xmidt-Max-Devices"

var authStatus *websocket.PreparedMessage

func init() {
	var err error
	authStatus, err = websocket.NewPreparedMessage(
		websocket.BinaryMessage,
		wrp.MustEncode(&wrp.AuthorizationStatus{Status: wrp.AuthStatusAuthorized}, wrp.Msgpack),
	)

	if err != nil {
		panic(err)
	}
}

// Connector is a strategy interface for managing device connections to a server.
// Implementations are responsible for upgrading websocket connections and providing
// for explicit disconnection.
type Connector interface {
	// Connect upgrade an HTTP connection to a websocket and begins concurrent
	// management of the device.
	Connect(http.ResponseWriter, *http.Request, http.Header) (Interface, error)

	// Disconnect disconnects the device associated with the given id.
	// If the id was found, this method returns true.
	Disconnect(ID) bool

	// DisconnectIf iterates over all devices known to this manager, applying the
	// given predicate.  For any devices that result in true, this method disconnects them.
	// Note that this method may pause connections and disconnections while it is executing.
	// This method returns the number of devices that were disconnected.
	//
	// Only disconnection by ID is supported, which means that any identifier matching
	// the predicate will result in *all* duplicate devices under that ID being removed.
	//
	// No methods on this Manager should be called from within the predicate function, or
	// a deadlock will likely occur.
	DisconnectIf(func(ID) bool) int

	// DisconnectAll disconnects all devices from this instance, and returns the count of
	// devices disconnected.
	DisconnectAll() int
}

// Router handles dispatching messages to devices.
type Router interface {
	// Route dispatches a WRP request to exactly one device, identified by the ID
	// field of the request.  Route is synchronous, and honors the cancellation semantics
	// of the Request's context.
	Route(*Request) (*Response, error)
}

// Registry is the strategy interface for querying the set of connected devices.  Methods
// in this interface follow the Visitor pattern and are typically executed under a read lock.
type Registry interface {
	// Get returns the device associated with the given ID, if any
	Get(ID) (Interface, bool)

	// VisitAll applies the given visitor function to each device known to this manager.
	//
	// No methods on this Manager should be called from within the visitor function, or
	// a deadlock will likely occur.
	VisitAll(func(Interface)) int
}

// Manager supplies a hub for connecting and disconnecting devices as well as
// an access point for obtaining device metadata.
type Manager interface {
	Connector
	Router
	Registry
}

// NewManager constructs a Manager from a set of options.  A ConnectionFactory will be
// created from the options if one is not supplied.
func NewManager(o *Options) Manager {
	var (
		logger   = o.logger()
		measures = NewMeasures(o.metricsProvider())
	)

	return &manager{
		logger:   logger,
		errorLog: logging.Error(logger),
		debugLog: logging.Debug(logger),

		readDeadline:     NewDeadline(o.idlePeriod(), o.now()),
		writeDeadline:    NewDeadline(o.writeTimeout(), o.now()),
		upgrader:         o.upgrader(),
		conveyTranslator: conveyhttp.NewHeaderTranslator("", nil),
		devices: newRegistry(registryOptions{
			Logger:   logger,
			Limit:    o.maxDevices(),
			Measures: measures,
		}),
		deviceMessageQueueSize: o.deviceMessageQueueSize(),
		pingPeriod:             o.pingPeriod(),
		authDelay:              o.authDelay(),

		listeners: o.listeners(),
		measures:  measures,
	}
}

// manager is the internal Manager implementation.
type manager struct {
	logger   log.Logger
	errorLog log.Logger
	debugLog log.Logger

	readDeadline     func() time.Time
	writeDeadline    func() time.Time
	upgrader         *websocket.Upgrader
	conveyTranslator conveyhttp.HeaderTranslator

	devices *registry

	deviceMessageQueueSize int
	pingPeriod             time.Duration
	authDelay              time.Duration

	listeners []Listener
	measures  Measures
}

func (m *manager) Connect(response http.ResponseWriter, request *http.Request, responseHeader http.Header) (Interface, error) {
	m.debugLog.Log(logging.MessageKey(), "device connect", "url", request.URL)
	id, ok := GetID(request.Context())
	if !ok {
		xhttp.WriteError(
			response,
			http.StatusInternalServerError,
			ErrorMissingDeviceNameContext,
		)

		return nil, ErrorMissingDeviceNameContext
	}

	d := newDevice(deviceOptions{ID: id, QueueSize: m.deviceMessageQueueSize, Logger: m.logger})
	if convey, err := m.conveyTranslator.FromHeader(request.Header); err == nil {
		d.infoLog.Log("convey", convey)
	} else if err != conveyhttp.ErrMissingHeader {
		d.errorLog.Log(logging.MessageKey(), "badly formatted convey data", logging.ErrorKey(), err)
	}

	c, err := m.upgrader.Upgrade(response, request, responseHeader)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "failed websocket upgrade", logging.ErrorKey(), err)
		return nil, err
	}

	d.debugLog.Log(logging.MessageKey(), "websocket upgrade complete", "localAddress", c.LocalAddr().String())

	pinger, err := NewPinger(c, m.measures.Ping, []byte(d.ID()), m.writeDeadline)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to create pinger", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	if err := m.devices.add(d); err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to register device", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	m.dispatch(
		&Event{
			Type:   Connect,
			Device: d,
		},
	)

	SetPongHandler(c, m.measures.Pong, m.readDeadline)
	closeOnce := new(sync.Once)
	go m.readPump(d, InstrumentReader(c, d.statistics), closeOnce)
	go m.writePump(d, InstrumentWriter(c, d.statistics), pinger, closeOnce)

	return d, nil
}

func (m *manager) dispatch(e *Event) {
	for _, listener := range m.listeners {
		listener(e)
	}
}

// pumpClose handles the proper shutdown and logging of a device's pumps.
// This method should be executed within a sync.Once, so that it only executes
// once for a given device.
//
// Note that the write pump does additional cleanup.  In particular, the write pump
// dispatches message failed events for any messages that were waiting to be delivered
// at the time of pump closure.
func (m *manager) pumpClose(d *device, c io.Closer, pumpError error) {
	// remove will invoke requestClose()
	m.devices.remove(d.id)

	closeError := c.Close()

	d.errorLog.Log(logging.MessageKey(), "Closed device connection",
		"closeError", closeError, "pumpError", pumpError,
		"finalStatistics", d.Statistics().String())

	m.dispatch(
		&Event{
			Type:   Disconnect,
			Device: d,
		},
	)
}

// readPump is the goroutine which handles the stream of WRP messages from a device.
// This goroutine exits when any error occurs on the connection.
func (m *manager) readPump(d *device, r ReadCloser, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "readPump exiting")
	d.debugLog.Log(logging.MessageKey(), "readPump starting")

	var (
		readError error
		decoder   = wrp.NewDecoder(nil, wrp.Msgpack)
	)

	// all the read pump has to do is ensure the device and the connection are closed
	// it is the write pump's responsibility to do further cleanup
	defer closeOnce.Do(func() { m.pumpClose(d, r, readError) })

	for {
		messageType, data, readError := r.ReadMessage()
		if readError != nil {
			d.errorLog.Log(logging.MessageKey(), "read error", logging.ErrorKey(), readError)
			return
		}

		if messageType != websocket.BinaryMessage {
			d.errorLog.Log(logging.MessageKey(), "skipping non-binary frame", "messageType", messageType)
			continue
		}

		var (
			message = new(wrp.Message)
			event   = Event{
				Type:     MessageReceived,
				Device:   d,
				Message:  message,
				Format:   wrp.Msgpack,
				Contents: data,
			}
		)

		decoder.ResetBytes(data)
		err := decoder.Decode(message)
		decoder.ResetBytes(nil)
		if err != nil {
			d.errorLog.Log(logging.MessageKey(), "skipping malformed WRP message", logging.ErrorKey(), err)
			continue
		}

		if message.Type == wrp.SimpleRequestResponseMessageType {
			m.measures.RequestResponse.Add(1.0)
		}

		// update any waiting transaction
		if message.IsTransactionPart() {
			err := d.transactions.Complete(
				message.TransactionKey(),
				&Response{
					Device:   d,
					Message:  message,
					Format:   wrp.Msgpack,
					Contents: data,
				},
			)

			if err != nil {
				d.errorLog.Log(logging.MessageKey(), "Error while completing transaction", logging.ErrorKey(), err)
				event.Type = TransactionBroken
				event.Error = err
			} else {
				event.Type = TransactionComplete
			}
		}

		m.dispatch(&event)
	}
}

// writePump is the goroutine which services messages addressed to the device.
// this goroutine exits when either an explicit shutdown is requested or any
// error occurs on the connection.
func (m *manager) writePump(d *device, w WriteCloser, pinger func() error, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "writePump exiting")
	d.debugLog.Log(logging.MessageKey(), "writePump starting")

	var (
		envelope   *envelope
		encoder    = wrp.NewEncoder(nil, wrp.Msgpack)
		writeError error

		pingTicker = time.NewTicker(m.pingPeriod)

		// wait for the delay, then send an auth status request to the device
		authStatusTimer = time.AfterFunc(m.authDelay, func() {
			// TODO: This will keep the device from being garbage collected until the timer
			// triggers.  This is only a problem if a device connects then disconnects faster
			// than the authDelay setting.
			w.WritePreparedMessage(authStatus)
		})
	)

	// cleanup: we not only ensure that the device and connection are closed but also
	// ensure that any messages that were waiting and/or failed are dispatched to
	// the configured listener
	defer func() {
		pingTicker.Stop()
		authStatusTimer.Stop()
		closeOnce.Do(func() { m.pumpClose(d, w, writeError) })

		// notify listener of any message that just now failed
		// any writeError is passed via this event
		if envelope != nil {
			m.dispatch(&Event{
				Type:     MessageFailed,
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			})
		}

		// drain the messages, dispatching them as message failed events.  we never close
		// the message channel, so just drain until a receive would block.
		//
		// Nil is passed explicitly as the error to indicate that these messages failed due
		// to the device disconnecting, not due to an actual I/O error.
		for {
			select {
			case undeliverable := <-d.messages:
				d.errorLog.Log(logging.MessageKey(), "undeliverable message", "deviceMessage", undeliverable)
				m.dispatch(&Event{
					Type:     MessageFailed,
					Device:   d,
					Message:  undeliverable.request.Message,
					Format:   undeliverable.request.Format,
					Contents: undeliverable.request.Contents,
					Error:    writeError,
				})
			default:
				return
			}
		}
	}()

	for writeError == nil {
		envelope = nil

		select {
		case <-d.shutdown:
			d.debugLog.Log(logging.MessageKey(), "explicit shutdown")
			writeError = w.Close()
			return

		case envelope = <-d.messages:
			var frameContents []byte
			if envelope.request.Format == wrp.Msgpack && len(envelope.request.Contents) > 0 {
				frameContents = envelope.request.Contents
			} else {
				// if the request was in a format other than Msgpack, or if the caller did not pass
				// Contents, then do the encoding here.
				encoder.ResetBytes(&frameContents)
				writeError = encoder.Encode(envelope.request.Message)
				encoder.ResetBytes(nil)
			}

			if writeError == nil {
				writeError = w.WriteMessage(websocket.BinaryMessage, frameContents)
			}

			event := Event{
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			}

			if writeError != nil {
				envelope.complete <- writeError
				event.Type = MessageFailed
			} else {
				event.Type = MessageSent
			}

			close(envelope.complete)
			m.dispatch(&event)

		case <-pingTicker.C:
			writeError = pinger()
		}
	}
}

func (m *manager) Disconnect(id ID) bool {
	_, ok := m.devices.remove(id)
	return ok
}

func (m *manager) DisconnectIf(filter func(ID) bool) int {
	return m.devices.removeIf(func(d *device) bool {
		return filter(d.id)
	})
}

func (m *manager) DisconnectAll() int {
	return m.devices.removeAll()
}

func (m *manager) Get(id ID) (Interface, bool) {
	return m.devices.get(id)
}

func (m *manager) VisitAll(visitor func(Interface)) int {
	return m.devices.visit(func(d *device) {
		visitor(d)
	})
}

func (m *manager) Route(request *Request) (*Response, error) {
	if destination, err := request.ID(); err != nil {
		return nil, err
	} else if d, ok := m.devices.get(destination); ok {
		return d.Send(request)
	} else {
		return nil, ErrorDeviceNotFound
	}
}
<|endoftext|>


<commit_before>package goadb

import (
	"log"
	"runtime"
	"strings"
	"sync/atomic"

	"github.com/zach-klippenstein/goadb/util"
	"github.com/zach-klippenstein/goadb/wire"
)

/*
DeviceWatcher publishes device status change events.
If the server dies while listening for events, it restarts the server.
*/
type DeviceWatcher struct {
	*deviceWatcherImpl
}

// DeviceStateChangedEvent represents a device state transition.
type DeviceStateChangedEvent struct {
	Serial   string
	OldState DeviceState
	NewState DeviceState
}

// CameOnline returns true if this event represents a device coming online.
func (s DeviceStateChangedEvent) CameOnline() bool {
	return s.OldState != StateOnline && s.NewState == StateOnline
}

// WentOffline returns true if this event represents a device going offline.
func (s DeviceStateChangedEvent) WentOffline() bool {
	return s.OldState == StateOnline && s.NewState != StateOnline
}

// DeviceState represents one of the 3 possible states adb will report devices.
// A device can be communicated with when it's in StateOnline.
// A USB device will transition from StateDisconnected->StateOffline->StateOnline when
// plugged in, and then StateOnline->StateDisconnected when unplugged.
// If code doesn't care about specific states, DeviceStateChangedEvent provides methods
// to query at a higher level.
//go:generate stringer -type=DeviceState
type DeviceState int8

const (
	StateDisconnected DeviceState = iota
	StateOffline
	StateOnline
)

var deviceStateStrings = map[string]DeviceState{
	"":        StateDisconnected,
	"offline": StateOffline,
	"device":  StateOnline,
}

type deviceWatcherImpl struct {
	config ClientConfig

	// If an error occurs, it is stored here and eventChan is close immediately after.
	err atomic.Value

	eventChan chan DeviceStateChangedEvent

	// Function to start the server if it's not running or dies.
	startServer func() error
}

func NewDeviceWatcher(config ClientConfig) *DeviceWatcher {
	watcher := &DeviceWatcher{&deviceWatcherImpl{
		config:      config.sanitized(),
		eventChan:   make(chan DeviceStateChangedEvent),
		startServer: StartServer,
	}}

	runtime.SetFinalizer(watcher, func(watcher *DeviceWatcher) {
		watcher.Shutdown()
	})

	go publishDevices(watcher.deviceWatcherImpl)

	return watcher
}

/*
C returns a channel than can be received on to get events.
If an unrecoverable error occurs, or Shutdown is called, the channel will be closed.
*/
func (w *DeviceWatcher) C() <-chan DeviceStateChangedEvent {
	return w.eventChan
}

// Err returns the error that caused the channel returned by C to be closed, if C is closed.
// If C is not closed, its return value is undefined.
func (w *DeviceWatcher) Err() error {
	if err, ok := w.err.Load().(error); ok {
		return err
	}
	return nil
}

// Shutdown stops the watcher from listening for events and closes the channel returned
// from C.
func (w *DeviceWatcher) Shutdown() {
	// TODO(z): Implement.
}

func (w *deviceWatcherImpl) reportErr(err error) {
	w.err.Store(err)
}

/*
publishDevices reads device lists from scanner, calculates diffs, and publishes events on
eventChan.
Returns when scanner returns an error.
Doesn't refer directly to a *DeviceWatcher so it can be GCed (which will,
in turn, close Scanner and stop this goroutine).

TODO: to support shutdown, spawn a new goroutine each time a server connection is established.
This goroutine should read messages and send them to a message channel. Can write errors directly
to errVal. publisHDevicesUntilError should take the msg chan and the scanner and select on the msg chan and stop chan, and if the stop
chan sends, close the scanner and return true. If the msg chan closes, just return false.
publishDevices can look at ret val: if false and err == EOF, reconnect. If false and other error, report err
and abort. If true, report no error and stop.
*/
func publishDevices(watcher *deviceWatcherImpl) {
	defer close(watcher.eventChan)

	var lastKnownStates map[string]DeviceState
	finished := false

	for {
		scanner, err := connectToTrackDevices(watcher.config.Dialer)
		if err != nil {
			watcher.reportErr(err)
			return
		}

		finished, err = publishDevicesUntilError(scanner, watcher.eventChan, &lastKnownStates)

		if finished {
			scanner.Close()
			return
		}

		if util.HasErrCode(err, util.ConnectionResetError) {
			// The server died, restart and reconnect.
			log.Println("[DeviceWatcher] server died, restarting…")
			if err := watcher.startServer(); err != nil {
				log.Println("[DeviceWatcher] error restarting server, giving up")
				watcher.reportErr(err)
				return
			} // Else server should be running, continue listening.
		} else {
			// Unknown error, don't retry.
			watcher.reportErr(err)
			return
		}
	}
}

func connectToTrackDevices(dialer Dialer) (wire.Scanner, error) {
	conn, err := dialer.Dial()
	if err != nil {
		return nil, err
	}

	if err := wire.SendMessageString(conn, "host:track-devices"); err != nil {
		conn.Close()
		return nil, err
	}

	if err := wire.ReadStatusFailureAsError(conn, "host:track-devices"); err != nil {
		conn.Close()
		return nil, err
	}

	return conn, nil
}

func publishDevicesUntilError(scanner wire.Scanner, eventChan chan<- DeviceStateChangedEvent, lastKnownStates *map[string]DeviceState) (finished bool, err error) {
	for {
		msg, err := scanner.ReadMessage()
		if err != nil {
			return false, err
		}

		deviceStates, err := parseDeviceStates(string(msg))
		if err != nil {
			return false, err
		}

		for _, event := range calculateStateDiffs(*lastKnownStates, deviceStates) {
			eventChan <- event
		}
		*lastKnownStates = deviceStates
	}
}

func parseDeviceStates(msg string) (states map[string]DeviceState, err error) {
	states = make(map[string]DeviceState)

	for lineNum, line := range strings.Split(msg, "\n") {
		if len(line) == 0 {
			continue
		}

		fields := strings.Split(line, "\t")
		if len(fields) != 2 {
			err = util.Errorf(util.ParseError, "invalid device state line %d: %s", lineNum, line)
			return
		}

		serial, stateString := fields[0], fields[1]
		state, ok := deviceStateStrings[stateString]
		if !ok {
			err = util.Errorf(util.ParseError, "invalid device state: %s", state)
		}
		states[serial] = state
	}

	return
}

func calculateStateDiffs(oldStates, newStates map[string]DeviceState) (events []DeviceStateChangedEvent) {
	for serial, oldState := range oldStates {
		newState, ok := newStates[serial]

		if oldState != newState {
			if ok {
				// Device present in both lists: state changed.
				events = append(events, DeviceStateChangedEvent{serial, oldState, newState})
			} else {
				// Device only present in old list: device removed.
				events = append(events, DeviceStateChangedEvent{serial, oldState, StateDisconnected})
			}
		}
	}

	for serial, newState := range newStates {
		if _, ok := oldStates[serial]; !ok {
			// Device only present in new list: device added.
			events = append(events, DeviceStateChangedEvent{serial, StateDisconnected, newState})
		}
	}

	return events
}
<commit_msg>Randomly backoff restarting adb server from DeviceWatcher.<commit_after>package goadb

import (
	"log"
	"runtime"
	"strings"
	"sync/atomic"
	"math/rand"
	"time"

	"github.com/zach-klippenstein/goadb/util"
	"github.com/zach-klippenstein/goadb/wire"
)

/*
DeviceWatcher publishes device status change events.
If the server dies while listening for events, it restarts the server.
*/
type DeviceWatcher struct {
	*deviceWatcherImpl
}

// DeviceStateChangedEvent represents a device state transition.
type DeviceStateChangedEvent struct {
	Serial   string
	OldState DeviceState
	NewState DeviceState
}

// CameOnline returns true if this event represents a device coming online.
func (s DeviceStateChangedEvent) CameOnline() bool {
	return s.OldState != StateOnline && s.NewState == StateOnline
}

// WentOffline returns true if this event represents a device going offline.
func (s DeviceStateChangedEvent) WentOffline() bool {
	return s.OldState == StateOnline && s.NewState != StateOnline
}

// DeviceState represents one of the 3 possible states adb will report devices.
// A device can be communicated with when it's in StateOnline.
// A USB device will transition from StateDisconnected->StateOffline->StateOnline when
// plugged in, and then StateOnline->StateDisconnected when unplugged.
// If code doesn't care about specific states, DeviceStateChangedEvent provides methods
// to query at a higher level.
//go:generate stringer -type=DeviceState
type DeviceState int8

const (
	StateDisconnected DeviceState = iota
	StateOffline
	StateOnline
)

var deviceStateStrings = map[string]DeviceState{
	"":        StateDisconnected,
	"offline": StateOffline,
	"device":  StateOnline,
}

type deviceWatcherImpl struct {
	config ClientConfig

	// If an error occurs, it is stored here and eventChan is close immediately after.
	err atomic.Value

	eventChan chan DeviceStateChangedEvent

	// Function to start the server if it's not running or dies.
	startServer func() error
}

func NewDeviceWatcher(config ClientConfig) *DeviceWatcher {
	watcher := &DeviceWatcher{&deviceWatcherImpl{
		config:      config.sanitized(),
		eventChan:   make(chan DeviceStateChangedEvent),
		startServer: StartServer,
	}}

	runtime.SetFinalizer(watcher, func(watcher *DeviceWatcher) {
		watcher.Shutdown()
	})

	go publishDevices(watcher.deviceWatcherImpl)

	return watcher
}

/*
C returns a channel than can be received on to get events.
If an unrecoverable error occurs, or Shutdown is called, the channel will be closed.
*/
func (w *DeviceWatcher) C() <-chan DeviceStateChangedEvent {
	return w.eventChan
}

// Err returns the error that caused the channel returned by C to be closed, if C is closed.
// If C is not closed, its return value is undefined.
func (w *DeviceWatcher) Err() error {
	if err, ok := w.err.Load().(error); ok {
		return err
	}
	return nil
}

// Shutdown stops the watcher from listening for events and closes the channel returned
// from C.
func (w *DeviceWatcher) Shutdown() {
	// TODO(z): Implement.
}

func (w *deviceWatcherImpl) reportErr(err error) {
	w.err.Store(err)
}

/*
publishDevices reads device lists from scanner, calculates diffs, and publishes events on
eventChan.
Returns when scanner returns an error.
Doesn't refer directly to a *DeviceWatcher so it can be GCed (which will,
in turn, close Scanner and stop this goroutine).

TODO: to support shutdown, spawn a new goroutine each time a server connection is established.
This goroutine should read messages and send them to a message channel. Can write errors directly
to errVal. publishDevicesUntilError should take the msg chan and the scanner and select on the msg chan and stop chan, and if the stop
chan sends, close the scanner and return true. If the msg chan closes, just return false.
publishDevices can look at ret val: if false and err == EOF, reconnect. If false and other error, report err
and abort. If true, report no error and stop.
*/
func publishDevices(watcher *deviceWatcherImpl) {
	defer close(watcher.eventChan)

	var lastKnownStates map[string]DeviceState
	finished := false

	for {
		scanner, err := connectToTrackDevices(watcher.config.Dialer)
		if err != nil {
			watcher.reportErr(err)
			return
		}

		finished, err = publishDevicesUntilError(scanner, watcher.eventChan, &lastKnownStates)

		if finished {
			scanner.Close()
			return
		}

		if util.HasErrCode(err, util.ConnectionResetError) {
			// The server died, restart and reconnect.

			// Delay by a random [0ms, 500ms) in case multiple DeviceWatchers are trying to
			// start the same server.
			delay := time.Duration(rand.Intn(500)) * time.Millisecond

			log.Printf("[DeviceWatcher] server died, restarting in %s…", delay)
			time.Sleep(delay)
			if err := watcher.startServer(); err != nil {
				log.Println("[DeviceWatcher] error restarting server, giving up")
				watcher.reportErr(err)
				return
			} // Else server should be running, continue listening.
		} else {
			// Unknown error, don't retry.
			watcher.reportErr(err)
			return
		}
	}
}

func connectToTrackDevices(dialer Dialer) (wire.Scanner, error) {
	conn, err := dialer.Dial()
	if err != nil {
		return nil, err
	}

	if err := wire.SendMessageString(conn, "host:track-devices"); err != nil {
		conn.Close()
		return nil, err
	}

	if err := wire.ReadStatusFailureAsError(conn, "host:track-devices"); err != nil {
		conn.Close()
		return nil, err
	}

	return conn, nil
}

func publishDevicesUntilError(scanner wire.Scanner, eventChan chan<- DeviceStateChangedEvent, lastKnownStates *map[string]DeviceState) (finished bool, err error) {
	for {
		msg, err := scanner.ReadMessage()
		if err != nil {
			return false, err
		}

		deviceStates, err := parseDeviceStates(string(msg))
		if err != nil {
			return false, err
		}

		for _, event := range calculateStateDiffs(*lastKnownStates, deviceStates) {
			eventChan <- event
		}
		*lastKnownStates = deviceStates
	}
}

func parseDeviceStates(msg string) (states map[string]DeviceState, err error) {
	states = make(map[string]DeviceState)

	for lineNum, line := range strings.Split(msg, "\n") {
		if len(line) == 0 {
			continue
		}

		fields := strings.Split(line, "\t")
		if len(fields) != 2 {
			err = util.Errorf(util.ParseError, "invalid device state line %d: %s", lineNum, line)
			return
		}

		serial, stateString := fields[0], fields[1]
		state, ok := deviceStateStrings[stateString]
		if !ok {
			err = util.Errorf(util.ParseError, "invalid device state: %s", state)
		}
		states[serial] = state
	}

	return
}

func calculateStateDiffs(oldStates, newStates map[string]DeviceState) (events []DeviceStateChangedEvent) {
	for serial, oldState := range oldStates {
		newState, ok := newStates[serial]

		if oldState != newState {
			if ok {
				// Device present in both lists: state changed.
				events = append(events, DeviceStateChangedEvent{serial, oldState, newState})
			} else {
				// Device only present in old list: device removed.
				events = append(events, DeviceStateChangedEvent{serial, oldState, StateDisconnected})
			}
		}
	}

	for serial, newState := range newStates {
		if _, ok := oldStates[serial]; !ok {
			// Device only present in new list: device added.
			events = append(events, DeviceStateChangedEvent{serial, StateDisconnected, newState})
		}
	}

	return events
}
<|endoftext|>


<commit_before>// Deploy implements the actual deployment process for successfully build jobs.
package deployment

import (
	"leeroy/config"
	"leeroy/logging"
	"log"
	"os/exec"
)

// Deploy a job if all tests are passed.
func Deploy(j *logging.Job, c *config.Config) {
	if j.Success() != true {
		log.Println("Not deploying", j.Branch, "build did not succeed.")
		return
	}

	r, err := c.ConfigForRepo(j.URL)

	if err != nil {
		log.Println("Repo", j.URL, "does not exist.")
		return
	}

	d, err := r.DeployTarget(j.Branch)

	if err != nil {
		log.Println("Deployment target for", j.Branch, "does not exist")
		return
	}

	announceStart(j, c, &d)

	o, err := call(d.Execute, r.URL, j.Branch)
	j.Deployed = o

	if err != nil {
		log.Println(err.Error())
	}

	//announceComplete(j, c, &d, o)
}

// Log and announce a started deployment.
func announceStart(j *logging.Job, c *config.Config, d *config.Deploy) {
	log.Println(j.Name, "triggered a deploy for", j.Branch, "to", d.Name)
}

// Log and announce a completed deployment.
func announceComplete(j *logging.Job, c *config.Config, d *config.Deploy, output string) {
	log.Println("finished deploying", j.Branch, "to", d.Name, "with", output)
}

// Call a deployment script and return the output.
func call(app string, repo string, branch string) (string, error) {
	cmd := exec.Command(app, repo, branch)
	out, err := cmd.CombinedOutput()
	return string(out), err
}
<commit_msg>enable 'deploy complete announcements' again<commit_after>// Deploy implements the actual deployment process for successfully build jobs.
package deployment

import (
	"leeroy/config"
	"leeroy/logging"
	"log"
	"os/exec"
)

// Deploy a job if all tests are passed.
func Deploy(j *logging.Job, c *config.Config) {
	if j.Success() != true {
		log.Println("Not deploying", j.Branch, "build did not succeed.")
		return
	}

	r, err := c.ConfigForRepo(j.URL)

	if err != nil {
		log.Println("Repo", j.URL, "does not exist.")
		return
	}

	d, err := r.DeployTarget(j.Branch)

	if err != nil {
		log.Println("Deployment target for", j.Branch, "does not exist")
		return
	}

	announceStart(j, c, &d)

	o, err := call(d.Execute, r.URL, j.Branch)
	j.Deployed = o

	if err != nil {
		log.Println(err.Error())
	}

	announceComplete(j, c, &d, o)
}

// Log and announce a started deployment.
func announceStart(j *logging.Job, c *config.Config, d *config.Deploy) {
	log.Println(j.Name, "triggered a deploy for", j.Branch, "to", d.Name)
}

// Log and announce a completed deployment.
func announceComplete(j *logging.Job, c *config.Config, d *config.Deploy, output string) {
	log.Println("finished deploying", j.Branch, "to", d.Name, "with", output)
}

// Call a deployment script and return the output.
func call(app string, repo string, branch string) (string, error) {
	cmd := exec.Command(app, repo, branch)
	out, err := cmd.CombinedOutput()
	return string(out), err
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 disk implements a blob store on the local disk.
package disk

import (
	"github.com/jacobsa/comeback/blob"
	"io/ioutil"
	"path"
)

type blobStore struct {
	basePath string
}

func (s *blobStore) Store(b []byte) (blob.Score, error) {
	score := blob.ComputeScore(b)
	hexScore := blob.HexScore(score)
	filePath := path.Join(s.basePath, hexScore)

	if err := ioutil.WriteFile(filePath, b, 0600); err != nil {
		return nil, err
	}

	return score, nil
}

func (s *blobStore) Load(score blob.Score) ([]byte, error) {
	hexScore := blob.HexScore(score)
	filePath := path.Join(s.basePath, hexScore)

	return ioutil.ReadFile(filePath)
}

// Return a blob store that stores its blobs in the directory with the supplied
// path.
func NewBlobStore(path string) (s blob.Store, err error) {
	return &blobStore{basePath: path}, nil
}
<commit_msg>Made disk build again.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 disk implements a blob store on the local disk.
package disk

import (
	"github.com/jacobsa/comeback/blob"
	"io/ioutil"
	"path"
)

type blobStore struct {
	basePath string
}

func (s *blobStore) Store(b []byte) (blob.Score, error) {
	score := blob.ComputeScore(b)
	filePath := path.Join(s.basePath, score.Hex())

	if err := ioutil.WriteFile(filePath, b, 0600); err != nil {
		return nil, err
	}

	return score, nil
}

func (s *blobStore) Load(score blob.Score) ([]byte, error) {
	filePath := path.Join(s.basePath, score.Hex())

	return ioutil.ReadFile(filePath)
}

// Return a blob store that stores its blobs in the directory with the supplied
// path.
func NewBlobStore(path string) (s blob.Store, err error) {
	return &blobStore{basePath: path}, nil
}
<|endoftext|>


<commit_before>package GameServer

import (
	"NetSession"
	"fmt"
)

type ServerManager struct {
}

/*
register gameserver
*/
func (manager *ServerManager) RegisterServer(session *NetSession.NetSession, server interface{}) {

}

/*
send all server to client
*/
func (manager *ServerManager) SendServerList(session *NetSession.NetSession) {

}
<commit_msg>remove package frm form ServerManager.go<commit_after>package GameServer

import (
	"NetSession"
	//"fmt"
)

type ServerManager struct {
}

/*
register gameserver
*/
func (manager *ServerManager) RegisterServer(session *NetSession.NetSession, server interface{}) {

}

/*
send all server to client
*/
func (manager *ServerManager) SendServerList(session *NetSession.NetSession) {

}
<|endoftext|>


<commit_before>package main

import (
        "encoding/json"
        "fmt"
        "os"
        "os/exec"
        "log"
        "regexp"
        "sync"
)

// The configuration file struct
type Configuration struct {
        DriveSyncDirectory string
        GoogleDriveRemoteDirectory string
        HugoPostDirectory string
}

// Read the configuration JSON file in order to get some settings and directories
func read_cfg(filename string, wg *sync.WaitGroup, conf_message chan string) {
        file, _ := os.Open(filename)
        decoder := json.NewDecoder(file)
        configuration := Configuration{}
        err := decoder.Decode(&configuration)
        if err != nil {
                fmt.Println("error: ", err)
        }
        drive_sync_dir := fmt.Sprintf(configuration.DriveSyncDirectory)
        conf_message <- drive_sync_dir
        drive_remote_dir := fmt.Sprintf(configuration.GoogleDriveRemoteDirectory)
        conf_message <- drive_remote_dir
        hugo_post_dir := fmt.Sprintf(configuration.HugoPostDirectory)
        conf_message <- hugo_post_dir
        wg.Done()
}

// Sync google drive remote folder to the configured local directory.
// Then send the output from drive CLI to a function to intepret the output
// by stripping the full output down to an array of string paths to docx files.
func sync_google_drive(sync_dir string, drive_remote_dir string, wg *sync.WaitGroup) {
        sync_gd := new(sync.WaitGroup)
        output := make(chan string)
        sync := exec.Command("drive pull -no-prompt -desktop-links=false -export docx", drive_remote_dir)
        sync.Dir = sync_dir
        fmt.Println("Syncing Google Drive...")
        out, err := sync.Output()
        if err != nil {
                fmt.Println("[ERROR] Error syncing Google Drive: ", err)
        }
        fmt.Println("Done syncing!")
        sync_gd.Add(1)
        go intepret_drive_output(sync_gd, output)
        output <- string(out)
        sync_gd.Wait()
        wg.Done()
}

// Find all Exported file paths via a regex expression and then add them to an array
func interpret_drive_output(sync_gd *sync.WaitGroup, output chan string) {
        results := <-output
        re := regexp.MustCompile(`to '(.*?)'`)
        matches := re.FindAllString(results, -1)
        sync_gd.Done()
}

// Convert from docx to markdown with pandoc
func convert_to_markdown_with_pandoc(docx_file_path string, md_file_path string, wg *sync.WaitGroup) {
        convert := exec.Command("pandoc --atx-headers --smart --normalize --email-obfuscation=references --mathjax -t markdown_strict -o", md_file_path, docx_file_path)
        out, err := convert.Output()
        if err != nil {
                fmt.Println("[ERROR] Error converting files to markdown with pandoc: ", err)
        }
        wg.Done()
}

func compile_and_serve_hugo_site(hugo_dir string, prod_dir string, use_hugo bool, wg *sync.WaitGroup) {
}

func main() {
        conf_message := make(chan string)
        wg := new(sync.WaitGroup)
        wg.Add(2)
        go exe_cmd("echo", "ping", wg)
        go read_cfg("conf.json", wg, conf_message)
        drive_sync_dir := <-conf_message
        fmt.Println(drive_sync_dir)
        drive_remote_dir := <-conf_message
        fmt.Println(drive_remote_dir)
        hugo_post_dir := <-conf_message
        fmt.Println(hugo_post_dir)
        wg.Wait()
}
<commit_msg>Adding today's progress<commit_after>package main

import (
        "crypt/md5"
        "encoding/json"
        "encoding/hex"
        "fmt"
        "os"
        "os/exec"
        "log"
        "regexp"
        "sync"
)

//The mutex lock for reading/writing to the hashtable
var counter = struct {
        sync.RWMutex
        hashtable map[string]string
}{hashtable: make(map[string]string)}

// The configuration file struct
type Configuration struct {
        DriveSyncDirectory string
        GoogleDriveRemoteDirectory string
        HugoPostDirectory string
}

// Read the configuration JSON file in order to get some settings and directories
func read_cfg(filename string, wg *sync.WaitGroup, conf_message chan string) {
        file, _ := os.Open(filename)
        decoder := json.NewDecoder(file)
        configuration := Configuration{}
        err := decoder.Decode(&configuration)
        if err != nil {
                fmt.Println("[ERROR] Error reading the JSON confguration: ", err)
                return
        }
        drive_sync_dir := fmt.Sprintf(configuration.DriveSyncDirectory)
        conf_message <- drive_sync_dir
        drive_remote_dir := fmt.Sprintf(configuration.GoogleDriveRemoteDirectory)
        conf_message <- drive_remote_dir
        hugo_post_dir := fmt.Sprintf(configuration.HugoPostDirectory)
        conf_message <- hugo_post_dir
        wg.Done()
}

// Use md5 hash sums for the filepaths
func md5hash(text string, DriveSyncDirectory string) string {
        r := strings.NewReplacer(DriveSyncDirectory, "")
        relative-path := r.Replace(text)
        hasher := md5.New()
        hasher.Write([]byte(relative-path))
        return hex.EncodeToString(hasher.Sum(nil))
}

// Write the filepath string to the hashtable
func write_path_to_hashtable(path string, hashtable map[string]string) {
        key := md5hash(path)
        counter.Lock()
        counter.m[key]path
        counter.Unlock()
}

// Lookup a path in the hashtable, if it exists return true, otherwise false
func check_for_path(path string, hashtable map[string]string, read *sync.WaitGroup) bool {
        key := md5hash(path)
        counter.RLock()
        read_path := counter.hashtable[key]
        if (read_path == path) {
                return true
        }
        return false
}

// Write the hashtable to a json file in order to keep a backup in case of system reboot
func write_hashtable_to_json(hashtable map[string]string) {
        bytes, err := json.Marshal(hashtable)
        if err != nil {
                fmt.Println("[ERROR] Error writing hashtable to JSON: ", err)
                return
        }
        text := string(bytes)
        fmt.Println(bytes)
}

// Read the saved hashtable from the saved json file to restart the syncing after reboot
func read_hashtable_from_json(hashtable map[string]string, table_dir string) {
}

// Sync google drive remote folder to the configured local directory.
// Then send the output from drive CLI to a function to intepret the output
// by stripping the full output down to an array of string paths to docx files.
func sync_google_drive(sync_dir string, drive_remote_dir string, wg *sync.WaitGroup) {
        sync_gd := new(sync.WaitGroup)
        output := make(chan string)
        sync := exec.Command("drive pull -no-prompt -desktop-links=false -export docx", drive_remote_dir)
        sync.Dir = sync_dir
        fmt.Println("Syncing Google Drive...")
        out, err := sync.Output()
        if err != nil {
                fmt.Println("[ERROR] Error syncing Google Drive: ", err)
                return
        }
        fmt.Println("Done syncing!")
        sync_gd.Add(1)
        go intepret_drive_output(sync_gd, output)
        output <- string(out)
        sync_gd.Wait()
        wg.Done()
}

// Find all Exported file paths via a regex expression and then add them to an array
func interpret_drive_output(sync_gd *sync.WaitGroup, output chan string) {
        results := <-output
        re := regexp.MustCompile(`to '(.*?)'`)
        matches := re.FindAllString(results, -1)
        sync_gd.Done()
}

// Convert from docx to markdown with pandoc
func convert_to_markdown_with_pandoc(docx_file_path string, md_file_path string, wg *sync.WaitGroup) {
        convert := exec.Command("pandoc --atx-headers --smart --normalize --email-obfuscation=references --mathjax -t markdown_strict -o", md_file_path, docx_file_path)
        out, err := convert.Output()
        if err != nil {
                fmt.Println("[ERROR] Error converting files to markdown with pandoc: ", err)
        }
        wg.Done()
}

// Use hugo to compile the markdown files into html and then serve with hugo or with nginx
func compile_and_serve_hugo_site(hugo_dir string, prod_dir string, use_hugo bool, wg *sync.WaitGroup) {
}

/* func main() {
        conf_message := make(chan string)
        wg := new(sync.WaitGroup)
        wg.Add(2)
        go exe_cmd("echo", "ping", wg)
        go read_cfg("conf.json", wg, conf_message)
        drive_sync_dir := <-conf_message
        fmt.Println(drive_sync_dir)
        drive_remote_dir := <-conf_message
        fmt.Println(drive_remote_dir)
        hugo_post_dir := <-conf_message
        fmt.Println(hugo_post_dir)
        wg.Wait()
} */
<|endoftext|>


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

package libkbfs

import (
	"fmt"
	"sync"
	"sync/atomic"

	lru "github.com/hashicorp/golang-lru"
	"github.com/keybase/kbfs/kbfsblock"
	"github.com/keybase/kbfs/kbfshash"
	"github.com/keybase/kbfs/tlf"
)

type blockContainer struct {
	block         Block
	hasPrefetched bool
}

type idCacheKey struct {
	tlf           tlf.ID
	plaintextHash kbfshash.RawDefaultHash
}

// BlockCacheStandard implements the BlockCache interface by storing
// blocks in an in-memory LRU cache.  Clean blocks are identified
// internally by just their block ID (since blocks are immutable and
// content-addressable).
type BlockCacheStandard struct {
	cleanBytesCapacity uint64

	ids *lru.Cache

	cleanTransient *lru.Cache

	cleanLock      sync.RWMutex
	cleanPermanent map[kbfsblock.ID]Block

	bytesLock       sync.Mutex
	cleanTotalBytes uint64
}

// NewBlockCacheStandard constructs a new BlockCacheStandard instance
// with the given transient capacity (in number of entries) and the
// clean bytes capacity, which is the total of number of bytes allowed
// between the transient and permanent clean caches.  If putting a
// block will exceed this bytes capacity, transient entries are
// evicted until the block will fit in capacity.
func NewBlockCacheStandard(transientCapacity int,
	cleanBytesCapacity uint64) *BlockCacheStandard {
	b := &BlockCacheStandard{
		cleanBytesCapacity: cleanBytesCapacity,
		cleanPermanent:     make(map[kbfsblock.ID]Block),
	}

	if transientCapacity > 0 {
		var err error
		// TODO: Plumb error up.
		b.ids, err = lru.New(transientCapacity)
		if err != nil {
			return nil
		}

		b.cleanTransient, err = lru.NewWithEvict(transientCapacity, b.onEvict)
		if err != nil {
			return nil
		}
	}
	return b
}

// GetWithPrefetch implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) GetWithPrefetch(ptr BlockPointer) (Block, bool, error) {
	if b.cleanTransient != nil {
		if tmp, ok := b.cleanTransient.Get(ptr.ID); ok {
			bc, ok := tmp.(blockContainer)
			if !ok {
				return nil, false, BadDataError{ptr.ID}
			}
			return bc.block, bc.hasPrefetched, nil
		}
	}

	block := func() Block {
		b.cleanLock.RLock()
		defer b.cleanLock.RUnlock()
		return b.cleanPermanent[ptr.ID]
	}()
	if block != nil {
		// A permanent entry can only be created if this client is performing a
		// write. Since the client is writing, it knows what goes into it,
		// including any potential directory entries or indirect blocks.
		// Thus, it is treated as already prefetched.
		return block, true, nil
	}

	return nil, false, NoSuchBlockError{ptr.ID}
}

// Get implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) Get(ptr BlockPointer) (Block, error) {
	block, _, err := b.GetWithPrefetch(ptr)
	return block, err
}

func getCachedBlockSize(block Block) uint32 {
	// Get the size of the block.  For direct file blocks, use the
	// length of the plaintext contents.  For everything else, just
	// approximate the plaintext size using the encoding size.
	switch b := block.(type) {
	case *FileBlock:
		if b.IsInd {
			return b.GetEncodedSize()
		}
		return uint32(len(b.Contents))
	default:
		return block.GetEncodedSize()
	}
}

func (b *BlockCacheStandard) onEvict(key interface{}, value interface{}) {
	bc, ok := value.(blockContainer)
	if !ok {
		return
	}
	block := bc.block

	b.bytesLock.Lock()
	defer b.bytesLock.Unlock()
	b.cleanTotalBytes -= uint64(getCachedBlockSize(block))
}

// CheckForKnownPtr implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) CheckForKnownPtr(tlf tlf.ID, block *FileBlock) (
	BlockPointer, error) {
	if block.IsInd {
		return BlockPointer{}, NotDirectFileBlockError{}
	}

	if b.ids == nil {
		return BlockPointer{}, nil
	}

	key := idCacheKey{tlf, block.UpdateHash()}
	tmp, ok := b.ids.Get(key)
	if !ok {
		return BlockPointer{}, nil
	}

	ptr, ok := tmp.(BlockPointer)
	if !ok {
		return BlockPointer{}, fmt.Errorf("Unexpected cached id: %v", tmp)
	}
	return ptr, nil
}

// SetCleanBytesCapacity implements the BlockCache interface for
// BlockCacheStandard.
func (b *BlockCacheStandard) SetCleanBytesCapacity(capacity uint64) {
	atomic.StoreUint64(&b.cleanBytesCapacity, capacity)
}

// GetCleanBytesCapacity implements the BlockCache interface for
// BlockCacheStandard.
func (b *BlockCacheStandard) GetCleanBytesCapacity() (capacity uint64) {
	return atomic.LoadUint64(&b.cleanBytesCapacity)
}

func (b *BlockCacheStandard) makeRoomForSize(size uint64) bool {
	if b.cleanTransient == nil {
		return false
	}

	oldLen := b.cleanTransient.Len() + 1
	doUnlock := true
	b.bytesLock.Lock()
	defer func() {
		if doUnlock {
			b.bytesLock.Unlock()
		}
	}()

	cleanBytesCapacity := b.GetCleanBytesCapacity()

	// Evict items from the cache until the bytes capacity is lower
	// than the total capacity (or until no items are removed).
	for b.cleanTotalBytes+size > cleanBytesCapacity {
		// Unlock while removing, since onEvict needs the lock and
		// cleanTransient.Len() takes the LRU mutex (which could lead
		// to a deadlock with onEvict).  TODO: either change
		// `cleanTransient` into an `lru.SimpleLRU` and protect it
		// with our own lock, or build our own LRU that can evict
		// based on total bytes.  See #250 and KBFS-1404 for a longer
		// discussion.
		b.bytesLock.Unlock()
		doUnlock = false
		if oldLen == b.cleanTransient.Len() {
			break
		}
		oldLen = b.cleanTransient.Len()
		b.cleanTransient.RemoveOldest()
		doUnlock = true
		b.bytesLock.Lock()
	}
	if b.cleanTotalBytes+size > cleanBytesCapacity {
		// There must be too many permanent clean blocks, so we
		// couldn't make room.
		return false
	}
	// Only count clean bytes if we actually have a transient cache.
	b.cleanTotalBytes += size
	return true
}

// PutWithPrefetch implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) PutWithPrefetch(
	ptr BlockPointer, tlf tlf.ID, block Block, lifetime BlockCacheLifetime,
	hasPrefetched bool) (err error) {

	_, cacheGetError := b.Get(ptr)
	wasInCache := cacheGetError == nil
	madeRoom := false

	switch lifetime {
	case NoCacheEntry:
		return nil

	case TransientEntry:
		// If it's the right type of block, store the hash -> ID mapping.
		if fBlock, isFileBlock := block.(*FileBlock); b.ids != nil && isFileBlock && !fBlock.IsInd {

			key := idCacheKey{tlf, fBlock.GetHash()}
			// zero out the refnonce, it doesn't matter
			ptr.RefNonce = kbfsblock.ZeroRefNonce
			b.ids.Add(key, ptr)
		}
		if b.cleanTransient == nil {
			return nil
		}
		// Cache it later, once we know there's room
		defer func() {
			if madeRoom || wasInCache {
				b.cleanTransient.Add(ptr.ID, blockContainer{block, hasPrefetched})
			} else {
				err = cachePutCacheFullError{ptr}
			}
		}()

	case PermanentEntry:
		func() {
			b.cleanLock.Lock()
			defer b.cleanLock.Unlock()
			b.cleanPermanent[ptr.ID] = block
		}()

	default:
		return fmt.Errorf("Unknown lifetime %v", lifetime)
	}

	// We must make room whether the cache is transient or permanent, but only
	// if it wasn't already in the cache.
	if !wasInCache {
		size := uint64(getCachedBlockSize(block))
		madeRoom = b.makeRoomForSize(size)
	}

	return nil
}

// Put implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) Put(
	ptr BlockPointer, tlf tlf.ID, block Block, lifetime BlockCacheLifetime) error {
	// Default should be to assume that a prefetch has happened, and thus it
	// won't trigger prefetches in the future.
	return b.PutWithPrefetch(ptr, tlf, block, lifetime, true)
}

// DeletePermanent implements the BlockCache interface for
// BlockCacheStandard.
func (b *BlockCacheStandard) DeletePermanent(id kbfsblock.ID) error {
	b.cleanLock.Lock()
	defer b.cleanLock.Unlock()
	block, ok := b.cleanPermanent[id]
	if ok {
		delete(b.cleanPermanent, id)
		b.bytesLock.Lock()
		defer b.bytesLock.Unlock()
		b.cleanTotalBytes -= uint64(getCachedBlockSize(block))
	}
	return nil
}

// DeleteTransient implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) DeleteTransient(
	ptr BlockPointer, tlf tlf.ID) error {
	if b.cleanTransient == nil {
		return nil
	}

	// If the block is cached and a file block, delete the known
	// pointer as well.
	if tmp, ok := b.cleanTransient.Get(ptr.ID); ok {
		bc, ok := tmp.(blockContainer)
		if !ok {
			return BadDataError{ptr.ID}
		}
		block := bc.block

		// Remove the key if it exists
		if fBlock, ok := block.(*FileBlock); b.ids != nil && ok &&
			!fBlock.IsInd {
			_, hash := kbfshash.DoRawDefaultHash(fBlock.Contents)
			key := idCacheKey{tlf, hash}
			b.ids.Remove(key)
		}

		b.cleanTransient.Remove(ptr.ID)
	}
	return nil
}

// DeleteKnownPtr implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) DeleteKnownPtr(tlf tlf.ID, block *FileBlock) error {
	if block.IsInd {
		return NotDirectFileBlockError{}
	}

	if b.ids == nil {
		return nil
	}

	_, hash := kbfshash.DoRawDefaultHash(block.Contents)
	key := idCacheKey{tlf, hash}
	b.ids.Remove(key)
	return nil
}
<commit_msg>bcache: Implement strib's PR feedback to clarify the PutWithPrefetch implementation.<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"fmt"
	"sync"
	"sync/atomic"

	lru "github.com/hashicorp/golang-lru"
	"github.com/keybase/kbfs/kbfsblock"
	"github.com/keybase/kbfs/kbfshash"
	"github.com/keybase/kbfs/tlf"
)

type blockContainer struct {
	block         Block
	hasPrefetched bool
}

type idCacheKey struct {
	tlf           tlf.ID
	plaintextHash kbfshash.RawDefaultHash
}

// BlockCacheStandard implements the BlockCache interface by storing
// blocks in an in-memory LRU cache.  Clean blocks are identified
// internally by just their block ID (since blocks are immutable and
// content-addressable).
type BlockCacheStandard struct {
	cleanBytesCapacity uint64

	ids *lru.Cache

	cleanTransient *lru.Cache

	cleanLock      sync.RWMutex
	cleanPermanent map[kbfsblock.ID]Block

	bytesLock       sync.Mutex
	cleanTotalBytes uint64
}

// NewBlockCacheStandard constructs a new BlockCacheStandard instance
// with the given transient capacity (in number of entries) and the
// clean bytes capacity, which is the total of number of bytes allowed
// between the transient and permanent clean caches.  If putting a
// block will exceed this bytes capacity, transient entries are
// evicted until the block will fit in capacity.
func NewBlockCacheStandard(transientCapacity int,
	cleanBytesCapacity uint64) *BlockCacheStandard {
	b := &BlockCacheStandard{
		cleanBytesCapacity: cleanBytesCapacity,
		cleanPermanent:     make(map[kbfsblock.ID]Block),
	}

	if transientCapacity > 0 {
		var err error
		// TODO: Plumb error up.
		b.ids, err = lru.New(transientCapacity)
		if err != nil {
			return nil
		}

		b.cleanTransient, err = lru.NewWithEvict(transientCapacity, b.onEvict)
		if err != nil {
			return nil
		}
	}
	return b
}

// GetWithPrefetch implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) GetWithPrefetch(ptr BlockPointer) (Block, bool, error) {
	if b.cleanTransient != nil {
		if tmp, ok := b.cleanTransient.Get(ptr.ID); ok {
			bc, ok := tmp.(blockContainer)
			if !ok {
				return nil, false, BadDataError{ptr.ID}
			}
			return bc.block, bc.hasPrefetched, nil
		}
	}

	block := func() Block {
		b.cleanLock.RLock()
		defer b.cleanLock.RUnlock()
		return b.cleanPermanent[ptr.ID]
	}()
	if block != nil {
		// A permanent entry can only be created if this client is performing a
		// write. Since the client is writing, it knows what goes into it,
		// including any potential directory entries or indirect blocks.
		// Thus, it is treated as already prefetched.
		return block, true, nil
	}

	return nil, false, NoSuchBlockError{ptr.ID}
}

// Get implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) Get(ptr BlockPointer) (Block, error) {
	block, _, err := b.GetWithPrefetch(ptr)
	return block, err
}

func getCachedBlockSize(block Block) uint32 {
	// Get the size of the block.  For direct file blocks, use the
	// length of the plaintext contents.  For everything else, just
	// approximate the plaintext size using the encoding size.
	switch b := block.(type) {
	case *FileBlock:
		if b.IsInd {
			return b.GetEncodedSize()
		}
		return uint32(len(b.Contents))
	default:
		return block.GetEncodedSize()
	}
}

func (b *BlockCacheStandard) onEvict(key interface{}, value interface{}) {
	bc, ok := value.(blockContainer)
	if !ok {
		return
	}
	block := bc.block

	b.bytesLock.Lock()
	defer b.bytesLock.Unlock()
	b.cleanTotalBytes -= uint64(getCachedBlockSize(block))
}

// CheckForKnownPtr implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) CheckForKnownPtr(tlf tlf.ID, block *FileBlock) (
	BlockPointer, error) {
	if block.IsInd {
		return BlockPointer{}, NotDirectFileBlockError{}
	}

	if b.ids == nil {
		return BlockPointer{}, nil
	}

	key := idCacheKey{tlf, block.UpdateHash()}
	tmp, ok := b.ids.Get(key)
	if !ok {
		return BlockPointer{}, nil
	}

	ptr, ok := tmp.(BlockPointer)
	if !ok {
		return BlockPointer{}, fmt.Errorf("Unexpected cached id: %v", tmp)
	}
	return ptr, nil
}

// SetCleanBytesCapacity implements the BlockCache interface for
// BlockCacheStandard.
func (b *BlockCacheStandard) SetCleanBytesCapacity(capacity uint64) {
	atomic.StoreUint64(&b.cleanBytesCapacity, capacity)
}

// GetCleanBytesCapacity implements the BlockCache interface for
// BlockCacheStandard.
func (b *BlockCacheStandard) GetCleanBytesCapacity() (capacity uint64) {
	return atomic.LoadUint64(&b.cleanBytesCapacity)
}

func (b *BlockCacheStandard) makeRoomForSize(size uint64) bool {
	if b.cleanTransient == nil {
		return false
	}

	oldLen := b.cleanTransient.Len() + 1
	doUnlock := true
	b.bytesLock.Lock()
	defer func() {
		if doUnlock {
			b.bytesLock.Unlock()
		}
	}()

	cleanBytesCapacity := b.GetCleanBytesCapacity()

	// Evict items from the cache until the bytes capacity is lower
	// than the total capacity (or until no items are removed).
	for b.cleanTotalBytes+size > cleanBytesCapacity {
		// Unlock while removing, since onEvict needs the lock and
		// cleanTransient.Len() takes the LRU mutex (which could lead
		// to a deadlock with onEvict).  TODO: either change
		// `cleanTransient` into an `lru.SimpleLRU` and protect it
		// with our own lock, or build our own LRU that can evict
		// based on total bytes.  See #250 and KBFS-1404 for a longer
		// discussion.
		b.bytesLock.Unlock()
		doUnlock = false
		if oldLen == b.cleanTransient.Len() {
			break
		}
		oldLen = b.cleanTransient.Len()
		b.cleanTransient.RemoveOldest()
		doUnlock = true
		b.bytesLock.Lock()
	}
	if b.cleanTotalBytes+size > cleanBytesCapacity {
		// There must be too many permanent clean blocks, so we
		// couldn't make room.
		return false
	}
	// Only count clean bytes if we actually have a transient cache.
	b.cleanTotalBytes += size
	return true
}

// PutWithPrefetch implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) PutWithPrefetch(
	ptr BlockPointer, tlf tlf.ID, block Block, lifetime BlockCacheLifetime,
	hasPrefetched bool) (err error) {

	var wasInCache bool

	switch lifetime {
	case NoCacheEntry:
		return nil

	case TransientEntry:
		// If it's the right type of block, store the hash -> ID mapping.
		if fBlock, isFileBlock := block.(*FileBlock); b.ids != nil && isFileBlock && !fBlock.IsInd {

			key := idCacheKey{tlf, fBlock.GetHash()}
			// zero out the refnonce, it doesn't matter
			ptr.RefNonce = kbfsblock.ZeroRefNonce
			b.ids.Add(key, ptr)
		}
		if b.cleanTransient == nil {
			return nil
		}
		wasInCache = b.cleanTransient.Contains(ptr.ID)
		// Cache it later, once we know there's room

	case PermanentEntry:
		func() {
			b.cleanLock.Lock()
			defer b.cleanLock.Unlock()
			_, wasInCache = b.cleanPermanent[ptr.ID]
			b.cleanPermanent[ptr.ID] = block
		}()

	default:
		return fmt.Errorf("Unknown lifetime %v", lifetime)
	}

	transientCacheHasRoom := true
	// We must make room whether the cache is transient or permanent, but only
	// if it wasn't already in the cache.
	if !wasInCache {
		size := uint64(getCachedBlockSize(block))
		transientCacheHasRoom = b.makeRoomForSize(size)
	}
	if lifetime == TransientEntry {
		if !transientCacheHasRoom {
			return cachePutCacheFullError{ptr}
		}
		b.cleanTransient.Add(ptr.ID, blockContainer{block, hasPrefetched})
	}

	return nil
}

// Put implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) Put(
	ptr BlockPointer, tlf tlf.ID, block Block, lifetime BlockCacheLifetime) error {
	// Default should be to assume that a prefetch has happened, and thus it
	// won't trigger prefetches in the future.
	return b.PutWithPrefetch(ptr, tlf, block, lifetime, true)
}

// DeletePermanent implements the BlockCache interface for
// BlockCacheStandard.
func (b *BlockCacheStandard) DeletePermanent(id kbfsblock.ID) error {
	b.cleanLock.Lock()
	defer b.cleanLock.Unlock()
	block, ok := b.cleanPermanent[id]
	if ok {
		delete(b.cleanPermanent, id)
		b.bytesLock.Lock()
		defer b.bytesLock.Unlock()
		b.cleanTotalBytes -= uint64(getCachedBlockSize(block))
	}
	return nil
}

// DeleteTransient implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) DeleteTransient(
	ptr BlockPointer, tlf tlf.ID) error {
	if b.cleanTransient == nil {
		return nil
	}

	// If the block is cached and a file block, delete the known
	// pointer as well.
	if tmp, ok := b.cleanTransient.Get(ptr.ID); ok {
		bc, ok := tmp.(blockContainer)
		if !ok {
			return BadDataError{ptr.ID}
		}
		block := bc.block

		// Remove the key if it exists
		if fBlock, ok := block.(*FileBlock); b.ids != nil && ok &&
			!fBlock.IsInd {
			_, hash := kbfshash.DoRawDefaultHash(fBlock.Contents)
			key := idCacheKey{tlf, hash}
			b.ids.Remove(key)
		}

		b.cleanTransient.Remove(ptr.ID)
	}
	return nil
}

// DeleteKnownPtr implements the BlockCache interface for BlockCacheStandard.
func (b *BlockCacheStandard) DeleteKnownPtr(tlf tlf.ID, block *FileBlock) error {
	if block.IsInd {
		return NotDirectFileBlockError{}
	}

	if b.ids == nil {
		return nil
	}

	_, hash := kbfshash.DoRawDefaultHash(block.Contents)
	key := idCacheKey{tlf, hash}
	b.ids.Remove(key)
	return nil
}
<|endoftext|>