<commit_before>package htlcswitch

import (
	"bytes"
	"fmt"

	sphinx "github.com/lightningnetwork/lightning-onion"
	"github.com/lightningnetwork/lnd/htlcswitch/hop"
	"github.com/lightningnetwork/lnd/lnwire"
)

// ForwardingError wraps an lnwire.FailureMessage in a struct that also
// includes the source of the error.
type ForwardingError struct {
	// FailureSourceIdx is the index of the node that sent the failure. With
	// this information, the dispatcher of a payment can modify their set of
	// candidate routes in response to the type of failure extracted. Index
	// zero is the self node.
	FailureSourceIdx int

	// ExtraMsg is an additional error message that callers can provide in
	// order to provide context specific error details.
	ExtraMsg string

	lnwire.FailureMessage
}

// Error implements the built-in error interface. We use this method to allow
// the switch or any callers to insert additional context to the error message
// returned.
func (f *ForwardingError) Error() string {
	if f.ExtraMsg == "" {
		return fmt.Sprintf("%v", f.FailureMessage)
	}

	return fmt.Sprintf("%v: %v", f.FailureMessage, f.ExtraMsg)
}

// ErrorDecrypter is an interface that is used to decrypt the onion encrypted
// failure reason an extra out a well formed error.
type ErrorDecrypter interface {
	// DecryptError peels off each layer of onion encryption from the first
	// hop, to the source of the error. A fully populated
	// lnwire.FailureMessage is returned along with the source of the
	// error.
	DecryptError(lnwire.OpaqueReason) (*ForwardingError, error)
}

// UnknownEncrypterType is an error message used to signal that an unexpected
// EncrypterType was encountered during decoding.
type UnknownEncrypterType hop.EncrypterType

// Error returns a formatted error indicating the invalid EncrypterType.
func (e UnknownEncrypterType) Error() string {
	return fmt.Sprintf("unknown error encrypter type: %d", e)
}

// OnionErrorDecrypter is the interface that provides onion level error
// decryption.
type OnionErrorDecrypter interface {
	// DecryptError attempts to decrypt the passed encrypted error response.
	// The onion failure is encrypted in backward manner, starting from the
	// node where error have occurred. As a result, in order to decrypt the
	// error we need get all shared secret and apply decryption in the
	// reverse order.
	DecryptError(encryptedData []byte) (*sphinx.DecryptedError, error)
}

// SphinxErrorDecrypter wraps the sphinx data SphinxErrorDecrypter and maps the
// returned errors to concrete lnwire.FailureMessage instances.
type SphinxErrorDecrypter struct {
	OnionErrorDecrypter
}

// DecryptError peels off each layer of onion encryption from the first hop, to
// the source of the error. A fully populated lnwire.FailureMessage is returned
// along with the source of the error.
//
// NOTE: Part of the ErrorDecrypter interface.
func (s *SphinxErrorDecrypter) DecryptError(reason lnwire.OpaqueReason) (
	*ForwardingError, error) {

	failure, err := s.OnionErrorDecrypter.DecryptError(reason)
	if err != nil {
		return nil, err
	}

	// Decode the failure. If an error occurs, we leave the failure message
	// field nil.
	r := bytes.NewReader(failure.Message)
	failureMsg, err := lnwire.DecodeFailure(r, 0)
	if err != nil {
		return &ForwardingError{
			FailureSourceIdx: failure.SenderIdx,
		}, nil
	}

	return &ForwardingError{
		FailureSourceIdx: failure.SenderIdx,
		FailureMessage:   failureMsg,
	}, nil
}

// A compile time check to ensure ErrorDecrypter implements the Deobfuscator
// interface.
var _ ErrorDecrypter = (*SphinxErrorDecrypter)(nil)
<commit_msg>htlcswitch: report failure source index for forwarding errors<commit_after>package htlcswitch

import (
	"bytes"
	"fmt"

	sphinx "github.com/lightningnetwork/lightning-onion"
	"github.com/lightningnetwork/lnd/htlcswitch/hop"
	"github.com/lightningnetwork/lnd/lnwire"
)

// ForwardingError wraps an lnwire.FailureMessage in a struct that also
// includes the source of the error.
type ForwardingError struct {
	// FailureSourceIdx is the index of the node that sent the failure. With
	// this information, the dispatcher of a payment can modify their set of
	// candidate routes in response to the type of failure extracted. Index
	// zero is the self node.
	FailureSourceIdx int

	// ExtraMsg is an additional error message that callers can provide in
	// order to provide context specific error details.
	ExtraMsg string

	lnwire.FailureMessage
}

// Error implements the built-in error interface. We use this method to allow
// the switch or any callers to insert additional context to the error message
// returned.
func (f *ForwardingError) Error() string {
	if f.ExtraMsg == "" {
		return fmt.Sprintf(
			"%v@%v", f.FailureMessage, f.FailureSourceIdx,
		)
	}

	return fmt.Sprintf(
		"%v@%v: %v", f.FailureMessage, f.FailureSourceIdx, f.ExtraMsg,
	)
}

// ErrorDecrypter is an interface that is used to decrypt the onion encrypted
// failure reason an extra out a well formed error.
type ErrorDecrypter interface {
	// DecryptError peels off each layer of onion encryption from the first
	// hop, to the source of the error. A fully populated
	// lnwire.FailureMessage is returned along with the source of the
	// error.
	DecryptError(lnwire.OpaqueReason) (*ForwardingError, error)
}

// UnknownEncrypterType is an error message used to signal that an unexpected
// EncrypterType was encountered during decoding.
type UnknownEncrypterType hop.EncrypterType

// Error returns a formatted error indicating the invalid EncrypterType.
func (e UnknownEncrypterType) Error() string {
	return fmt.Sprintf("unknown error encrypter type: %d", e)
}

// OnionErrorDecrypter is the interface that provides onion level error
// decryption.
type OnionErrorDecrypter interface {
	// DecryptError attempts to decrypt the passed encrypted error response.
	// The onion failure is encrypted in backward manner, starting from the
	// node where error have occurred. As a result, in order to decrypt the
	// error we need get all shared secret and apply decryption in the
	// reverse order.
	DecryptError(encryptedData []byte) (*sphinx.DecryptedError, error)
}

// SphinxErrorDecrypter wraps the sphinx data SphinxErrorDecrypter and maps the
// returned errors to concrete lnwire.FailureMessage instances.
type SphinxErrorDecrypter struct {
	OnionErrorDecrypter
}

// DecryptError peels off each layer of onion encryption from the first hop, to
// the source of the error. A fully populated lnwire.FailureMessage is returned
// along with the source of the error.
//
// NOTE: Part of the ErrorDecrypter interface.
func (s *SphinxErrorDecrypter) DecryptError(reason lnwire.OpaqueReason) (
	*ForwardingError, error) {

	failure, err := s.OnionErrorDecrypter.DecryptError(reason)
	if err != nil {
		return nil, err
	}

	// Decode the failure. If an error occurs, we leave the failure message
	// field nil.
	r := bytes.NewReader(failure.Message)
	failureMsg, err := lnwire.DecodeFailure(r, 0)
	if err != nil {
		return &ForwardingError{
			FailureSourceIdx: failure.SenderIdx,
		}, nil
	}

	return &ForwardingError{
		FailureSourceIdx: failure.SenderIdx,
		FailureMessage:   failureMsg,
	}, nil
}

// A compile time check to ensure ErrorDecrypter implements the Deobfuscator
// interface.
var _ ErrorDecrypter = (*SphinxErrorDecrypter)(nil)
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"net/http"
	"time"

	"github.com/TeaMeow/KitSvc/module/logger"
	"github.com/TeaMeow/KitSvc/module/sd"
	"github.com/TeaMeow/KitSvc/router"
	"github.com/TeaMeow/KitSvc/router/middleware"
	"github.com/TeaMeow/KitSvc/shared/eventutil"
	"github.com/TeaMeow/KitSvc/shared/mqutil"
	"github.com/TeaMeow/KitSvc/shared/wsutil"
	"github.com/codegangsta/cli"
	"github.com/gin-gonic/gin"
)

var serverFlags = []cli.Flag{
	// Common flags.
	cli.StringFlag{
		EnvVar: "KITSVC_NAME",
		Name:   "name",
		Usage:  "the name of the service, exposed for service discovery.",
		Value:  "Service",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_URL",
		Name:   "url",
		Usage:  "the url of the service.",
		Value:  "http://127.0.0.1:8080",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_ADDR",
		Name:   "addr",
		Usage:  "the address of the service (with the port).",
		Value:  "127.0.0.1:8080",
	},
	cli.IntFlag{
		EnvVar: "KITSVC_PORT",
		Name:   "port",
		Usage:  "the port of the service.",
		Value:  8080,
	},
	cli.StringFlag{
		EnvVar: "KITSVC_USAGE",
		Name:   "usage",
		Usage:  "the usage of the service, exposed for service discovery.",
		Value:  "Operations about the users.",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_JWT_SECRET",
		Name:   "jwt-secret",
		Usage:  "the secert used to encode the json web token.",
		Value:  "4Rtg8BPKwixXy2ktDPxoMMAhRzmo9mmuZjvKONGPZZQSaJWNLijxR42qRgq0iBb5",
	},
	cli.IntFlag{
		EnvVar: "KITSVC_MAX_PING_COUNT",
		Name:   "max-ping-count",
		Usage:  "",
		Value:  20,
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DEBUG",
		Name:   "debug",
		Usage:  "enable the debug mode.",
	},

	// Database flags.
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_NAME",
		Name:   "database-name",
		Usage:  "the name of the database.",
		Value:  "service",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_HOST",
		Name:   "database-host",
		Usage:  "the host of the database (with the port).",
		Value:  "127.0.0.1:3306",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_USER",
		Name:   "database-user",
		Usage:  "the user of the database.",
		Value:  "root",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_PASSWORD",
		Name:   "database-password",
		Usage:  "the password of the database.",
		Value:  "root",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_CHARSET",
		Name:   "database-charset",
		Usage:  "the charset of the database.",
		Value:  "utf8",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_LOC",
		Name:   "database-loc",
		Usage:  "the timezone of the database.",
		Value:  "Local",
	},
	cli.BoolFlag{
		EnvVar: "KITSVC_DATABASE_PARSE_TIME",
		Name:   "database-parse_time",
		Usage:  "parse the time.",
	},

	// NSQ flags.
	cli.StringFlag{
		EnvVar: "KITSVC_NSQ_PRODUCER",
		Name:   "nsq-producer",
		Usage:  "the address of the TCP NSQ producer (with the port).",
		Value:  "127.0.0.1:4150",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_NSQ_PRODUCER_HTTP",
		Name:   "nsq-producer-http",
		Usage:  "the address of the HTTP NSQ producer (with the port).",
		Value:  "127.0.0.1:4151",
	},
	cli.StringSliceFlag{
		EnvVar: "KITSVC_NSQ_LOOKUPDS",
		Name:   "nsq-lookupds",
		Usage:  "the address of the NSQ lookupds (with the port).",
		Value: &cli.StringSlice{
			"127.0.0.1:4161",
		},
	},

	// Event store flags.
	cli.StringFlag{
		EnvVar: "KITSVC_ES_SERVER_URL",
		Name:   "es-url",
		Usage:  "the url of the event store server.",
		Value:  "http://127.0.0.1:2113",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_ES_USERNAME",
		Name:   "es-username",
		Usage:  "the username of the event store.",
		Value:  "admin",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_ES_PASSWORD",
		Name:   "es-password",
		Usage:  "the password of the event store.",
		Value:  "changeit",
	},

	// Prometheus flags.
	cli.StringFlag{
		EnvVar: "KITSVC_PROMETHEUS_NAMESPACE",
		Name:   "prometheus-namespace",
		Usage:  "the prometheus namespace.",
		Value:  "service",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_PROMETHEUS_SUBSYSTEM",
		Name:   "prometheus-subsystem",
		Usage:  "the subsystem of the promethues.",
		Value:  "user",
	},

	// Consul flags.
	cli.StringFlag{
		EnvVar: "KITSVC_CONSUL_CHECK_INTERVAL",
		Name:   "consul-check_interval",
		Usage:  "the interval of consul health check.",
		Value:  "30s",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_CONSUL_CHECK_TIMEOUT",
		Name:   "consul-check_timeout",
		Usage:  "the timeout of consul health check.",
		Value:  "1s",
	},
	cli.StringSliceFlag{
		EnvVar: "KITSVC_CONSUL_TAGS",
		Name:   "consul-tags",
		Usage:  "the service tags for consul.",
		Value: &cli.StringSlice{
			"user",
			"micro",
		},
	},
}

// server runs the server.
func server(c *cli.Context, started chan bool) error {
	// `deployed` will be closed when the router is deployed.
	deployed := make(chan bool)
	// `replayed` will be closed after the events are all replayed.
	replayed := make(chan bool)

	// Debug mode.
	if !c.Bool("debug") {
		gin.SetMode(gin.ReleaseMode)
	}

	// Initialize the logger.
	logger.Init(c)
	// Create the Gin engine.
	g := gin.New()
	// Event handlers.
	event := eventutil.New(g)
	// Websocket handlers.
	ws := wsutil.New(g)
	// Message queue handlers.
	mq := mqutil.New(g)

	// Routes.
	router.Load(
		// Cores.
		g, event, ws, mq,
		// Middlwares.
		middleware.Config(c),
		middleware.Store(c),
		middleware.Logging(),
		middleware.Event(c, event, replayed, deployed),
		middleware.MQ(c, mq, deployed),
		middleware.Metrics(),
	)

	// Register to the service registry when the events were replayed.
	go func() {
		<-replayed

		sd.Register(c)
		// After the service is registered to the consul,
		// close the `started` channel to make it non-blocking.
		close(started)
	}()

	// Ping the server to make sure the router is working.
	go func() {
		if err := pingServer(c); err != nil {
			logger.Fatal("The router has no response, or it might took too long to start up.")
		}
		logger.Info("The router has been deployed successfully.")
		// Close the `deployed` channel to make it non-blocking.
		close(deployed)
	}()

	// Start to listening the incoming requests.
	return http.ListenAndServe(c.String("addr"), g)
}

// pingServer pings the http server to make sure the router is working.
func pingServer(c *cli.Context) error {
	for i := 0; i < c.Int("max-ping-count"); i++ {
		// Ping the server by sending a GET request to `/health`.
		resp, err := http.Get(c.String("url") + "/sd/health")
		if err == nil && resp.StatusCode == 200 {
			return nil
		}

		// Sleep for a second to continue the next ping.
		logger.Info("Waiting for the router, retry in 1 second.")
		time.Sleep(time.Second)
	}
	return errors.New("Cannot connect to the router.")
}
<commit_msg>Added database driver env<commit_after>package main

import (
	"errors"
	"net/http"
	"time"

	"github.com/TeaMeow/KitSvc/module/logger"
	"github.com/TeaMeow/KitSvc/module/sd"
	"github.com/TeaMeow/KitSvc/router"
	"github.com/TeaMeow/KitSvc/router/middleware"
	"github.com/TeaMeow/KitSvc/shared/eventutil"
	"github.com/TeaMeow/KitSvc/shared/mqutil"
	"github.com/TeaMeow/KitSvc/shared/wsutil"
	"github.com/codegangsta/cli"
	"github.com/gin-gonic/gin"
)

var serverFlags = []cli.Flag{
	// Common flags.
	cli.StringFlag{
		EnvVar: "KITSVC_NAME",
		Name:   "name",
		Usage:  "the name of the service, exposed for service discovery.",
		Value:  "Service",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_URL",
		Name:   "url",
		Usage:  "the url of the service.",
		Value:  "http://127.0.0.1:8080",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_ADDR",
		Name:   "addr",
		Usage:  "the address of the service (with the port).",
		Value:  "127.0.0.1:8080",
	},
	cli.IntFlag{
		EnvVar: "KITSVC_PORT",
		Name:   "port",
		Usage:  "the port of the service.",
		Value:  8080,
	},
	cli.StringFlag{
		EnvVar: "KITSVC_USAGE",
		Name:   "usage",
		Usage:  "the usage of the service, exposed for service discovery.",
		Value:  "Operations about the users.",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_JWT_SECRET",
		Name:   "jwt-secret",
		Usage:  "the secert used to encode the json web token.",
		Value:  "4Rtg8BPKwixXy2ktDPxoMMAhRzmo9mmuZjvKONGPZZQSaJWNLijxR42qRgq0iBb5",
	},
	cli.IntFlag{
		EnvVar: "KITSVC_MAX_PING_COUNT",
		Name:   "max-ping-count",
		Usage:  "",
		Value:  20,
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DEBUG",
		Name:   "debug",
		Usage:  "enable the debug mode.",
	},

	// Database flags.
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_DRIVER",
		Name:   "database-driver",
		Usage:  "the driver of the database.",
		Value:  "mysql",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_NAME",
		Name:   "database-name",
		Usage:  "the name of the database.",
		Value:  "service",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_HOST",
		Name:   "database-host",
		Usage:  "the host of the database (with the port).",
		Value:  "127.0.0.1:3306",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_USER",
		Name:   "database-user",
		Usage:  "the user of the database.",
		Value:  "root",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_PASSWORD",
		Name:   "database-password",
		Usage:  "the password of the database.",
		Value:  "root",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_CHARSET",
		Name:   "database-charset",
		Usage:  "the charset of the database.",
		Value:  "utf8",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_DATABASE_LOC",
		Name:   "database-loc",
		Usage:  "the timezone of the database.",
		Value:  "Local",
	},
	cli.BoolFlag{
		EnvVar: "KITSVC_DATABASE_PARSE_TIME",
		Name:   "database-parse_time",
		Usage:  "parse the time.",
	},

	// NSQ flags.
	cli.StringFlag{
		EnvVar: "KITSVC_NSQ_PRODUCER",
		Name:   "nsq-producer",
		Usage:  "the address of the TCP NSQ producer (with the port).",
		Value:  "127.0.0.1:4150",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_NSQ_PRODUCER_HTTP",
		Name:   "nsq-producer-http",
		Usage:  "the address of the HTTP NSQ producer (with the port).",
		Value:  "127.0.0.1:4151",
	},
	cli.StringSliceFlag{
		EnvVar: "KITSVC_NSQ_LOOKUPDS",
		Name:   "nsq-lookupds",
		Usage:  "the address of the NSQ lookupds (with the port).",
		Value: &cli.StringSlice{
			"127.0.0.1:4161",
		},
	},

	// Event store flags.
	cli.StringFlag{
		EnvVar: "KITSVC_ES_SERVER_URL",
		Name:   "es-url",
		Usage:  "the url of the event store server.",
		Value:  "http://127.0.0.1:2113",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_ES_USERNAME",
		Name:   "es-username",
		Usage:  "the username of the event store.",
		Value:  "admin",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_ES_PASSWORD",
		Name:   "es-password",
		Usage:  "the password of the event store.",
		Value:  "changeit",
	},

	// Prometheus flags.
	cli.StringFlag{
		EnvVar: "KITSVC_PROMETHEUS_NAMESPACE",
		Name:   "prometheus-namespace",
		Usage:  "the prometheus namespace.",
		Value:  "service",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_PROMETHEUS_SUBSYSTEM",
		Name:   "prometheus-subsystem",
		Usage:  "the subsystem of the promethues.",
		Value:  "user",
	},

	// Consul flags.
	cli.StringFlag{
		EnvVar: "KITSVC_CONSUL_CHECK_INTERVAL",
		Name:   "consul-check_interval",
		Usage:  "the interval of consul health check.",
		Value:  "30s",
	},
	cli.StringFlag{
		EnvVar: "KITSVC_CONSUL_CHECK_TIMEOUT",
		Name:   "consul-check_timeout",
		Usage:  "the timeout of consul health check.",
		Value:  "1s",
	},
	cli.StringSliceFlag{
		EnvVar: "KITSVC_CONSUL_TAGS",
		Name:   "consul-tags",
		Usage:  "the service tags for consul.",
		Value: &cli.StringSlice{
			"user",
			"micro",
		},
	},
}

// server runs the server.
func server(c *cli.Context, started chan bool) error {
	// `deployed` will be closed when the router is deployed.
	deployed := make(chan bool)
	// `replayed` will be closed after the events are all replayed.
	replayed := make(chan bool)

	// Debug mode.
	if !c.Bool("debug") {
		gin.SetMode(gin.ReleaseMode)
	}

	// Initialize the logger.
	logger.Init(c)
	// Create the Gin engine.
	g := gin.New()
	// Event handlers.
	event := eventutil.New(g)
	// Websocket handlers.
	ws := wsutil.New(g)
	// Message queue handlers.
	mq := mqutil.New(g)

	// Routes.
	router.Load(
		// Cores.
		g, event, ws, mq,
		// Middlwares.
		middleware.Config(c),
		middleware.Store(c),
		middleware.Logging(),
		middleware.Event(c, event, replayed, deployed),
		middleware.MQ(c, mq, deployed),
		middleware.Metrics(),
	)

	// Register to the service registry when the events were replayed.
	go func() {
		<-replayed

		sd.Register(c)
		// After the service is registered to the consul,
		// close the `started` channel to make it non-blocking.
		close(started)
	}()

	// Ping the server to make sure the router is working.
	go func() {
		if err := pingServer(c); err != nil {
			logger.Fatal("The router has no response, or it might took too long to start up.")
		}
		logger.Info("The router has been deployed successfully.")
		// Close the `deployed` channel to make it non-blocking.
		close(deployed)
	}()

	// Start to listening the incoming requests.
	return http.ListenAndServe(c.String("addr"), g)
}

// pingServer pings the http server to make sure the router is working.
func pingServer(c *cli.Context) error {
	for i := 0; i < c.Int("max-ping-count"); i++ {
		// Ping the server by sending a GET request to `/health`.
		resp, err := http.Get(c.String("url") + "/sd/health")
		if err == nil && resp.StatusCode == 200 {
			return nil
		}

		// Sleep for a second to continue the next ping.
		logger.Info("Waiting for the router, retry in 1 second.")
		time.Sleep(time.Second)
	}
	return errors.New("Cannot connect to the router.")
}
<|endoftext|>


<commit_before>package server

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

	"github.com/gorilla/mux"
	"github.com/nytlabs/st-core/core"
	"github.com/thejerf/suture"
)

type SourceLedger struct {
	Label      string              `json:"label"`
	Type       string              `json:"type"`
	Id         int                 `json:"id"`
	Source     core.Source         `json:"-"`
	Parent     *Group              `json:"-"`
	Token      suture.ServiceToken `json:"-"`
	Position   Position            `json:"position"`
	Parameters map[string]string   `json:"params,omitempty"`
}

type ProtoSource struct {
	Label    string   `json:"label"`
	Type     string   `json:"type"`
	Position Position `json:"position"`
	Parent   int      `json:"group"`
}

func (sl *SourceLedger) GetID() int {
	return sl.Id
}

func (sl *SourceLedger) GetParent() *Group {
	return sl.Parent
}

func (sl *SourceLedger) SetParent(group *Group) {
	sl.Parent = group
}

func (s *Server) ListSources() []SourceLedger {
	sources := []SourceLedger{}
	for _, source := range s.sources {
		sources = append(sources, *source)
	}
	return sources
}

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

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(s.ListSources()); err != nil {
		panic(err)
	}
}

func (s *Server) SourceHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}
	s.Lock()
	defer s.Unlock()
	source, ok := s.sources[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find source"})
		return
	}
	w.WriteHeader(http.StatusOK)
	writeJSON(w, source)
}

func (s *Server) CreateSource(p ProtoSource) (*SourceLedger, error) {
	f, ok := s.sourceLibrary[p.Type]
	if !ok {
		return nil, errors.New("source type " + p.Type + " does not exist")
	}

	source := f.New()

	sl := &SourceLedger{
		Label:    p.Label,
		Position: p.Position,
		Source:   source,
		Type:     p.Type,
		Id:       s.GetNextID(),
	}

	if i, ok := source.(core.Interface); ok {
		// Describe() is not thread-safe it must be put ahead of supervior...
		sl.Parameters = i.Describe()
		sl.Token = s.supervisor.Add(i)
	}

	s.sources[sl.Id] = sl
	s.websocketBroadcast(Update{Action: CREATE, Type: SOURCE, Data: sl})

	err := s.AddChildToGroup(p.Parent, sl)
	if err != nil {
		return nil, err

	}

	return sl, nil
}

func (s *Server) DeleteSource(id int) error {
	source, ok := s.sources[id]
	if !ok {
		return errors.New("could not find source")
	}

	for _, l := range s.links {
		if l.Source == id {
			err := s.DeleteLink(l.Id)
			if err != nil {
				return err
			}
		}
	}

	if _, ok := source.Source.(core.Interface); ok {
		s.supervisor.Remove(source.Token)
	}

	s.DetachChild(source)

	s.websocketBroadcast(Update{Action: DELETE, Type: SOURCE, Data: s.sources[id]})
	delete(s.sources, source.Id)
	return nil
}

func (s *Server) SourceCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var m ProtoSource
	err = json.Unmarshal(body, &m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, err := s.CreateSource(m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, b)
}

func (s *Server) ModifySource(id int, m map[string]string) error {
	source, ok := s.sources[id]
	if !ok {
		return errors.New("no source found")
	}

	i, ok := source.Source.(core.Interface)
	if !ok {
		return errors.New("cannot modify store")
	}

	s.supervisor.Remove(source.Token)
	for k, _ := range source.Parameters {
		if v, ok := m[k]; ok {
			i.SetSourceParameter(k, v)
			source.Parameters[k] = v
			update := struct {
				Id    int    `json:"id"`
				Key   string `json:"param"`
				Value string `json:"value"`
			}{
				id, k, v,
			}
			s.websocketBroadcast(Update{Action: UPDATE, Type: SOURCE, Data: update})
		}
	}
	source.Token = s.supervisor.Add(i)
	return nil
}

func (s *Server) SourceModifyHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var m map[string]string
	err = json.Unmarshal(body, &m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.ModifySource(id, m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) SourceDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteSource(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) SourceGetValueHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	val, err := s.GetSourceValue(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	w.Write(val)
}

func (s *Server) SourceModifyPositionHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var p Position
	err = json.Unmarshal(body, &p)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read JSON"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, ok := s.sources[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find block"})
		return
	}

	b.Position = p

	update := struct {
		Position
		Id int
	}{
		p,
		id,
	}

	s.websocketBroadcast(Update{Action: UPDATE, Type: SOURCE, Data: update})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) SourceModifyNameHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	s.Lock()
	defer s.Unlock()

	_, ok := s.sources[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"block not found"})
		return
	}

	var label string
	err = json.Unmarshal(body, &label)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.sources[id].Label = label

	update := struct {
		Id    int    `json:"id"`
		Label string `json:"label"`
	}{
		id, label,
	}

	s.websocketBroadcast(Update{Action: UPDATE, Type: SOURCE, Data: update})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) SourceSetValueHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.SetSourceValue(id, body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) GetSourceValue(id int) ([]byte, error) {
	source, ok := s.sources[id]
	if !ok {
		return nil, errors.New("source does not exist")
	}

	store, ok := source.Source.(core.Store)
	if !ok {
		return nil, errors.New("can only get values from stores")
	}

	store.Lock()
	defer store.Unlock()
	out, err := json.Marshal(store.Get())
	if err != nil {
		return nil, err
	}

	return out, nil
}

func (s *Server) SetSourceValue(id int, body []byte) error {
	source, ok := s.sources[id]
	if !ok {
		return errors.New("source does not exist")
	}

	store, ok := source.Source.(core.Store)
	if !ok {
		return errors.New("can only get values from stores")
	}

	var m interface{}
	err := json.Unmarshal(body, &m)
	if err != nil {
		return err
	}

	store.Lock()
	defer store.Unlock()
	err = store.Set(m)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>fixing source api<commit_after>package server

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

	"github.com/gorilla/mux"
	"github.com/nytlabs/st-core/core"
	"github.com/thejerf/suture"
)

type SourceLedger struct {
	Label      string              `json:"label"`
	Type       string              `json:"type"`
	Id         int                 `json:"id"`
	Source     core.Source         `json:"-"`
	Parent     *Group              `json:"-"`
	Token      suture.ServiceToken `json:"-"`
	Position   Position            `json:"position"`
	Parameters map[string]string   `json:"params,omitempty"`
}

type ProtoSource struct {
	Label    string   `json:"label"`
	Type     string   `json:"type"`
	Position Position `json:"position"`
	Parent   int      `json:"group"`
}

func (sl *SourceLedger) GetID() int {
	return sl.Id
}

func (sl *SourceLedger) GetParent() *Group {
	return sl.Parent
}

func (sl *SourceLedger) SetParent(group *Group) {
	sl.Parent = group
}

func (s *Server) ListSources() []SourceLedger {
	sources := []SourceLedger{}
	for _, source := range s.sources {
		sources = append(sources, *source)
	}
	return sources
}

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

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(s.ListSources()); err != nil {
		panic(err)
	}
}

func (s *Server) SourceHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}
	s.Lock()
	defer s.Unlock()
	source, ok := s.sources[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find source"})
		return
	}
	w.WriteHeader(http.StatusOK)
	writeJSON(w, source)
}

func (s *Server) CreateSource(p ProtoSource) (*SourceLedger, error) {
	f, ok := s.sourceLibrary[p.Type]
	if !ok {
		return nil, errors.New("source type " + p.Type + " does not exist")
	}

	source := f.New()

	sl := &SourceLedger{
		Label:    p.Label,
		Position: p.Position,
		Source:   source,
		Type:     p.Type,
		Id:       s.GetNextID(),
	}

	if i, ok := source.(core.Interface); ok {
		// Describe() is not thread-safe it must be put ahead of supervior...
		sl.Parameters = i.Describe()
		sl.Token = s.supervisor.Add(i)
	}

	s.sources[sl.Id] = sl
	s.websocketBroadcast(Update{Action: CREATE, Type: SOURCE, Data: wsSource{*sl}})

	err := s.AddChildToGroup(p.Parent, sl)
	if err != nil {
		return nil, err

	}

	return sl, nil
}

func (s *Server) DeleteSource(id int) error {
	source, ok := s.sources[id]
	if !ok {
		return errors.New("could not find source")
	}

	for _, l := range s.links {
		if l.Source.Id == id {
			err := s.DeleteLink(l.Id)
			if err != nil {
				return err
			}
		}
	}

	if _, ok := source.Source.(core.Interface); ok {
		s.supervisor.Remove(source.Token)
	}

	s.DetachChild(source)

	s.websocketBroadcast(Update{Action: DELETE, Type: SOURCE, Data: wsSource{wsId{id}}})
	delete(s.sources, source.Id)
	return nil
}

func (s *Server) SourceCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var m ProtoSource
	err = json.Unmarshal(body, &m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, err := s.CreateSource(m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, b)
}

func (s *Server) ModifySource(id int, m map[string]string) error {
	source, ok := s.sources[id]
	if !ok {
		return errors.New("no source found")
	}

	i, ok := source.Source.(core.Interface)
	if !ok {
		return errors.New("cannot modify store")
	}

	s.supervisor.Remove(source.Token)
	for k, _ := range source.Parameters {
		if v, ok := m[k]; ok {
			i.SetSourceParameter(k, v)
			source.Parameters[k] = v
			s.websocketBroadcast(Update{Action: UPDATE, Type: PARAM, Data: wsSourceModify{wsId{id}, k, v}})
		}
	}
	source.Token = s.supervisor.Add(i)
	return nil
}

func (s *Server) SourceModifyHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var m map[string]string
	err = json.Unmarshal(body, &m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.ModifySource(id, m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) SourceDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteSource(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) SourceGetValueHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	val, err := s.GetSourceValue(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	w.Write(val)
}

func (s *Server) SourceModifyPositionHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var p Position
	err = json.Unmarshal(body, &p)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read JSON"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, ok := s.sources[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find block"})
		return
	}

	b.Position = p

	s.websocketBroadcast(Update{Action: UPDATE, Type: SOURCE, Data: wsSource{wsPosition{wsId{id}, p}}})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) SourceModifyNameHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	s.Lock()
	defer s.Unlock()

	_, ok := s.sources[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"block not found"})
		return
	}

	var label string
	err = json.Unmarshal(body, &label)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.sources[id].Label = label

	s.websocketBroadcast(Update{Action: UPDATE, Type: SOURCE, Data: wsSource{wsLabel{wsId{id}, label}}})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) SourceSetValueHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.SetSourceValue(id, body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) GetSourceValue(id int) ([]byte, error) {
	source, ok := s.sources[id]
	if !ok {
		return nil, errors.New("source does not exist")
	}

	store, ok := source.Source.(core.Store)
	if !ok {
		return nil, errors.New("can only get values from stores")
	}

	store.Lock()
	defer store.Unlock()
	out, err := json.Marshal(store.Get())
	if err != nil {
		return nil, err
	}

	return out, nil
}

func (s *Server) SetSourceValue(id int, body []byte) error {
	source, ok := s.sources[id]
	if !ok {
		return errors.New("source does not exist")
	}

	store, ok := source.Source.(core.Store)
	if !ok {
		return errors.New("can only get values from stores")
	}

	var m interface{}
	err := json.Unmarshal(body, &m)
	if err != nil {
		return err
	}

	store.Lock()
	defer store.Unlock()
	err = store.Set(m)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


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

package main

import (
	"bytes"
	"github.com/globocom/tsuru/cmd"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"net/http"
	"os"
)

func (s *S) TestServiceCreateInfo(c *C) {
	desc := "Creates a service based on a passed manifest. The manifest format should be a yaml and follow the standard described in the documentation (should link to it here)"
	cmd := ServiceCreate{}
	i := cmd.Info()
	c.Assert(i.Name, Equals, "create")
	c.Assert(i.Usage, Equals, "create path/to/manifesto")
	c.Assert(i.Desc, Equals, desc)
	c.Assert(i.MinArgs, Equals, 1)
}

func (s *S) TestServiceCreateRun(c *C) {
	var stdout, stderr bytes.Buffer
	args := []string{"testdata/manifest.yml"}
	context := cmd.Context{
		Args:   args,
		Stdout: &stdout,
		Stderr: &stderr,
	}
	client := cmd.NewClient(&http.Client{Transport: &transport{msg: "success", status: http.StatusOK}}, nil, manager)
	err := (&ServiceCreate{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(stdout.String(), Equals, "success")
}

func (s *S) TestServiceRemoveRun(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	context := cmd.Context{
		Args:   []string{"my-service"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	trans := &conditionalTransport{
		transport{
			msg:    "",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "DELETE" && req.URL.Path == "/services/my-service"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	err := (&ServiceRemove{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(stdout.String(), Equals, "Service successfully removed.\n")
}

func (s *S) TestServiceRemoveRunWithRequestFailure(c *C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Args:   []string{"my-service"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	trans := transport{
		msg:    "This service cannot be removed because it has instances.\nPlease remove these instances before removing the service.",
		status: http.StatusForbidden,
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	err := (&ServiceRemove{}).Run(&context, client)
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, trans.msg)
}

func (s *S) TestServiceRemoveIsACommand(c *C) {
	var _ cmd.Command = &ServiceRemove{}
}

func (s *S) TestServiceRemoveInfo(c *C) {
	expected := &cmd.Info{
		Name:    "remove",
		Usage:   "remove <servicename>",
		Desc:    "removes a service from catalog",
		MinArgs: 1,
	}
	c.Assert((&ServiceRemove{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceRemoveIsAnInfor(c *C) {
	var _ cmd.Infoer = &ServiceRemove{}
}

func (s *S) TestServiceListInfo(c *C) {
	cmd := ServiceList{}
	i := cmd.Info()
	c.Assert(i.Name, Equals, "list")
	c.Assert(i.Usage, Equals, "list")
	c.Assert(i.Desc, Equals, "list services that belongs to user's team and it's service instances.")
}

func (s *S) TestServiceListRun(c *C) {
	var stdout, stderr bytes.Buffer
	response := `[{"service": "mysql", "instances": ["my_db"]}]`
	expected := `+----------+-----------+
| Services | Instances |
+----------+-----------+
| mysql    | my_db     |
+----------+-----------+
`
	trans := transport{msg: response, status: http.StatusOK}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceList{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(stdout.String(), Equals, expected)
}

func (s *S) TestServiceListRunWithNoServicesReturned(c *C) {
	var stdout, stderr bytes.Buffer
	response := `[]`
	expected := ""
	trans := transport{msg: response, status: http.StatusOK}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceList{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(stdout.String(), Equals, expected)
}

func (s *S) TestServiceUpdate(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	trans := conditionalTransport{
		transport{
			msg:    "",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "PUT" && req.URL.Path == "/services"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{"testdata/manifest.yml"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceUpdate{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(stdout.String(), Equals, "Service successfully updated.\n")
}

func (s *S) TestServiceUpdateIsACommand(c *C) {
	var _ cmd.Command = &ServiceUpdate{}
}

func (s *S) TestServiceUpdateInfo(c *C) {
	expected := &cmd.Info{
		Name:    "update",
		Usage:   "update <path/to/manifesto>",
		Desc:    "Update service data, extracting it from the given manifesto file.",
		MinArgs: 1,
	}
	c.Assert((&ServiceUpdate{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceUpdateIsAnInfoer(c *C) {
	var _ cmd.Infoer = &ServiceUpdate{}
}

func (s *S) TestServiceDocAdd(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	trans := conditionalTransport{
		transport{
			msg:    "",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "PUT" && req.URL.Path == "/services/serv/doc"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{"serv", "testdata/doc.md"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceDocAdd{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(stdout.String(), Equals, "Documentation for 'serv' successfully updated.\n")
}

func (s *S) TestServiceDocAddInfo(c *C) {
	expected := &cmd.Info{
		Name:    "doc-add",
		Usage:   "service doc-add <service> <path/to/docfile>",
		Desc:    "Update service documentation, extracting it from the given file.",
		MinArgs: 2,
	}
	c.Assert((&ServiceDocAdd{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceDocGet(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	trans := conditionalTransport{
		transport{
			msg:    "some doc",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "GET" && req.URL.Path == "/services/serv/doc"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{"serv"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceDocGet{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(context.Stdout.(*bytes.Buffer).String(), Equals, "some doc")
}

func (s *S) TestServiceDocGetInfo(c *C) {
	expected := &cmd.Info{
		Name:    "doc-get",
		Usage:   "service doc-get <service>",
		Desc:    "Shows service documentation.",
		MinArgs: 1,
	}
	c.Assert((&ServiceDocGet{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceTemplateInfo(c *C) {
	got := (&ServiceTemplate{}).Info()
	usg := `template
e.g.: $ crane template`
	expected := &cmd.Info{
		Name:  "template",
		Usage: usg,
		Desc:  "Generates a manifest template file and places it in current path",
	}
	c.Assert(got, DeepEquals, expected)
}

func (s *S) TestServiceTemplateRun(c *C) {
	var stdout, stderr bytes.Buffer
	trans := transport{msg: "", status: http.StatusOK}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	ctx := cmd.Context{
		Args:   []string{},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceTemplate{}).Run(&ctx, client)
	defer os.Remove("./manifest.yaml")
	c.Assert(err, IsNil)
	expected := "Generated file \"manifest.yaml\" in current path\n"
	c.Assert(stdout.String(), Equals, expected)
	f, err := os.Open("./manifest.yaml")
	c.Assert(err, IsNil)
	fc, err := ioutil.ReadAll(f)
	manifest := `id: servicename
endpoint:
  production: production-endpoint.com
  test: test-endpoint.com:8080`
	c.Assert(string(fc), Equals, manifest)
}
<commit_msg>cmd/crane: fix tests that used Infoer<commit_after>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"github.com/globocom/tsuru/cmd"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"net/http"
	"os"
)

func (s *S) TestServiceCreateInfo(c *C) {
	desc := "Creates a service based on a passed manifest. The manifest format should be a yaml and follow the standard described in the documentation (should link to it here)"
	cmd := ServiceCreate{}
	i := cmd.Info()
	c.Assert(i.Name, Equals, "create")
	c.Assert(i.Usage, Equals, "create path/to/manifesto")
	c.Assert(i.Desc, Equals, desc)
	c.Assert(i.MinArgs, Equals, 1)
}

func (s *S) TestServiceCreateRun(c *C) {
	var stdout, stderr bytes.Buffer
	args := []string{"testdata/manifest.yml"}
	context := cmd.Context{
		Args:   args,
		Stdout: &stdout,
		Stderr: &stderr,
	}
	client := cmd.NewClient(&http.Client{Transport: &transport{msg: "success", status: http.StatusOK}}, nil, manager)
	err := (&ServiceCreate{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(stdout.String(), Equals, "success")
}

func (s *S) TestServiceRemoveRun(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	context := cmd.Context{
		Args:   []string{"my-service"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	trans := &conditionalTransport{
		transport{
			msg:    "",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "DELETE" && req.URL.Path == "/services/my-service"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	err := (&ServiceRemove{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(stdout.String(), Equals, "Service successfully removed.\n")
}

func (s *S) TestServiceRemoveRunWithRequestFailure(c *C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Args:   []string{"my-service"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	trans := transport{
		msg:    "This service cannot be removed because it has instances.\nPlease remove these instances before removing the service.",
		status: http.StatusForbidden,
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	err := (&ServiceRemove{}).Run(&context, client)
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, trans.msg)
}

func (s *S) TestServiceRemoveIsACommand(c *C) {
	var _ cmd.Command = &ServiceRemove{}
}

func (s *S) TestServiceRemoveInfo(c *C) {
	expected := &cmd.Info{
		Name:    "remove",
		Usage:   "remove <servicename>",
		Desc:    "removes a service from catalog",
		MinArgs: 1,
	}
	c.Assert((&ServiceRemove{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceListInfo(c *C) {
	cmd := ServiceList{}
	i := cmd.Info()
	c.Assert(i.Name, Equals, "list")
	c.Assert(i.Usage, Equals, "list")
	c.Assert(i.Desc, Equals, "list services that belongs to user's team and it's service instances.")
}

func (s *S) TestServiceListRun(c *C) {
	var stdout, stderr bytes.Buffer
	response := `[{"service": "mysql", "instances": ["my_db"]}]`
	expected := `+----------+-----------+
| Services | Instances |
+----------+-----------+
| mysql    | my_db     |
+----------+-----------+
`
	trans := transport{msg: response, status: http.StatusOK}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceList{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(stdout.String(), Equals, expected)
}

func (s *S) TestServiceListRunWithNoServicesReturned(c *C) {
	var stdout, stderr bytes.Buffer
	response := `[]`
	expected := ""
	trans := transport{msg: response, status: http.StatusOK}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceList{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(stdout.String(), Equals, expected)
}

func (s *S) TestServiceUpdate(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	trans := conditionalTransport{
		transport{
			msg:    "",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "PUT" && req.URL.Path == "/services"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{"testdata/manifest.yml"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceUpdate{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(stdout.String(), Equals, "Service successfully updated.\n")
}

func (s *S) TestServiceUpdateIsACommand(c *C) {
	var _ cmd.Command = &ServiceUpdate{}
}

func (s *S) TestServiceUpdateInfo(c *C) {
	expected := &cmd.Info{
		Name:    "update",
		Usage:   "update <path/to/manifesto>",
		Desc:    "Update service data, extracting it from the given manifesto file.",
		MinArgs: 1,
	}
	c.Assert((&ServiceUpdate{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceDocAdd(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	trans := conditionalTransport{
		transport{
			msg:    "",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "PUT" && req.URL.Path == "/services/serv/doc"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{"serv", "testdata/doc.md"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceDocAdd{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(stdout.String(), Equals, "Documentation for 'serv' successfully updated.\n")
}

func (s *S) TestServiceDocAddInfo(c *C) {
	expected := &cmd.Info{
		Name:    "doc-add",
		Usage:   "service doc-add <service> <path/to/docfile>",
		Desc:    "Update service documentation, extracting it from the given file.",
		MinArgs: 2,
	}
	c.Assert((&ServiceDocAdd{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceDocGet(c *C) {
	var (
		called         bool
		stdout, stderr bytes.Buffer
	)
	trans := conditionalTransport{
		transport{
			msg:    "some doc",
			status: http.StatusNoContent,
		},
		func(req *http.Request) bool {
			called = true
			return req.Method == "GET" && req.URL.Path == "/services/serv/doc"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	context := cmd.Context{
		Args:   []string{"serv"},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceDocGet{}).Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(called, Equals, true)
	c.Assert(context.Stdout.(*bytes.Buffer).String(), Equals, "some doc")
}

func (s *S) TestServiceDocGetInfo(c *C) {
	expected := &cmd.Info{
		Name:    "doc-get",
		Usage:   "service doc-get <service>",
		Desc:    "Shows service documentation.",
		MinArgs: 1,
	}
	c.Assert((&ServiceDocGet{}).Info(), DeepEquals, expected)
}

func (s *S) TestServiceTemplateInfo(c *C) {
	got := (&ServiceTemplate{}).Info()
	usg := `template
e.g.: $ crane template`
	expected := &cmd.Info{
		Name:  "template",
		Usage: usg,
		Desc:  "Generates a manifest template file and places it in current path",
	}
	c.Assert(got, DeepEquals, expected)
}

func (s *S) TestServiceTemplateRun(c *C) {
	var stdout, stderr bytes.Buffer
	trans := transport{msg: "", status: http.StatusOK}
	client := cmd.NewClient(&http.Client{Transport: &trans}, nil, manager)
	ctx := cmd.Context{
		Args:   []string{},
		Stdout: &stdout,
		Stderr: &stderr,
	}
	err := (&ServiceTemplate{}).Run(&ctx, client)
	defer os.Remove("./manifest.yaml")
	c.Assert(err, IsNil)
	expected := "Generated file \"manifest.yaml\" in current path\n"
	c.Assert(stdout.String(), Equals, expected)
	f, err := os.Open("./manifest.yaml")
	c.Assert(err, IsNil)
	fc, err := ioutil.ReadAll(f)
	manifest := `id: servicename
endpoint:
  production: production-endpoint.com
  test: test-endpoint.com:8080`
	c.Assert(string(fc), Equals, manifest)
}
<|endoftext|>


<commit_before>package data

// Evil things happen here. Rippled needs a V2 API...

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
)

type ledgerJSON Ledger

// adds all the legacy fields
type ledgerExtraJSON struct {
	ledgerJSON
	HumanCloseTime *RippleHumanTime `json:"close_time_human"`
	Hash           Hash256          `json:"hash"`
	LedgerHash     Hash256          `json:"ledger_hash"`
	TotalCoins     uint64           `json:"totalCoins,string"`
	SequenceNumber uint32           `json:"seqNum,string"`
}

func (l Ledger) MarshalJSON() ([]byte, error) {
	return json.Marshal(ledgerExtraJSON{
		ledgerJSON:     ledgerJSON(l),
		HumanCloseTime: l.CloseTime.Human(),
		Hash:           l.Hash(),
		LedgerHash:     l.Hash(),
		TotalCoins:     l.TotalXRP,
		SequenceNumber: l.LedgerSequence,
	})
}

func (l *Ledger) UnmarshalJSON(b []byte) error {
	var ledger ledgerExtraJSON
	if err := json.Unmarshal(b, &ledger); err != nil {
		return err
	}
	*l = Ledger(ledger.ledgerJSON)
	l.SetHash(ledger.Hash[:])
	return nil
}

// Wrapper types to enable second level of marshalling
// when found in ledger API call
type txmLedger struct {
	MetaData MetaData `json:"metaData"`
}

// Wrapper types to enable second level of marshalling
// when found in tx API call
type txmNormal TransactionWithMetaData

var (
	txmTransactionTypeRegex = regexp.MustCompile(`"TransactionType":\s*"(\w+)"`)
	txmHashRegex            = regexp.MustCompile(`"hash":\s*"(\w+)"`)
	txmMetaTypeRegex        = regexp.MustCompile(`"(meta|metaData)"`)
)

func (txm *TransactionWithMetaData) UnmarshalJSON(b []byte) error {
	txTypeMatch := txmTransactionTypeRegex.FindStringSubmatch(string(b))
	hashMatch := txmHashRegex.FindStringSubmatch(string(b))
	metaTypeMatch := txmMetaTypeRegex.FindStringSubmatch(string(b))
	var txType, hash, metaType string
	if txTypeMatch == nil {
		return fmt.Errorf("Not a valid transaction with metadata: Missing TransactionType")
	}
	txType = txTypeMatch[1]
	if hashMatch == nil {
		return fmt.Errorf("Not a valid transaction with metadata: Missing Hash")
	}
	hash = hashMatch[1]
	if metaTypeMatch != nil {
		metaType = metaTypeMatch[1]
	}
	txm.Transaction = GetTxFactoryByType(txType)()
	h, err := hex.DecodeString(hash)
	if err != nil {
		return fmt.Errorf("Bad hash: %s", hash)
	}
	txm.SetHash(h)
	if err := json.Unmarshal(b, txm.Transaction); err != nil {
		return err
	}
	switch metaType {
	case "meta":
		return json.Unmarshal(b, (*txmNormal)(txm))
	case "metaData":
		var meta txmLedger
		if err := json.Unmarshal(b, &meta); err != nil {
			return err
		}
		txm.MetaData = meta.MetaData
		return nil
	default:
		return nil
	}
}

func (txm TransactionWithMetaData) marshalJSON() ([]byte, []byte, error) {
	tx, err := json.Marshal(txm.Transaction)
	if err != nil {
		return nil, nil, err
	}
	meta, err := json.Marshal(txm.MetaData)
	if err != nil {
		return nil, nil, err
	}
	return tx, meta, nil
}

type extractTxm struct {
	Tx   json.RawMessage `json:"transaction"`
	Meta json.RawMessage `json:"meta"`
}

const extractTxmFormat = `%s,"meta":%s}`

func UnmarshalTransactionWithMetadata(b []byte, txm *TransactionWithMetaData) error {
	var extract extractTxm
	if err := json.Unmarshal(b, &extract); err != nil {
		return err
	}
	raw := fmt.Sprintf(extractTxmFormat, extract.Tx[:len(extract.Tx)-1], extract.Meta)
	return json.Unmarshal([]byte(raw), txm)
}

const txmFormat = `%s,"hash":"%s","inLedger":%d,"ledger_index":%d,"meta":%s}`

func (txm TransactionWithMetaData) MarshalJSON() ([]byte, error) {
	tx, meta, err := txm.marshalJSON()
	if err != nil {
		return nil, err
	}
	out := fmt.Sprintf(txmFormat, string(tx[:len(tx)-1]), txm.Hash().String(), txm.LedgerSequence, txm.LedgerSequence, string(meta))
	return []byte(out), nil
}

const txmSliceFormat = `%s,"hash":"%s","metaData":%s}`

func (s TransactionSlice) MarshalJSON() ([]byte, error) {
	raw := make([]json.RawMessage, len(s))
	var err error
	var tx, meta []byte
	for i, txm := range s {
		if tx, meta, err = txm.marshalJSON(); err != nil {
			return nil, err
		}
		extra := fmt.Sprintf(txmSliceFormat, string(tx[:len(tx)-1]), txm.Hash().String(), meta)
		raw[i] = json.RawMessage(extra)
	}
	return json.Marshal(raw)
}

var (
	leTypeRegex  = regexp.MustCompile(`"LedgerEntryType":\s*"(\w+)"`)
	leIndexRegex = regexp.MustCompile(`"index":\s*"(\w+)"`)
)

func (l *LedgerEntrySlice) UnmarshalJSON(b []byte) error {
	var s []json.RawMessage
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	for _, raw := range s {
		leTypeMatch := leTypeRegex.FindStringSubmatch(string(raw))
		indexMatch := leIndexRegex.FindStringSubmatch(string(raw))
		if leTypeMatch == nil {
			return fmt.Errorf("Bad LedgerEntryType")
		}
		if indexMatch == nil {
			return fmt.Errorf("Missing LedgerEntry index")
		}
		le := GetLedgerEntryFactoryByType(leTypeMatch[1])()
		index, err := hex.DecodeString(indexMatch[1])
		if err != nil {
			return fmt.Errorf("Bad index: %s", index)
		}
		le.SetHash(index)
		if err := json.Unmarshal(raw, &le); err != nil {
			return err
		}
		*l = append(*l, le)
	}
	return nil
}

const leSliceFormat = `%s,"index":"%s"}`

func (s LedgerEntrySlice) MarshalJSON() ([]byte, error) {
	raw := make([]json.RawMessage, len(s))
	var err error
	for i, le := range s {
		if raw[i], err = json.Marshal(le); err != nil {
			return nil, err
		}
		extra := fmt.Sprintf(leSliceFormat, string(raw[i][:len(raw[i])-1]), le.Hash().String())
		raw[i] = json.RawMessage(extra)
	}
	return json.Marshal(raw)
}

func (i NodeIndex) MarshalText() ([]byte, error) {
	return []byte(fmt.Sprintf("%016X", i)), nil
}

func (i *NodeIndex) UnmarshalText(b []byte) error {
	n, err := strconv.ParseUint(string(b), 16, 64)
	*i = NodeIndex(n)
	return err
}

func (r TransactionResult) MarshalText() ([]byte, error) {
	return []byte(resultNames[r]), nil
}

func (r *TransactionResult) UnmarshalText(b []byte) error {
	if result, ok := reverseResults[string(b)]; ok {
		*r = result
		return nil
	}
	return fmt.Errorf("Unknown TransactionResult: %s", string(b))
}

func (l LedgerEntryType) MarshalText() ([]byte, error) {
	return []byte(ledgerEntryNames[l]), nil
}

func (l *LedgerEntryType) UnmarshalText(b []byte) error {
	if leType, ok := ledgerEntryTypes[string(b)]; ok {
		*l = leType
		return nil
	}
	return fmt.Errorf("Unknown LedgerEntryType: %s", string(b))
}

func (t TransactionType) MarshalText() ([]byte, error) {
	return []byte(txNames[t]), nil
}

func (t *TransactionType) UnmarshalText(b []byte) error {
	if txType, ok := txTypes[string(b)]; ok {
		*t = txType
		return nil
	}
	return fmt.Errorf("Unknown TransactionType: %s", string(b))
}

func (t RippleTime) MarshalJSON() ([]byte, error) {
	return []byte(strconv.FormatUint(uint64(t.Uint32()), 10)), nil
}

func (t *RippleTime) UnmarshalJSON(b []byte) error {
	n, err := strconv.ParseUint(string(b), 10, 32)
	if err != nil {
		return err
	}
	t.SetUint32(uint32(n))
	return nil
}

func (t RippleHumanTime) MarshalJSON() ([]byte, error) {
	return []byte(`"` + t.String() + `"`), nil
}

func (t *RippleHumanTime) UnmarshalJSON(b []byte) error {
	t.RippleTime = &RippleTime{}
	return t.SetString(string(b[1 : len(b)-1]))
}

func (v *Value) MarshalText() ([]byte, error) {
	if v.Native {
		return []byte(strconv.FormatUint(v.Num, 10)), nil
	}
	return []byte(v.String()), nil
}

func (v *Value) UnmarshalText(b []byte) error {
	value, err := NewValue(string(b), true)
	if err != nil {
		return err
	}
	*v = *value
	return nil
}

type nonNativeValue Value

func (v *nonNativeValue) UnmarshalText(b []byte) error {
	value, err := NewValue(string(b), false)
	if err != nil {
		return err
	}
	*v = nonNativeValue(*value)
	return nil
	// return (*Value)(v).Parse(string(b))
}

func (v *nonNativeValue) MarshalText() ([]byte, error) {
	return (*Value)(v).MarshalText()
}

type amountJSON struct {
	Value    *nonNativeValue `json:"value"`
	Currency Currency        `json:"currency"`
	Issuer   Account         `json:"issuer"`
}

func (a Amount) MarshalJSON() ([]byte, error) {
	if a.Native {
		return []byte(`"` + strconv.FormatUint(a.Num, 10) + `"`), nil
	}
	return json.Marshal(amountJSON{(*nonNativeValue)(a.Value), a.Currency, a.Issuer})
}

func (a *Amount) UnmarshalJSON(b []byte) (err error) {
	if b[0] != '{' {
		a.Value = new(Value)
		return json.Unmarshal(b, a.Value)
	}
	var dummy amountJSON
	if err := json.Unmarshal(b, &dummy); err != nil {
		return err
	}
	a.Value, a.Currency, a.Issuer = (*Value)(dummy.Value), dummy.Currency, dummy.Issuer
	return nil
}

func (c Currency) MarshalText() ([]byte, error) {
	return []byte(c.String()), nil
}

func (c *Currency) UnmarshalText(text []byte) error {
	var err error
	*c, err = NewCurrency(string(text))
	return err
}

func (h Hash128) MarshalText() ([]byte, error) {
	return b2h(h[:]), nil
}

func (h *Hash128) UnmarshalText(b []byte) error {
	_, err := hex.Decode(h[:], b)
	return err
}

func (h Hash160) MarshalText() ([]byte, error) {
	return b2h(h[:]), nil
}

func (h *Hash160) UnmarshalText(b []byte) error {
	_, err := hex.Decode(h[:], b)
	return err
}

func (h Hash256) MarshalText() ([]byte, error) {
	return b2h(h[:]), nil
}

func (h *Hash256) UnmarshalText(b []byte) error {
	_, err := hex.Decode(h[:], b)
	return err
}

func (a Account) MarshalText() ([]byte, error) {
	address, err := a.Hash()
	if err != nil {
		return nil, err
	}
	return address.MarshalText()
}

// Expects base58-encoded account id
func (a *Account) UnmarshalText(b []byte) error {
	account, err := NewAccountFromAddress(string(b))
	if err != nil {
		return err
	}
	copy(a[:], account[:])
	return nil
}

func (r RegularKey) MarshalText() ([]byte, error) {
	address, err := r.Hash()
	if err != nil {
		return nil, err
	}
	return address.MarshalText()
}

// Expects base58-encoded account id
func (r *RegularKey) UnmarshalText(b []byte) error {
	account, err := NewRegularKeyFromAddress(string(b))
	if err != nil {
		return err
	}
	copy(r[:], account[:])
	return nil
}

func (v VariableLength) MarshalText() ([]byte, error) {
	return b2h(v), nil
}

// Expects variable length hex
func (v *VariableLength) UnmarshalText(b []byte) error {
	var err error
	*v, err = hex.DecodeString(string(b))
	return err
}

func (p PublicKey) MarshalText() ([]byte, error) {
	return b2h(p[:]), nil
}

// Expects public key hex
func (p *PublicKey) UnmarshalText(b []byte) error {
	_, err := hex.Decode(p[:], b)
	return err
}

type affectedNodeJSON AffectedNode

type affectedFields struct {
	*affectedNodeJSON
	FinalFields    *json.RawMessage
	PreviousFields *json.RawMessage
	NewFields      *json.RawMessage
}

func (n *AffectedNode) UnmarshalJSON(b []byte) error {
	affected := affectedFields{
		affectedNodeJSON: (*affectedNodeJSON)(n),
	}
	if err := json.Unmarshal(b, &affected); err != nil {
		return err
	}
	if affected.FinalFields != nil {
		n.FinalFields = FieldsFactory[n.LedgerEntryType]()
		if err := json.Unmarshal(*affected.FinalFields, n.FinalFields); err != nil {
			return err
		}
	}
	if affected.PreviousFields != nil {
		n.PreviousFields = FieldsFactory[n.LedgerEntryType]()
		if err := json.Unmarshal(*affected.PreviousFields, n.PreviousFields); err != nil {
			return err
		}
	}
	if affected.NewFields != nil {
		n.NewFields = FieldsFactory[n.LedgerEntryType]()
		if err := json.Unmarshal(*affected.NewFields, n.NewFields); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Marshal non-standard currencies as hex so they can be unmarshaled again<commit_after>package data

// Evil things happen here. Rippled needs a V2 API...

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
)

type ledgerJSON Ledger

// adds all the legacy fields
type ledgerExtraJSON struct {
	ledgerJSON
	HumanCloseTime *RippleHumanTime `json:"close_time_human"`
	Hash           Hash256          `json:"hash"`
	LedgerHash     Hash256          `json:"ledger_hash"`
	TotalCoins     uint64           `json:"totalCoins,string"`
	SequenceNumber uint32           `json:"seqNum,string"`
}

func (l Ledger) MarshalJSON() ([]byte, error) {
	return json.Marshal(ledgerExtraJSON{
		ledgerJSON:     ledgerJSON(l),
		HumanCloseTime: l.CloseTime.Human(),
		Hash:           l.Hash(),
		LedgerHash:     l.Hash(),
		TotalCoins:     l.TotalXRP,
		SequenceNumber: l.LedgerSequence,
	})
}

func (l *Ledger) UnmarshalJSON(b []byte) error {
	var ledger ledgerExtraJSON
	if err := json.Unmarshal(b, &ledger); err != nil {
		return err
	}
	*l = Ledger(ledger.ledgerJSON)
	l.SetHash(ledger.Hash[:])
	return nil
}

// Wrapper types to enable second level of marshalling
// when found in ledger API call
type txmLedger struct {
	MetaData MetaData `json:"metaData"`
}

// Wrapper types to enable second level of marshalling
// when found in tx API call
type txmNormal TransactionWithMetaData

var (
	txmTransactionTypeRegex = regexp.MustCompile(`"TransactionType":\s*"(\w+)"`)
	txmHashRegex            = regexp.MustCompile(`"hash":\s*"(\w+)"`)
	txmMetaTypeRegex        = regexp.MustCompile(`"(meta|metaData)"`)
)

func (txm *TransactionWithMetaData) UnmarshalJSON(b []byte) error {
	txTypeMatch := txmTransactionTypeRegex.FindStringSubmatch(string(b))
	hashMatch := txmHashRegex.FindStringSubmatch(string(b))
	metaTypeMatch := txmMetaTypeRegex.FindStringSubmatch(string(b))
	var txType, hash, metaType string
	if txTypeMatch == nil {
		return fmt.Errorf("Not a valid transaction with metadata: Missing TransactionType")
	}
	txType = txTypeMatch[1]
	if hashMatch == nil {
		return fmt.Errorf("Not a valid transaction with metadata: Missing Hash")
	}
	hash = hashMatch[1]
	if metaTypeMatch != nil {
		metaType = metaTypeMatch[1]
	}
	txm.Transaction = GetTxFactoryByType(txType)()
	h, err := hex.DecodeString(hash)
	if err != nil {
		return fmt.Errorf("Bad hash: %s", hash)
	}
	txm.SetHash(h)
	if err := json.Unmarshal(b, txm.Transaction); err != nil {
		return err
	}
	switch metaType {
	case "meta":
		return json.Unmarshal(b, (*txmNormal)(txm))
	case "metaData":
		var meta txmLedger
		if err := json.Unmarshal(b, &meta); err != nil {
			return err
		}
		txm.MetaData = meta.MetaData
		return nil
	default:
		return nil
	}
}

func (txm TransactionWithMetaData) marshalJSON() ([]byte, []byte, error) {
	tx, err := json.Marshal(txm.Transaction)
	if err != nil {
		return nil, nil, err
	}
	meta, err := json.Marshal(txm.MetaData)
	if err != nil {
		return nil, nil, err
	}
	return tx, meta, nil
}

type extractTxm struct {
	Tx   json.RawMessage `json:"transaction"`
	Meta json.RawMessage `json:"meta"`
}

const extractTxmFormat = `%s,"meta":%s}`

func UnmarshalTransactionWithMetadata(b []byte, txm *TransactionWithMetaData) error {
	var extract extractTxm
	if err := json.Unmarshal(b, &extract); err != nil {
		return err
	}
	raw := fmt.Sprintf(extractTxmFormat, extract.Tx[:len(extract.Tx)-1], extract.Meta)
	return json.Unmarshal([]byte(raw), txm)
}

const txmFormat = `%s,"hash":"%s","inLedger":%d,"ledger_index":%d,"meta":%s}`

func (txm TransactionWithMetaData) MarshalJSON() ([]byte, error) {
	tx, meta, err := txm.marshalJSON()
	if err != nil {
		return nil, err
	}
	out := fmt.Sprintf(txmFormat, string(tx[:len(tx)-1]), txm.Hash().String(), txm.LedgerSequence, txm.LedgerSequence, string(meta))
	return []byte(out), nil
}

const txmSliceFormat = `%s,"hash":"%s","metaData":%s}`

func (s TransactionSlice) MarshalJSON() ([]byte, error) {
	raw := make([]json.RawMessage, len(s))
	var err error
	var tx, meta []byte
	for i, txm := range s {
		if tx, meta, err = txm.marshalJSON(); err != nil {
			return nil, err
		}
		extra := fmt.Sprintf(txmSliceFormat, string(tx[:len(tx)-1]), txm.Hash().String(), meta)
		raw[i] = json.RawMessage(extra)
	}
	return json.Marshal(raw)
}

var (
	leTypeRegex  = regexp.MustCompile(`"LedgerEntryType":\s*"(\w+)"`)
	leIndexRegex = regexp.MustCompile(`"index":\s*"(\w+)"`)
)

func (l *LedgerEntrySlice) UnmarshalJSON(b []byte) error {
	var s []json.RawMessage
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	for _, raw := range s {
		leTypeMatch := leTypeRegex.FindStringSubmatch(string(raw))
		indexMatch := leIndexRegex.FindStringSubmatch(string(raw))
		if leTypeMatch == nil {
			return fmt.Errorf("Bad LedgerEntryType")
		}
		if indexMatch == nil {
			return fmt.Errorf("Missing LedgerEntry index")
		}
		le := GetLedgerEntryFactoryByType(leTypeMatch[1])()
		index, err := hex.DecodeString(indexMatch[1])
		if err != nil {
			return fmt.Errorf("Bad index: %s", index)
		}
		le.SetHash(index)
		if err := json.Unmarshal(raw, &le); err != nil {
			return err
		}
		*l = append(*l, le)
	}
	return nil
}

const leSliceFormat = `%s,"index":"%s"}`

func (s LedgerEntrySlice) MarshalJSON() ([]byte, error) {
	raw := make([]json.RawMessage, len(s))
	var err error
	for i, le := range s {
		if raw[i], err = json.Marshal(le); err != nil {
			return nil, err
		}
		extra := fmt.Sprintf(leSliceFormat, string(raw[i][:len(raw[i])-1]), le.Hash().String())
		raw[i] = json.RawMessage(extra)
	}
	return json.Marshal(raw)
}

func (i NodeIndex) MarshalText() ([]byte, error) {
	return []byte(fmt.Sprintf("%016X", i)), nil
}

func (i *NodeIndex) UnmarshalText(b []byte) error {
	n, err := strconv.ParseUint(string(b), 16, 64)
	*i = NodeIndex(n)
	return err
}

func (r TransactionResult) MarshalText() ([]byte, error) {
	return []byte(resultNames[r]), nil
}

func (r *TransactionResult) UnmarshalText(b []byte) error {
	if result, ok := reverseResults[string(b)]; ok {
		*r = result
		return nil
	}
	return fmt.Errorf("Unknown TransactionResult: %s", string(b))
}

func (l LedgerEntryType) MarshalText() ([]byte, error) {
	return []byte(ledgerEntryNames[l]), nil
}

func (l *LedgerEntryType) UnmarshalText(b []byte) error {
	if leType, ok := ledgerEntryTypes[string(b)]; ok {
		*l = leType
		return nil
	}
	return fmt.Errorf("Unknown LedgerEntryType: %s", string(b))
}

func (t TransactionType) MarshalText() ([]byte, error) {
	return []byte(txNames[t]), nil
}

func (t *TransactionType) UnmarshalText(b []byte) error {
	if txType, ok := txTypes[string(b)]; ok {
		*t = txType
		return nil
	}
	return fmt.Errorf("Unknown TransactionType: %s", string(b))
}

func (t RippleTime) MarshalJSON() ([]byte, error) {
	return []byte(strconv.FormatUint(uint64(t.Uint32()), 10)), nil
}

func (t *RippleTime) UnmarshalJSON(b []byte) error {
	n, err := strconv.ParseUint(string(b), 10, 32)
	if err != nil {
		return err
	}
	t.SetUint32(uint32(n))
	return nil
}

func (t RippleHumanTime) MarshalJSON() ([]byte, error) {
	return []byte(`"` + t.String() + `"`), nil
}

func (t *RippleHumanTime) UnmarshalJSON(b []byte) error {
	t.RippleTime = &RippleTime{}
	return t.SetString(string(b[1 : len(b)-1]))
}

func (v *Value) MarshalText() ([]byte, error) {
	if v.Native {
		return []byte(strconv.FormatUint(v.Num, 10)), nil
	}
	return []byte(v.String()), nil
}

func (v *Value) UnmarshalText(b []byte) error {
	value, err := NewValue(string(b), true)
	if err != nil {
		return err
	}
	*v = *value
	return nil
}

type nonNativeValue Value

func (v *nonNativeValue) UnmarshalText(b []byte) error {
	value, err := NewValue(string(b), false)
	if err != nil {
		return err
	}
	*v = nonNativeValue(*value)
	return nil
	// return (*Value)(v).Parse(string(b))
}

func (v *nonNativeValue) MarshalText() ([]byte, error) {
	return (*Value)(v).MarshalText()
}

type amountJSON struct {
	Value    *nonNativeValue `json:"value"`
	Currency Currency        `json:"currency"`
	Issuer   Account         `json:"issuer"`
}

func (a Amount) MarshalJSON() ([]byte, error) {
	if a.Native {
		return []byte(`"` + strconv.FormatUint(a.Num, 10) + `"`), nil
	}
	return json.Marshal(amountJSON{(*nonNativeValue)(a.Value), a.Currency, a.Issuer})
}

func (a *Amount) UnmarshalJSON(b []byte) (err error) {
	if b[0] != '{' {
		a.Value = new(Value)
		return json.Unmarshal(b, a.Value)
	}
	var dummy amountJSON
	if err := json.Unmarshal(b, &dummy); err != nil {
		return err
	}
	a.Value, a.Currency, a.Issuer = (*Value)(dummy.Value), dummy.Currency, dummy.Issuer
	return nil
}

func (c Currency) MarshalText() ([]byte, error) {
	return []byte(c.Machine()), nil
}

func (c *Currency) UnmarshalText(text []byte) error {
	var err error
	*c, err = NewCurrency(string(text))
	return err
}

func (h Hash128) MarshalText() ([]byte, error) {
	return b2h(h[:]), nil
}

func (h *Hash128) UnmarshalText(b []byte) error {
	_, err := hex.Decode(h[:], b)
	return err
}

func (h Hash160) MarshalText() ([]byte, error) {
	return b2h(h[:]), nil
}

func (h *Hash160) UnmarshalText(b []byte) error {
	_, err := hex.Decode(h[:], b)
	return err
}

func (h Hash256) MarshalText() ([]byte, error) {
	return b2h(h[:]), nil
}

func (h *Hash256) UnmarshalText(b []byte) error {
	_, err := hex.Decode(h[:], b)
	return err
}

func (a Account) MarshalText() ([]byte, error) {
	address, err := a.Hash()
	if err != nil {
		return nil, err
	}
	return address.MarshalText()
}

// Expects base58-encoded account id
func (a *Account) UnmarshalText(b []byte) error {
	account, err := NewAccountFromAddress(string(b))
	if err != nil {
		return err
	}
	copy(a[:], account[:])
	return nil
}

func (r RegularKey) MarshalText() ([]byte, error) {
	address, err := r.Hash()
	if err != nil {
		return nil, err
	}
	return address.MarshalText()
}

// Expects base58-encoded account id
func (r *RegularKey) UnmarshalText(b []byte) error {
	account, err := NewRegularKeyFromAddress(string(b))
	if err != nil {
		return err
	}
	copy(r[:], account[:])
	return nil
}

func (v VariableLength) MarshalText() ([]byte, error) {
	return b2h(v), nil
}

// Expects variable length hex
func (v *VariableLength) UnmarshalText(b []byte) error {
	var err error
	*v, err = hex.DecodeString(string(b))
	return err
}

func (p PublicKey) MarshalText() ([]byte, error) {
	return b2h(p[:]), nil
}

// Expects public key hex
func (p *PublicKey) UnmarshalText(b []byte) error {
	_, err := hex.Decode(p[:], b)
	return err
}

type affectedNodeJSON AffectedNode

type affectedFields struct {
	*affectedNodeJSON
	FinalFields    *json.RawMessage
	PreviousFields *json.RawMessage
	NewFields      *json.RawMessage
}

func (n *AffectedNode) UnmarshalJSON(b []byte) error {
	affected := affectedFields{
		affectedNodeJSON: (*affectedNodeJSON)(n),
	}
	if err := json.Unmarshal(b, &affected); err != nil {
		return err
	}
	if affected.FinalFields != nil {
		n.FinalFields = FieldsFactory[n.LedgerEntryType]()
		if err := json.Unmarshal(*affected.FinalFields, n.FinalFields); err != nil {
			return err
		}
	}
	if affected.PreviousFields != nil {
		n.PreviousFields = FieldsFactory[n.LedgerEntryType]()
		if err := json.Unmarshal(*affected.PreviousFields, n.PreviousFields); err != nil {
			return err
		}
	}
	if affected.NewFields != nil {
		n.NewFields = FieldsFactory[n.LedgerEntryType]()
		if err := json.Unmarshal(*affected.NewFields, n.NewFields); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package sgload

import (
	"fmt"
	"strings"
	"sync"
	"time"

	sgreplicate "github.com/couchbaselabs/sg-replicate"
)

type Reader struct {
	Agent
	SGChannels      []string // The Sync Gateway channels this reader is assigned to pull from
	NumDocsExpected int      // The total number of docs this reader is expected to pull
	BatchSize       int      // The number of docs to pull in batch (_changes feed and bulk_get)

}

func NewReader(wg *sync.WaitGroup, ID int, u UserCred, d DataStore, batchsize int) *Reader {

	reader := Reader{
		Agent: Agent{
			FinishedWg: wg,
			UserCred:   u,
			ID:         ID,
			DataStore:  d,
			BatchSize:  batchsize,
		},
	}

	return &reader

}

func (r *Reader) SetChannels(sgChannels []string) {
	r.SGChannels = sgChannels
}

func (r *Reader) SetNumDocsExpected(n int) {
	r.NumDocsExpected = n
}

func (r *Reader) SetBatchSize(batchSize int) {
	r.BatchSize = batchSize
}

func (r *Reader) pushPostRunTimingStats(numDocsPulled int, timeStartedCreatingDocs time.Time) {
	if r.StatsdClient == nil {
		return
	}
	delta := time.Since(timeStartedCreatingDocs)

	// How long it took for this reader to read all of its docs
	r.StatsdClient.Timing(
		statsdSampleRate,
		"get_all_documents",
		delta,
	)

	if numDocsPulled > 0 {
		// Average time it took to read each doc from
		// the changes feed and the doc itself
		deltaChangeAndDoc := time.Duration(int64(delta) / int64(numDocsPulled))
		r.StatsdClient.Timing(
			statsdSampleRate,
			"get_change_and_document",
			deltaChangeAndDoc,
		)
	}

	logger.Info("Reader finished", "agent.ID", r.ID, "numdocs", numDocsPulled)

}

func (r *Reader) Run() {

	since := StringSincer{}
	result := pullMoreDocsResult{}
	uniqueDocIdsPulled := map[string]struct{}{}
	var err error
	var timeStartedCreatingDocs time.Time

	defer r.FinishedWg.Done()
	defer func() {
		r.pushPostRunTimingStats(len(uniqueDocIdsPulled), timeStartedCreatingDocs)
	}()

	r.createSGUserIfNeeded(r.SGChannels)

	timeStartedCreatingDocs = time.Now()

	for {

		if r.isFinished(len(uniqueDocIdsPulled)) {
			break
		}
		result, err = r.pullMoreDocs(since)
		if err != nil {
			logger.Error("Error calling pullMoreDocs", "agent.ID", r.ID, "since", since, "err", err)
			panic(fmt.Sprintf("Error calling pullMoreDoc: %v", err))
		}
		since = result.since

		addNewUniqueDocIdsPulled(uniqueDocIdsPulled, result)

	}

}

func (r *Reader) isFinished(numDocsPulled int) bool {

	switch {
	case numDocsPulled > r.NumDocsExpected:
		panic(fmt.Sprintf("Reader was only expected to pull %d docs, but pulled %d.", r.NumDocsExpected, numDocsPulled))
	case numDocsPulled == r.NumDocsExpected:
		return true
	default:
		return false
	}

}

type pullMoreDocsResult struct {
	since        StringSincer
	uniqueDocIds map[string]sgreplicate.DocumentRevisionPair
}

func (r *Reader) pullMoreDocs(since Sincer) (pullMoreDocsResult, error) {

	// Create a retry sleeper which controls how many times to retry
	// and how long to wait in between retries
	numRetries := 14
	sleepMsBetweenRetry := 500
	retrySleeper := CreateDoublingSleeperFunc(numRetries, sleepMsBetweenRetry)

	// Create retry worker that knows how to do actual work
	retryWorker := func() (shouldRetry bool, err error, value interface{}) {

		result := pullMoreDocsResult{}

		changes, newSince, err := r.DataStore.Changes(since, r.BatchSize)
		if err != nil {
			return false, err, result
		}

		if len(changes.Results) == 0 {
			return true, nil, result
		}
		if newSince.Equals(since) {
			logger.Warn("Since value should have changed", "agent.ID", r.ID, "since", since, "newsince", newSince)
			return true, nil, result
		}

		// Strip out any changes with id "id":"_user/*"
		// since they are user docs and we don't care about them
		changes = stripUserDocChanges(changes)

		bulkGetRequest, uniqueDocIds := getBulkGetRequest(changes)

		docs, err := r.DataStore.BulkGetDocuments(bulkGetRequest)
		if err != nil {
			return false, err, result
		}
		if len(docs) != len(bulkGetRequest.Docs) {
			return false, fmt.Errorf("Expected %d docs, got %d", len(bulkGetRequest.Docs), len(docs)), result
		}

		docsMustBeInExpectedChannels(docs, r.SGChannels)

		result.since = newSince.(StringSincer)
		result.uniqueDocIds = uniqueDocIds
		return false, nil, result

	}

	// Invoke the retry worker / sleeper combo in a loop
	err, workerReturnVal := RetryLoop("pullMoreDocs", retryWorker, retrySleeper)
	if err != nil {
		return pullMoreDocsResult{}, err
	}

	return workerReturnVal.(pullMoreDocsResult), nil

}

func getBulkGetRequest(changes sgreplicate.Changes) (sgreplicate.BulkGetRequest, map[string]sgreplicate.DocumentRevisionPair) {

	uniqueDocIds := map[string]sgreplicate.DocumentRevisionPair{}

	bulkGetRequest := sgreplicate.BulkGetRequest{}
	docs := []sgreplicate.DocumentRevisionPair{}
	for _, change := range changes.Results {
		docRevPair := sgreplicate.DocumentRevisionPair{}
		docRevPair.Id = change.Id
		docRevPair.Revision = change.ChangedRevs[0].Revision
		docs = append(docs, docRevPair)
		uniqueDocIds[docRevPair.Id] = docRevPair
	}
	bulkGetRequest.Docs = docs

	// Validate expectation that doc id's only appear once in the changes feed response
	if len(uniqueDocIds) != len(docs) {
		logger.Error(
			"len(uniqueDocIds) != len(docs)",
			"len(uniqueDocIds)",
			len(uniqueDocIds),
			"len(docs)",
			len(docs),
		)
		for docId, docRevPair := range uniqueDocIds {
			logger.Error("uniqueDocIds", "docId", docId, "docRevPair", docRevPair)
		}
		for _, doc := range docs {
			logger.Error("docs", "doc", doc, "doc.id", doc.Id, "doc.rev", doc.Revision)
		}
		panic(
			fmt.Sprintf(
				"len(uniqueDocIds) != len(docs), %d != %d",
				len(uniqueDocIds),
				len(docs),
			),
		)

	}

	return bulkGetRequest, uniqueDocIds

}

func stripUserDocChanges(changes sgreplicate.Changes) (changesStripped sgreplicate.Changes) {
	changesStripped.LastSequence = changes.LastSequence

	for _, change := range changes.Results {
		if strings.Contains(change.Id, "_user") {
			continue
		}
		changesStripped.Results = append(changesStripped.Results, change)

	}

	return changesStripped
}

// A retry sleeper is called back by the retry loop and passed
// the current retryCount, and should return the amount of milliseconds
// that the retry should sleep.
type RetrySleeper func(retryCount int) (shouldContinue bool, timeTosleepMs int)

// A RetryWorker encapsulates the work being done in a Retry Loop.  The shouldRetry
// return value determines whether the worker will retry, regardless of the err value.
// If the worker has exceeded it's retry attempts, then it will not be called again
// even if it returns shouldRetry = true.
type RetryWorker func() (shouldRetry bool, err error, value interface{})

func RetryLoop(description string, worker RetryWorker, sleeper RetrySleeper) (error, interface{}) {

	numAttempts := 1

	for {
		shouldRetry, err, value := worker()
		if !shouldRetry {
			if err != nil {
				return err, nil
			}
			return nil, value
		}
		shouldContinue, sleepMs := sleeper(numAttempts)
		if !shouldContinue {
			if err == nil {
				err = fmt.Errorf("RetryLoop for %v giving up after %v attempts", description, numAttempts)
			}
			return err, value
		}

		<-time.After(time.Millisecond * time.Duration(sleepMs))

		numAttempts += 1

	}
}

// Create a RetrySleeper that will double the retry time on every iteration and
// use the given parameters
func CreateDoublingSleeperFunc(maxNumAttempts, initialTimeToSleepMs int) RetrySleeper {

	timeToSleepMs := initialTimeToSleepMs

	sleeper := func(numAttempts int) (bool, int) {
		if numAttempts > maxNumAttempts {
			return false, -1
		}
		if numAttempts > 1 {
			timeToSleepMs *= 2
		}
		return true, timeToSleepMs
	}
	return sleeper

}

func addNewUniqueDocIdsPulled(uniqueDocIdsPulled map[string]struct{}, r pullMoreDocsResult) {

	for docId, _ := range r.uniqueDocIds {
		uniqueDocIdsPulled[docId] = struct{}{}
	}

}
<commit_msg>Revert "Increase number of retries before giving up"<commit_after>package sgload

import (
	"fmt"
	"strings"
	"sync"
	"time"

	sgreplicate "github.com/couchbaselabs/sg-replicate"
)

type Reader struct {
	Agent
	SGChannels      []string // The Sync Gateway channels this reader is assigned to pull from
	NumDocsExpected int      // The total number of docs this reader is expected to pull
	BatchSize       int      // The number of docs to pull in batch (_changes feed and bulk_get)

}

func NewReader(wg *sync.WaitGroup, ID int, u UserCred, d DataStore, batchsize int) *Reader {

	reader := Reader{
		Agent: Agent{
			FinishedWg: wg,
			UserCred:   u,
			ID:         ID,
			DataStore:  d,
			BatchSize:  batchsize,
		},
	}

	return &reader

}

func (r *Reader) SetChannels(sgChannels []string) {
	r.SGChannels = sgChannels
}

func (r *Reader) SetNumDocsExpected(n int) {
	r.NumDocsExpected = n
}

func (r *Reader) SetBatchSize(batchSize int) {
	r.BatchSize = batchSize
}

func (r *Reader) pushPostRunTimingStats(numDocsPulled int, timeStartedCreatingDocs time.Time) {
	if r.StatsdClient == nil {
		return
	}
	delta := time.Since(timeStartedCreatingDocs)

	// How long it took for this reader to read all of its docs
	r.StatsdClient.Timing(
		statsdSampleRate,
		"get_all_documents",
		delta,
	)

	if numDocsPulled > 0 {
		// Average time it took to read each doc from
		// the changes feed and the doc itself
		deltaChangeAndDoc := time.Duration(int64(delta) / int64(numDocsPulled))
		r.StatsdClient.Timing(
			statsdSampleRate,
			"get_change_and_document",
			deltaChangeAndDoc,
		)
	}

	logger.Info("Reader finished", "agent.ID", r.ID, "numdocs", numDocsPulled)

}

func (r *Reader) Run() {

	since := StringSincer{}
	result := pullMoreDocsResult{}
	uniqueDocIdsPulled := map[string]struct{}{}
	var err error
	var timeStartedCreatingDocs time.Time

	defer r.FinishedWg.Done()
	defer func() {
		r.pushPostRunTimingStats(len(uniqueDocIdsPulled), timeStartedCreatingDocs)
	}()

	r.createSGUserIfNeeded(r.SGChannels)

	timeStartedCreatingDocs = time.Now()

	for {

		if r.isFinished(len(uniqueDocIdsPulled)) {
			break
		}
		result, err = r.pullMoreDocs(since)
		if err != nil {
			logger.Error("Error calling pullMoreDocs", "agent.ID", r.ID, "since", since, "err", err)
			panic(fmt.Sprintf("Error calling pullMoreDoc: %v", err))
		}
		since = result.since

		addNewUniqueDocIdsPulled(uniqueDocIdsPulled, result)

	}

}

func (r *Reader) isFinished(numDocsPulled int) bool {

	switch {
	case numDocsPulled > r.NumDocsExpected:
		panic(fmt.Sprintf("Reader was only expected to pull %d docs, but pulled %d.", r.NumDocsExpected, numDocsPulled))
	case numDocsPulled == r.NumDocsExpected:
		return true
	default:
		return false
	}

}

type pullMoreDocsResult struct {
	since        StringSincer
	uniqueDocIds map[string]sgreplicate.DocumentRevisionPair
}

func (r *Reader) pullMoreDocs(since Sincer) (pullMoreDocsResult, error) {

	// Create a retry sleeper which controls how many times to retry
	// and how long to wait in between retries
	numRetries := 7
	sleepMsBetweenRetry := 500
	retrySleeper := CreateDoublingSleeperFunc(numRetries, sleepMsBetweenRetry)

	// Create retry worker that knows how to do actual work
	retryWorker := func() (shouldRetry bool, err error, value interface{}) {

		result := pullMoreDocsResult{}

		changes, newSince, err := r.DataStore.Changes(since, r.BatchSize)
		if err != nil {
			return false, err, result
		}

		if len(changes.Results) == 0 {
			return true, nil, result
		}
		if newSince.Equals(since) {
			logger.Warn("Since value should have changed", "agent.ID", r.ID, "since", since, "newsince", newSince)
			return true, nil, result
		}

		// Strip out any changes with id "id":"_user/*"
		// since they are user docs and we don't care about them
		changes = stripUserDocChanges(changes)

		bulkGetRequest, uniqueDocIds := getBulkGetRequest(changes)

		docs, err := r.DataStore.BulkGetDocuments(bulkGetRequest)
		if err != nil {
			return false, err, result
		}
		if len(docs) != len(bulkGetRequest.Docs) {
			return false, fmt.Errorf("Expected %d docs, got %d", len(bulkGetRequest.Docs), len(docs)), result
		}

		docsMustBeInExpectedChannels(docs, r.SGChannels)

		result.since = newSince.(StringSincer)
		result.uniqueDocIds = uniqueDocIds
		return false, nil, result

	}

	// Invoke the retry worker / sleeper combo in a loop
	err, workerReturnVal := RetryLoop("pullMoreDocs", retryWorker, retrySleeper)
	if err != nil {
		return pullMoreDocsResult{}, err
	}

	return workerReturnVal.(pullMoreDocsResult), nil

}

func getBulkGetRequest(changes sgreplicate.Changes) (sgreplicate.BulkGetRequest, map[string]sgreplicate.DocumentRevisionPair) {

	uniqueDocIds := map[string]sgreplicate.DocumentRevisionPair{}

	bulkGetRequest := sgreplicate.BulkGetRequest{}
	docs := []sgreplicate.DocumentRevisionPair{}
	for _, change := range changes.Results {
		docRevPair := sgreplicate.DocumentRevisionPair{}
		docRevPair.Id = change.Id
		docRevPair.Revision = change.ChangedRevs[0].Revision
		docs = append(docs, docRevPair)
		uniqueDocIds[docRevPair.Id] = docRevPair
	}
	bulkGetRequest.Docs = docs

	// Validate expectation that doc id's only appear once in the changes feed response
	if len(uniqueDocIds) != len(docs) {
		logger.Error(
			"len(uniqueDocIds) != len(docs)",
			"len(uniqueDocIds)",
			len(uniqueDocIds),
			"len(docs)",
			len(docs),
		)
		for docId, docRevPair := range uniqueDocIds {
			logger.Error("uniqueDocIds", "docId", docId, "docRevPair", docRevPair)
		}
		for _, doc := range docs {
			logger.Error("docs", "doc", doc, "doc.id", doc.Id, "doc.rev", doc.Revision)
		}
		panic(
			fmt.Sprintf(
				"len(uniqueDocIds) != len(docs), %d != %d",
				len(uniqueDocIds),
				len(docs),
			),
		)

	}

	return bulkGetRequest, uniqueDocIds

}

func stripUserDocChanges(changes sgreplicate.Changes) (changesStripped sgreplicate.Changes) {
	changesStripped.LastSequence = changes.LastSequence

	for _, change := range changes.Results {
		if strings.Contains(change.Id, "_user") {
			continue
		}
		changesStripped.Results = append(changesStripped.Results, change)

	}

	return changesStripped
}

// A retry sleeper is called back by the retry loop and passed
// the current retryCount, and should return the amount of milliseconds
// that the retry should sleep.
type RetrySleeper func(retryCount int) (shouldContinue bool, timeTosleepMs int)

// A RetryWorker encapsulates the work being done in a Retry Loop.  The shouldRetry
// return value determines whether the worker will retry, regardless of the err value.
// If the worker has exceeded it's retry attempts, then it will not be called again
// even if it returns shouldRetry = true.
type RetryWorker func() (shouldRetry bool, err error, value interface{})

func RetryLoop(description string, worker RetryWorker, sleeper RetrySleeper) (error, interface{}) {

	numAttempts := 1

	for {
		shouldRetry, err, value := worker()
		if !shouldRetry {
			if err != nil {
				return err, nil
			}
			return nil, value
		}
		shouldContinue, sleepMs := sleeper(numAttempts)
		if !shouldContinue {
			if err == nil {
				err = fmt.Errorf("RetryLoop for %v giving up after %v attempts", description, numAttempts)
			}
			return err, value
		}

		<-time.After(time.Millisecond * time.Duration(sleepMs))

		numAttempts += 1

	}
}

// Create a RetrySleeper that will double the retry time on every iteration and
// use the given parameters
func CreateDoublingSleeperFunc(maxNumAttempts, initialTimeToSleepMs int) RetrySleeper {

	timeToSleepMs := initialTimeToSleepMs

	sleeper := func(numAttempts int) (bool, int) {
		if numAttempts > maxNumAttempts {
			return false, -1
		}
		if numAttempts > 1 {
			timeToSleepMs *= 2
		}
		return true, timeToSleepMs
	}
	return sleeper

}

func addNewUniqueDocIdsPulled(uniqueDocIdsPulled map[string]struct{}, r pullMoreDocsResult) {

	for docId, _ := range r.uniqueDocIds {
		uniqueDocIdsPulled[docId] = struct{}{}
	}

}
<|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
//
// 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 serial

import "io"
import "os"

func openInternal(options OpenOptions) (io.ReadWriteCloser, os.Error) {
	return nil, os.NewError("Not implemented.")
}

<commit_msg>Added some constants.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
//
// 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 serial

import "io"
import "os"

// OS-specific constants.
const (
	// sys/termios.h
	B9600 = 9600
	B14400 = 14400
	B19200 = 19200

	// sys/termios.h
	CS5 = 0x00000000
	CS6 = 0x00000100
	CS7 = 0x00000200
	CS8 = 0x00000300
	CLOCAL = 0x00008000
	CREAD = 0x00000800
	IGNPAR = 0x00000004
	VMIN = Tcflag_t (16);
	VTIME = Tcflag_t (17);

	NCCS = 20;
)

func openInternal(options OpenOptions) (io.ReadWriteCloser, os.Error) {
	return nil, os.NewError("Not implemented.")
}

<|endoftext|>