<commit_before>package bot

import (
	"sync"

	"github.com/jmoiron/sqlx"
)

type dict map[string]string

// The Dictionary is a glorified string/string map that's kept in sync with a database table.
type Dictionary struct {
	db    *sqlx.DB
	log   Logger
	data  dict
	mutex sync.RWMutex
}

func NewDictionary(db *sqlx.DB, log Logger) *Dictionary {
	return &Dictionary{db, log, make(dict), sync.RWMutex{}}
}

func (self *Dictionary) Keys() []string {
	self.mutex.RLock()
	defer self.mutex.RUnlock()

	list := make([]string, len(self.data))
	idx := 0

	for key, _ := range self.data {
		list[idx] = key
		idx = idx + 1
	}

	return list
}

func (self *Dictionary) Add(key string, value string) {
	self.mutex.Lock()
	defer self.mutex.Unlock()

	_, exists := self.data[key]
	if !exists {
		self.data[key] = value

		_, err := self.db.Exec("INSERT INTO dictionary (keyname, value) VALUES (?, ?)", key, value)
		if err != nil {
			self.log.Fatal("Could not add dictionary entry '" + key + "' to the database: " + err.Error())
		}

		self.log.Debug("Added dictionary entry '%s' as '%s'.", key, value)
	}
}

func (self *Dictionary) Set(key string, value string) {
	self.mutex.Lock()

	_, exists := self.data[key]
	if !exists {
		self.mutex.Unlock()
		self.Add(key, value)
		return
	}

	_, err := self.db.Exec("UPDATE dictionary SET value = ? WHERE keyname = ?", value, key)
	if err != nil {
		self.log.Fatal("Could not update dictionary entry '" + key + "' in the database: " + err.Error())
	}

	self.log.Debug("Updated dictionary entry '%s' with '%s'.", key, value)
	self.mutex.Unlock()
}

func (self *Dictionary) Get(key string) string {
	self.mutex.RLock()
	defer self.mutex.RUnlock()

	value, exists := self.data[key]
	if exists {
		return value
	}

	return ""
}

func (self *Dictionary) Has(key string) bool {
	self.mutex.RLock()
	defer self.mutex.RUnlock()

	_, exists := self.data[key]
	return exists
}

func (self *Dictionary) Delete(key string) *Dictionary {
	_, exists := self.data[key]
	if exists {
		delete(self.data, key)

		_, err := self.db.Exec("DELETE FROM dictionary WHERE keyname = ?", key)
		if err != nil {
			self.log.Fatal("Could not remove dictionary entry '" + key + "' from the database: " + err.Error())
		}

		self.log.Debug("Deleted dictionary entry '%s'.", key)
	}

	return self
}

type dictRow struct {
	Keyname string
	Value   string
}

func (self *Dictionary) load() {
	self.mutex.Lock()
	defer self.mutex.Unlock()

	list := make([]dictRow, 0)
	self.db.Select(&list, "SELECT * FROM dictionary ORDER BY keyname")

	for _, item := range list {
		self.data[item.Keyname] = item.Value
	}

	self.log.Debug("Loaded %d dictionary entries.", len(list))
}
<commit_msg>set the in-memory copy of dictionary values after updating them<commit_after>package bot

import (
	"sync"

	"github.com/jmoiron/sqlx"
)

type dict map[string]string

// The Dictionary is a glorified string/string map that's kept in sync with a database table.
type Dictionary struct {
	db    *sqlx.DB
	log   Logger
	data  dict
	mutex sync.RWMutex
}

func NewDictionary(db *sqlx.DB, log Logger) *Dictionary {
	return &Dictionary{db, log, make(dict), sync.RWMutex{}}
}

func (self *Dictionary) Keys() []string {
	self.mutex.RLock()
	defer self.mutex.RUnlock()

	list := make([]string, len(self.data))
	idx := 0

	for key, _ := range self.data {
		list[idx] = key
		idx = idx + 1
	}

	return list
}

func (self *Dictionary) Add(key string, value string) {
	self.mutex.Lock()
	defer self.mutex.Unlock()

	_, exists := self.data[key]
	if !exists {
		self.data[key] = value

		_, err := self.db.Exec("INSERT INTO dictionary (keyname, value) VALUES (?, ?)", key, value)
		if err != nil {
			self.log.Fatal("Could not add dictionary entry '" + key + "' to the database: " + err.Error())
		}

		self.log.Debug("Added dictionary entry '%s' as '%s'.", key, value)
	}
}

func (self *Dictionary) Set(key string, value string) {
	self.mutex.Lock()

	_, exists := self.data[key]
	if !exists {
		self.mutex.Unlock()
		self.Add(key, value)
		return
	}

	self.data[key] = value

	_, err := self.db.Exec("UPDATE dictionary SET value = ? WHERE keyname = ?", value, key)
	if err != nil {
		self.log.Fatal("Could not update dictionary entry '" + key + "' in the database: " + err.Error())
	}

	self.log.Debug("Updated dictionary entry '%s' with '%s'.", key, value)
	self.mutex.Unlock()
}

func (self *Dictionary) Get(key string) string {
	self.mutex.RLock()
	defer self.mutex.RUnlock()

	value, exists := self.data[key]
	if exists {
		return value
	}

	return ""
}

func (self *Dictionary) Has(key string) bool {
	self.mutex.RLock()
	defer self.mutex.RUnlock()

	_, exists := self.data[key]
	return exists
}

func (self *Dictionary) Delete(key string) *Dictionary {
	_, exists := self.data[key]
	if exists {
		delete(self.data, key)

		_, err := self.db.Exec("DELETE FROM dictionary WHERE keyname = ?", key)
		if err != nil {
			self.log.Fatal("Could not remove dictionary entry '" + key + "' from the database: " + err.Error())
		}

		self.log.Debug("Deleted dictionary entry '%s'.", key)
	}

	return self
}

type dictRow struct {
	Keyname string
	Value   string
}

func (self *Dictionary) load() {
	self.mutex.Lock()
	defer self.mutex.Unlock()

	list := make([]dictRow, 0)
	self.db.Select(&list, "SELECT * FROM dictionary ORDER BY keyname")

	for _, item := range list {
		self.data[item.Keyname] = item.Value
	}

	self.log.Debug("Loaded %d dictionary entries.", len(list))
}
<|endoftext|>


<commit_before>package core

import (
	"errors"
	"fmt"
	"math/big"
	"strings"
	"time"

	"github.com/OpenBazaar/openbazaar-go/pb"
	"github.com/OpenBazaar/openbazaar-go/repo"
	"github.com/OpenBazaar/wallet-interface"
	"github.com/btcsuite/btcutil"
)

// DefaultCurrencyDivisibility is the Divisibility of the Currency if not
// defined otherwise
const DefaultCurrencyDivisibility uint32 = 8

type SpendRequest struct {
	decodedAddress btcutil.Address

	Amount                 string                   `json:"amount"`
	Currency               *repo.CurrencyDefinition `json:"currency"`
	CurrencyCode           string                   `json:"currencyCode"`
	Address                string                   `json:"address"`
	FeeLevel               string                   `json:"feeLevel"`
	Memo                   string                   `json:"memo"`
	OrderID                string                   `json:"orderId"`
	RequireAssociatedOrder bool                     `json:"requireOrder"`
	SpendAll               bool                     `json:"spendAll"`
}

type SpendResponse struct {
	Amount             string                   `json:"amount"`
	ConfirmedBalance   string                   `json:"confirmedBalance"`
	UnconfirmedBalance string                   `json:"unconfirmedBalance"`
	Currency           *repo.CurrencyDefinition `json:"currency"`
	Memo               string                   `json:"memo"`
	OrderID            string                   `json:"orderId"`
	Timestamp          time.Time                `json:"timestamp"`
	Txid               string                   `json:"txid"`
	PeerID             string                   `json:"-"`
	ConsumedInput      bool                     `json:"-"`
}

// Spend will attempt to move funds from the node to the destination address described in the
// SpendRequest for the amount indicated.
func (n *OpenBazaarNode) Spend(args *SpendRequest) (*SpendResponse, error) {
	var (
		feeLevel wallet.FeeLevel
		peerID   string

		amt        = new(big.Int)
		lookupCode = args.CurrencyCode
	)

	if lookupCode == "" && args.Currency != nil {
		lookupCode = args.Currency.Code.String()
	}
	var currencyDef, err = repo.LoadCurrencyDefinitions().Lookup(lookupCode)
	if err != nil {
		return nil, repo.ErrCurrencyDefinitionUndefined
	}
	if args.Currency != nil && currencyDef.Divisibility != args.Currency.Divisibility {
		currencyDef.Divisibility = args.Currency.Divisibility
		if err := currencyDef.Valid(); err != nil {
			return nil, err
		}
	}

	amt, ok := amt.SetString(args.Amount, 10)
	if !ok {
		return nil, ErrInvalidAmount
	}

	wal, err := n.Multiwallet.WalletForCurrencyCode(args.CurrencyCode)
	if err != nil {
		return nil, ErrUnknownWallet
	}

	addr, err := wal.DecodeAddress(args.Address)
	if err != nil {
		return nil, ErrInvalidSpendAddress
	}
	args.decodedAddress = addr

	contract, err := n.getOrderContractBySpendRequest(args)
	if err != nil && args.RequireAssociatedOrder {
		return nil, ErrOrderNotFound
	}

	switch strings.ToUpper(args.FeeLevel) {
	case "PRIORITY":
		feeLevel = wallet.PRIOIRTY
	case "NORMAL":
		feeLevel = wallet.NORMAL
	case "ECONOMIC":
		feeLevel = wallet.ECONOMIC
	default:
		feeLevel = wallet.NORMAL
	}

	txid, err := wal.Spend(*amt, addr, feeLevel, args.OrderID, args.SpendAll)
	if err != nil {
		switch {
		case err == wallet.ErrorInsuffientFunds:
			return nil, ErrInsufficientFunds
		case err == wallet.ErrorDustAmount:
			return nil, ErrSpendAmountIsDust
		default:
			return nil, err
		}
	}

	txn, err := wal.GetTransaction(*txid)
	if err != nil {
		return nil, fmt.Errorf("failed retrieving new wallet balance: %s", err)
	}

	var (
		thumbnail string
		title     string
		memo      = args.Memo
		toAddress = args.Address
	)

	if txn.ToAddress != "" {
		toAddress = txn.ToAddress
	}

	if contract != nil && contract.VendorListings[0] != nil {
		if contract.VendorListings[0].Item != nil && len(contract.VendorListings[0].Item.Images) > 0 {
			thumbnail = contract.VendorListings[0].Item.Images[0].Tiny
			title = contract.VendorListings[0].Item.Title
		}
		if contract.VendorListings[0].VendorID != nil {
			peerID = contract.VendorListings[0].VendorID.PeerID
		}
	}
	if memo == "" && title != "" {
		memo = title
	}

	if err := n.Datastore.TxMetadata().Put(repo.Metadata{
		Txid:       txid.String(),
		Address:    toAddress,
		Memo:       memo,
		OrderId:    args.OrderID,
		Thumbnail:  thumbnail,
		CanBumpFee: false,
	}); err != nil {
		return nil, fmt.Errorf("failed persisting transaction metadata: %s", err)
	}

	confirmed, unconfirmed := wal.Balance()
	defn, err := repo.LoadCurrencyDefinitions().Lookup(wal.CurrencyCode())
	if err != nil {
		return nil, fmt.Errorf("wallet currency not found in dictionary")
	}

	return &SpendResponse{
		Txid:               txid.String(),
		ConfirmedBalance:   confirmed.Value.String(),
		UnconfirmedBalance: unconfirmed.Value.String(),
		Currency:           defn,
		Amount:             strings.TrimPrefix(txn.Value, "-"),
		Timestamp:          txn.Timestamp,
		Memo:               memo,
		OrderID:            args.OrderID,
		PeerID:             peerID,
	}, nil
}

func (n *OpenBazaarNode) getOrderContractBySpendRequest(args *SpendRequest) (*pb.RicardianContract, error) {
	var errorStr = "unable to find order from order id or spend address"
	if args.OrderID != "" {
		contract, _, _, _, _, _, err := n.Datastore.Purchases().GetByOrderId(args.OrderID)
		if err != nil {
			return nil, fmt.Errorf("%s: %s", errorStr, err)
		}
		if contract != nil {
			return contract, nil
		}
	}

	if args.decodedAddress != nil {
		contract, _, _, _, err := n.Datastore.Purchases().GetByPaymentAddress(args.decodedAddress)
		if err != nil {
			return nil, fmt.Errorf("%s: %s", errorStr, err)
		}
		if contract != nil {
			return contract, nil
		}
	}

	return nil, errors.New(errorStr)
}
<commit_msg>Fix wallet lookup value; Protect against spends from wrong network<commit_after>package core

import (
	"errors"
	"fmt"
	"math/big"
	"strings"
	"time"

	"github.com/OpenBazaar/openbazaar-go/pb"
	"github.com/OpenBazaar/openbazaar-go/repo"
	"github.com/OpenBazaar/wallet-interface"
	"github.com/btcsuite/btcutil"
)

// DefaultCurrencyDivisibility is the Divisibility of the Currency if not
// defined otherwise
const DefaultCurrencyDivisibility uint32 = 8

type SpendRequest struct {
	decodedAddress btcutil.Address

	Amount                 string                   `json:"amount"`
	Currency               *repo.CurrencyDefinition `json:"currency"`
	CurrencyCode           string                   `json:"currencyCode"`
	Address                string                   `json:"address"`
	FeeLevel               string                   `json:"feeLevel"`
	Memo                   string                   `json:"memo"`
	OrderID                string                   `json:"orderId"`
	RequireAssociatedOrder bool                     `json:"requireOrder"`
	SpendAll               bool                     `json:"spendAll"`
}

type SpendResponse struct {
	Amount             string                   `json:"amount"`
	ConfirmedBalance   string                   `json:"confirmedBalance"`
	UnconfirmedBalance string                   `json:"unconfirmedBalance"`
	Currency           *repo.CurrencyDefinition `json:"currency"`
	Memo               string                   `json:"memo"`
	OrderID            string                   `json:"orderId"`
	Timestamp          time.Time                `json:"timestamp"`
	Txid               string                   `json:"txid"`
	PeerID             string                   `json:"-"`
	ConsumedInput      bool                     `json:"-"`
}

// Spend will attempt to move funds from the node to the destination address described in the
// SpendRequest for the amount indicated.
func (n *OpenBazaarNode) Spend(args *SpendRequest) (*SpendResponse, error) {
	var (
		feeLevel wallet.FeeLevel
		peerID   string

		amt        = new(big.Int)
		lookupCode = args.CurrencyCode
	)

	if lookupCode == "" && args.Currency != nil {
		lookupCode = n.NormalizeCurrencyCode(args.Currency.Code.String())
	}
	var currencyDef, err = repo.LoadCurrencyDefinitions().Lookup(lookupCode)
	if err != nil {
		return nil, repo.ErrCurrencyDefinitionUndefined
	}
	if args.Currency != nil && currencyDef.Divisibility != args.Currency.Divisibility {
		currencyDef.Divisibility = args.Currency.Divisibility
		if err := currencyDef.Valid(); err != nil {
			return nil, err
		}
	}

	amt, ok := amt.SetString(args.Amount, 10)
	if !ok {
		return nil, ErrInvalidAmount
	}

	wal, err := n.Multiwallet.WalletForCurrencyCode(lookupCode)
	if err != nil {
		return nil, ErrUnknownWallet
	}

	addr, err := wal.DecodeAddress(args.Address)
	if err != nil {
		return nil, ErrInvalidSpendAddress
	}
	args.decodedAddress = addr

	contract, err := n.getOrderContractBySpendRequest(args)
	if err != nil && args.RequireAssociatedOrder {
		return nil, ErrOrderNotFound
	}

	switch strings.ToUpper(args.FeeLevel) {
	case "PRIORITY":
		feeLevel = wallet.PRIOIRTY
	case "NORMAL":
		feeLevel = wallet.NORMAL
	case "ECONOMIC":
		feeLevel = wallet.ECONOMIC
	default:
		feeLevel = wallet.NORMAL
	}

	txid, err := wal.Spend(*amt, addr, feeLevel, args.OrderID, args.SpendAll)
	if err != nil {
		switch {
		case err == wallet.ErrorInsuffientFunds:
			return nil, ErrInsufficientFunds
		case err == wallet.ErrorDustAmount:
			return nil, ErrSpendAmountIsDust
		default:
			return nil, err
		}
	}

	txn, err := wal.GetTransaction(*txid)
	if err != nil {
		return nil, fmt.Errorf("failed retrieving new wallet balance: %s", err)
	}

	var (
		thumbnail string
		title     string
		memo      = args.Memo
		toAddress = args.Address
	)

	if txn.ToAddress != "" {
		toAddress = txn.ToAddress
	}

	if contract != nil && contract.VendorListings[0] != nil {
		if contract.VendorListings[0].Item != nil && len(contract.VendorListings[0].Item.Images) > 0 {
			thumbnail = contract.VendorListings[0].Item.Images[0].Tiny
			title = contract.VendorListings[0].Item.Title
		}
		if contract.VendorListings[0].VendorID != nil {
			peerID = contract.VendorListings[0].VendorID.PeerID
		}
	}
	if memo == "" && title != "" {
		memo = title
	}

	if err := n.Datastore.TxMetadata().Put(repo.Metadata{
		Txid:       txid.String(),
		Address:    toAddress,
		Memo:       memo,
		OrderId:    args.OrderID,
		Thumbnail:  thumbnail,
		CanBumpFee: false,
	}); err != nil {
		return nil, fmt.Errorf("failed persisting transaction metadata: %s", err)
	}

	confirmed, unconfirmed := wal.Balance()
	defn, err := repo.LoadCurrencyDefinitions().Lookup(wal.CurrencyCode())
	if err != nil {
		return nil, fmt.Errorf("wallet currency not found in dictionary")
	}

	return &SpendResponse{
		Txid:               txid.String(),
		ConfirmedBalance:   confirmed.Value.String(),
		UnconfirmedBalance: unconfirmed.Value.String(),
		Currency:           defn,
		Amount:             strings.TrimPrefix(txn.Value, "-"),
		Timestamp:          txn.Timestamp,
		Memo:               memo,
		OrderID:            args.OrderID,
		PeerID:             peerID,
	}, nil
}

func (n *OpenBazaarNode) getOrderContractBySpendRequest(args *SpendRequest) (*pb.RicardianContract, error) {
	var errorStr = "unable to find order from order id or spend address"
	if args.OrderID != "" {
		contract, _, _, _, _, _, err := n.Datastore.Purchases().GetByOrderId(args.OrderID)
		if err != nil {
			return nil, fmt.Errorf("%s: %s", errorStr, err)
		}
		if contract != nil {
			return contract, nil
		}
	}

	if args.decodedAddress != nil {
		contract, _, _, _, err := n.Datastore.Purchases().GetByPaymentAddress(args.decodedAddress)
		if err != nil {
			return nil, fmt.Errorf("%s: %s", errorStr, err)
		}
		if contract != nil {
			return contract, nil
		}
	}

	return nil, errors.New(errorStr)
}
<|endoftext|>


<commit_before>// Config package provides an implementation of ConfigProvider. Currently it only reads from a JSON file, but
// this may be extended to support other forms.
package config

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
)

// Concrete storage for the configuration values read. The map keys are dot-delimited.
type Config map[string]interface{}

// Create a new, empty Config
func NewConfig() Config {
	return make(map[string]interface{})
}

// ReadFromFile is a helper that reads the configuration from JSON in the provided path with default
// options. It is equivalent to creating a new empty config, and calling AddFile with no prefix.
func ReadFromFile(path string) (Config, error) {
	result := NewConfig()

	e := result.AddFile(path, "", false)

	return result, e
}

// ReadFromEnv is a helper that reads the configuration from environment variables. It is equivalent
// to creating a new empty config, and calling AddEnvironment with no destination prefix
// (i.e. environment variables with the given prefix are added to config at the top level.)
func ReadFromEnv(prefix string) Config {
	result := NewConfig()

	result.AddEnvironment(prefix, "", false)

	return result
}

// AddFile will read a config file and merge configuration that it contains into the Config
// instance. destPrefix can be used to ensure all settings from this file appear within a
// specific namespace within the Config instance. It can be "", in which case the new config
// entries occur at the top level.
// If settings in this file already exist in config, override determines whether the new settings
// will override existing settings (yes if true, no if false.)
// The config file is JSON and will generally contain a single object. The properties of that
// object form the top-level name space for Get().
func (c Config) AddFile(path string, destPrefix string, override bool) error {
	// read file
	data, e := ioutil.ReadFile(path)
	if e != nil {
		return e
	}

	// decode json
	var v interface{}
	e = json.Unmarshal(data, &v)
	if e != nil {
		return e
	}

	// Get this as a map
	nested := v.(map[string]interface{})

	c.nestedMerge(nested, destPrefix, override)

	return nil

}

// AddEnvironment will add configuration properties from the environment. Only environment
// variables that start with sourcePrefix will be added; if sourcePrefix is "", all environment
// variables will be added. destPrefix can be used to ensure that all settings to be loaded from
// the environment appear within a specific namespace within the Config instance. That too can be ""
// in which case the new Config entries appear at the top level.
// If settings in this file already exist in config, override determines whether the new settings
// will override existing settings (yes if true, no if false.)
func (c Config) AddEnvironment(sourcePrefix string, destPrefix string, override bool) {
	envs := os.Environ()
	nested := make(map[string]interface{})

	fmt.Printf("AddEnvironment: %s\n", envs)
	for _, x := range envs {
		// split on =
		kv := strings.SplitN(x, "=", 2)
		fmt.Printf("kv=%s\n", kv)

		// if name starts with the prefix, add it to nested for merging
		if len(kv) >= 1 && strings.HasPrefix(kv[0], sourcePrefix) {
			// this is a candidate
			var v string
			if len(kv) == 2 {
				v = kv[1]
			} else {
				v = ""
			}
			nested[kv[0]] = v
		}
	}

	fmt.Printf("AddEnvironment: nested: %s\n", nested)
	// merge it into the config
	c.nestedMerge(nested, destPrefix, override)
}

func (c Config) nestedMerge(object map[string]interface{}, prefix string, override bool) {
	p := prefix + "."
	if p == "." {
		p = ""
	}

	for k, v := range object {
		if reflect.TypeOf(v).Kind() == reflect.Map {
			// if 'v' is a map of interface{}, recursively add.
			c.nestedMerge(v.(map[string]interface{}), p+k, override)
		} else {
			// otherwise just add the property, using the prefix. If the value exists, use
			// override.
			_, exists := c[p+k]
			if !exists || override {
				c[p+k] = v
			}
		}
	}
}

// Get looks up an object in the map via a key. The key can have "." separators for names;
// this will go into the structure as appropriate. It will return nil if a key maps to an undefined
// property, or where a partial key is not an object.
func (c Config) Get(key string) interface{} {
	return c[key]
}

func (c Config) HasKey(key string) bool {
	return c[key] != nil
}

// AsString returns a key from the configuration (using Get), but returning it as a string.
// If the key is not defined, it returns "".
func (c Config) AsString(key string) string {
	v := c.Get(key)
	if v == nil {
		return ""
	}
	return fmt.Sprintf("%v", v)
}

// AsInt returns a key from the configuration as an integer. Integer values in the
// json file are retrieved as float64. This will return an error if it's not a float64.
// Otherwise will convert it to an int (with truncation).
func (c Config) AsInt(key string) (int, error) {
	v := c.Get(key)
	vv, ok := v.(float64)
	if !ok {
		return 0, fmt.Errorf("Expected config property to be a numeric value, but wasn't: '%s'", v)
	}
	return int(vv), nil
}

func (c Config) AddDefault(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultOverride(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultString(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultStringOverride(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultInt(key string, value int) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultIntOverride(key string, value int) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,false)
	return true,nil
}

func (c Config) AddDefaultOverrideWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,true)
	return true,nil
}

func (c Config) AddDefaultStringWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,false)
	return true,nil
}

func (c Config) AddDefaultStringOverrideWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,true)
	return true,nil
}

func (c Config) AddDefaultIntWithSection(key string, value int, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,false)
	return true,nil
}

func (c Config) AddDefaultIntOverrideWithSection(key string, value int, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,true)
	return true,nil
}
<commit_msg>Update config.go<commit_after>// Config package provides an implementation of ConfigProvider. Currently it only reads from a JSON file, but
// this may be extended to support other forms.
package config

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
)

// Concrete storage for the configuration values read. The map keys are dot-delimited.
type Config map[string]interface{}

// Create a new, empty Config
func NewConfig() Config {
	return make(map[string]interface{})
}

// ReadFromFile is a helper that reads the configuration from JSON in the provided path with default
// options. It is equivalent to creating a new empty config, and calling AddFile with no prefix.
func ReadFromFile(path string) (Config, error) {
	result := NewConfig()

	e := result.AddFile(path, "", false)

	return result, e
}

// ReadFromEnv is a helper that reads the configuration from environment variables. It is equivalent
// to creating a new empty config, and calling AddEnvironment with no destination prefix
// (i.e. environment variables with the given prefix are added to config at the top level.)
func ReadFromEnv(prefix string) Config {
	result := NewConfig()

	result.AddEnvironment(prefix, "", false)

	return result
}

// AddFile will read a config file and merge configuration that it contains into the Config
// instance. destPrefix can be used to ensure all settings from this file appear within a
// specific namespace within the Config instance. It can be "", in which case the new config
// entries occur at the top level.
// If settings in this file already exist in config, override determines whether the new settings
// will override existing settings (yes if true, no if false.)
// The config file is JSON and will generally contain a single object. The properties of that
// object form the top-level name space for Get().
func (c Config) AddFile(path string, destPrefix string, override bool) error {
	// read file
	data, e := ioutil.ReadFile(path)
	if e != nil {
		return e
	}

	// decode json
	var v interface{}
	e = json.Unmarshal(data, &v)
	if e != nil {
		return e
	}

	// Get this as a map
	nested := v.(map[string]interface{})

	c.nestedMerge(nested, destPrefix, override)

	return nil

}

// AddEnvironment will add configuration properties from the environment. Only environment
// variables that start with sourcePrefix will be added; if sourcePrefix is "", all environment
// variables will be added. destPrefix can be used to ensure that all settings to be loaded from
// the environment appear within a specific namespace within the Config instance. That too can be ""
// in which case the new Config entries appear at the top level.
// If settings in this file already exist in config, override determines whether the new settings
// will override existing settings (yes if true, no if false.)
func (c Config) AddEnvironment(sourcePrefix string, destPrefix string, override bool) {
	envs := os.Environ()
	nested := make(map[string]interface{})

	fmt.Printf("AddEnvironment: %s\n", envs)
	for _, x := range envs {
		// split on =
		kv := strings.SplitN(x, "=", 2)
		fmt.Printf("kv=%s\n", kv)

		// if name starts with the prefix, add it to nested for merging
		if len(kv) >= 1 && strings.HasPrefix(kv[0], sourcePrefix) {
			// this is a candidate
			var v string
			if len(kv) == 2 {
				v = kv[1]
			} else {
				v = ""
			}
			nested[kv[0]] = v
		}
	}

	fmt.Printf("AddEnvironment: nested: %s\n", nested)
	// merge it into the config
	c.nestedMerge(nested, destPrefix, override)
}

func (c Config) nestedMerge(object map[string]interface{}, prefix string, override bool) {
	p := prefix + "."
	if p == "." {
		p = ""
	}

	for k, v := range object {
		if reflect.TypeOf(v).Kind() == reflect.Map {
			// if 'v' is a map of interface{}, recursively add.
			c.nestedMerge(v.(map[string]interface{}), p+k, override)
		} else {
			// otherwise just add the property, using the prefix. If the value exists, use
			// override.
			_, exists := c[p+k]
			if !exists || override {
				c[p+k] = v
			}
		}
	}
}

// Get looks up an object in the map via a key. The key can have "." separators for names;
// this will go into the structure as appropriate. It will return nil if a key maps to an undefined
// property, or where a partial key is not an object.
func (c Config) Get(key string) interface{} {
	return c[key]
}

func (c Config) HasKey(key string) bool {
	return c[key] != nil
}

// AsString returns a key from the configuration (using Get), but returning it as a string.
// If the key is not defined, it returns "".
func (c Config) GetString(key string) string {
	v := c.Get(key)
	if v == nil {
		return ""
	}
	return fmt.Sprintf("%v", v)
}

// AsInt returns a key from the configuration as an integer. Integer values in the
// json file are retrieved as float64. This will return an error if it's not a float64.
// Otherwise will convert it to an int (with truncation).
func (c Config) GetInt(key string) (int, error) {
	v := c.Get(key)
	vv, ok := v.(float64)
	if !ok {
		return 0, fmt.Errorf("Expected config property to be a numeric value, but wasn't: '%s'", v)
	}
	return int(vv), nil
}

func (c Config) GetBool(key string) (bool, error) {
	v := c.Get(key)
	vv, ok := v.(bool)
	if !ok {
		return 0, fmt.Errorf("Expected config property to be a numeric value, but wasn't: '%s'", v)
	}
	return bool(vv), nil
}

func (c Config) AddDefault(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultOverride(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultString(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultStringOverride(key string, value string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultBool(key string, value bool) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultBoolOverride(key string, value bool) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultInt(key string, value int) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",false)
	return true,nil
}

func (c Config) AddDefaultIntOverride(key string, value int) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,"",true)
	return true,nil
}

func (c Config) AddDefaultWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,false)
	return true,nil
}

func (c Config) AddDefaultOverrideWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,true)
	return true,nil
}

func (c Config) AddDefaultStringWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,false)
	return true,nil
}

func (c Config) AddDefaultStringOverrideWithSection(key string, value string, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,true)
	return true,nil
}

func (c Config) AddDefaultIntWithSection(key string, value int, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,false)
	return true,nil
}

func (c Config) AddDefaultIntOverrideWithSection(key string, value int, section string) (bool,error) {
	tmp := make(map[string]interface{})
	tmp[key] = value
	c.nestedMerge(tmp,section,true)
	return true,nil
}
<|endoftext|>


<commit_before>package storage

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/go/glog"
	"github.com/chrislusf/seaweedfs/go/operation"
	"github.com/chrislusf/seaweedfs/go/security"
	"github.com/chrislusf/seaweedfs/go/util"
	"github.com/golang/protobuf/proto"
)

const (
	MAX_TTL_VOLUME_REMOVAL_DELAY = 10 // 10 minutes
)

type DiskLocation struct {
	Directory      string
	MaxVolumeCount int
	volumes        map[VolumeId]*Volume
}

func (mn *DiskLocation) reset() {
}

type MasterNodes struct {
	nodes    []string
	lastNode int
}

func (mn *MasterNodes) String() string {
	return fmt.Sprintf("nodes:%v, lastNode:%d", mn.nodes, mn.lastNode)
}

func NewMasterNodes(bootstrapNode string) (mn *MasterNodes) {
	mn = &MasterNodes{nodes: []string{bootstrapNode}, lastNode: -1}
	return
}
func (mn *MasterNodes) reset() {
	if len(mn.nodes) > 1 && mn.lastNode > 0 {
		mn.lastNode = -mn.lastNode
	}
}
func (mn *MasterNodes) findMaster() (string, error) {
	if len(mn.nodes) == 0 {
		return "", errors.New("No master node found!")
	}
	if mn.lastNode < 0 {
		for _, m := range mn.nodes {
			glog.V(4).Infof("Listing masters on %s", m)
			if masters, e := operation.ListMasters(m); e == nil {
				if len(masters) == 0 {
					continue
				}
				mn.nodes = masters
				mn.lastNode = rand.Intn(len(mn.nodes))
				glog.V(2).Info("current master node is :", mn.nodes[mn.lastNode])
				break
			} else {
				glog.V(4).Infof("Failed listing masters on %s: %v", m, e)
			}
		}
	}
	if mn.lastNode < 0 {
		return "", errors.New("No master node available!")
	}
	return mn.nodes[mn.lastNode], nil
}

/*
 * A VolumeServer contains one Store
 */
type Store struct {
	Ip              string
	Port            int
	PublicUrl       string
	Locations       []*DiskLocation
	dataCenter      string //optional informaton, overwriting master setting if exists
	rack            string //optional information, overwriting master setting if exists
	connected       bool
	volumeSizeLimit uint64 //read from the master
	masterNodes     *MasterNodes
}

func (s *Store) String() (str string) {
	str = fmt.Sprintf("Ip:%s, Port:%d, PublicUrl:%s, dataCenter:%s, rack:%s, connected:%v, volumeSizeLimit:%d, masterNodes:%s", s.Ip, s.Port, s.PublicUrl, s.dataCenter, s.rack, s.connected, s.volumeSizeLimit, s.masterNodes)
	return
}

func NewStore(port int, ip, publicUrl string, dirnames []string, maxVolumeCounts []int, needleMapKind NeedleMapType) (s *Store) {
	s = &Store{Port: port, Ip: ip, PublicUrl: publicUrl}
	s.Locations = make([]*DiskLocation, 0)
	for i := 0; i < len(dirnames); i++ {
		location := &DiskLocation{Directory: dirnames[i], MaxVolumeCount: maxVolumeCounts[i]}
		location.volumes = make(map[VolumeId]*Volume)
		location.loadExistingVolumes(needleMapKind)
		s.Locations = append(s.Locations, location)
	}
	return
}
func (s *Store) AddVolume(volumeListString string, collection string, needleMapKind NeedleMapType, replicaPlacement string, ttlString string) error {
	rt, e := NewReplicaPlacementFromString(replicaPlacement)
	if e != nil {
		return e
	}
	ttl, e := ReadTTL(ttlString)
	if e != nil {
		return e
	}
	for _, range_string := range strings.Split(volumeListString, ",") {
		if strings.Index(range_string, "-") < 0 {
			id_string := range_string
			id, err := NewVolumeId(id_string)
			if err != nil {
				return fmt.Errorf("Volume Id %s is not a valid unsigned integer!", id_string)
			}
			e = s.addVolume(VolumeId(id), collection, needleMapKind, rt, ttl)
		} else {
			pair := strings.Split(range_string, "-")
			start, start_err := strconv.ParseUint(pair[0], 10, 64)
			if start_err != nil {
				return fmt.Errorf("Volume Start Id %s is not a valid unsigned integer!", pair[0])
			}
			end, end_err := strconv.ParseUint(pair[1], 10, 64)
			if end_err != nil {
				return fmt.Errorf("Volume End Id %s is not a valid unsigned integer!", pair[1])
			}
			for id := start; id <= end; id++ {
				if err := s.addVolume(VolumeId(id), collection, needleMapKind, rt, ttl); err != nil {
					e = err
				}
			}
		}
	}
	return e
}
func (s *Store) DeleteCollection(collection string) (e error) {
	for _, location := range s.Locations {
		for k, v := range location.volumes {
			if v.Collection == collection {
				e = v.Destroy()
				if e != nil {
					return
				}
				delete(location.volumes, k)
			}
		}
	}
	return
}
func (s *Store) DeleteVolume(volumes map[VolumeId]*Volume, v *Volume) (e error) {
	e = v.Destroy()
	if e != nil {
		return
	}
	delete(volumes, v.Id)
	return
}
func (s *Store) findVolume(vid VolumeId) *Volume {
	for _, location := range s.Locations {
		if v, found := location.volumes[vid]; found {
			return v
		}
	}
	return nil
}
func (s *Store) findFreeLocation() (ret *DiskLocation) {
	max := 0
	for _, location := range s.Locations {
		currentFreeCount := location.MaxVolumeCount - len(location.volumes)
		if currentFreeCount > max {
			max = currentFreeCount
			ret = location
		}
	}
	return ret
}
func (s *Store) addVolume(vid VolumeId, collection string, needleMapKind NeedleMapType, replicaPlacement *ReplicaPlacement, ttl *TTL) error {
	if s.findVolume(vid) != nil {
		return fmt.Errorf("Volume Id %d already exists!", vid)
	}
	if location := s.findFreeLocation(); location != nil {
		glog.V(0).Infof("In dir %s adds volume:%v collection:%s replicaPlacement:%v ttl:%v",
			location.Directory, vid, collection, replicaPlacement, ttl)
		if volume, err := NewVolume(location.Directory, collection, vid, needleMapKind, replicaPlacement, ttl); err == nil {
			location.volumes[vid] = volume
			return nil
		} else {
			return err
		}
	}
	return fmt.Errorf("No more free space left")
}

func (l *DiskLocation) loadExistingVolumes(needleMapKind NeedleMapType) {
	if dirs, err := ioutil.ReadDir(l.Directory); err == nil {
		for _, dir := range dirs {
			name := dir.Name()
			if !dir.IsDir() && strings.HasSuffix(name, ".dat") {
				collection := ""
				base := name[:len(name)-len(".dat")]
				i := strings.Index(base, "_")
				if i > 0 {
					collection, base = base[0:i], base[i+1:]
				}
				if vid, err := NewVolumeId(base); err == nil {
					if l.volumes[vid] == nil {
						if v, e := NewVolume(l.Directory, collection, vid, needleMapKind, nil, nil); e == nil {
							l.volumes[vid] = v
							glog.V(0).Infof("data file %s, replicaPlacement=%s v=%d size=%d ttl=%s", l.Directory+"/"+name, v.ReplicaPlacement, v.Version(), v.Size(), v.Ttl.String())
						}
					}
				}
			}
		}
	}
	glog.V(0).Infoln("Store started on dir:", l.Directory, "with", len(l.volumes), "volumes", "max", l.MaxVolumeCount)
}
func (s *Store) Status() []*VolumeInfo {
	var stats []*VolumeInfo
	for _, location := range s.Locations {
		for k, v := range location.volumes {
			s := &VolumeInfo{
				Id:               VolumeId(k),
				Size:             v.ContentSize(),
				Collection:       v.Collection,
				ReplicaPlacement: v.ReplicaPlacement,
				Version:          v.Version(),
				FileCount:        v.nm.FileCount(),
				DeleteCount:      v.nm.DeletedCount(),
				DeletedByteCount: v.nm.DeletedSize(),
				ReadOnly:         v.readOnly,
				Ttl:              v.Ttl}
			stats = append(stats, s)
		}
	}
	sortVolumeInfos(stats)
	return stats
}

func (s *Store) SetDataCenter(dataCenter string) {
	s.dataCenter = dataCenter
}
func (s *Store) SetRack(rack string) {
	s.rack = rack
}

func (s *Store) SetBootstrapMaster(bootstrapMaster string) {
	s.masterNodes = NewMasterNodes(bootstrapMaster)
}
func (s *Store) SendHeartbeatToMaster() (masterNode string, secretKey security.Secret, e error) {
	masterNode, e = s.masterNodes.findMaster()
	if e != nil {
		return
	}
	var volumeMessages []*operation.VolumeInformationMessage
	maxVolumeCount := 0
	var maxFileKey uint64
	for _, location := range s.Locations {
		maxVolumeCount = maxVolumeCount + location.MaxVolumeCount
		for k, v := range location.volumes {
			if maxFileKey < v.nm.MaxFileKey() {
				maxFileKey = v.nm.MaxFileKey()
			}
			if !v.expired(s.volumeSizeLimit) {
				volumeMessage := &operation.VolumeInformationMessage{
					Id:               proto.Uint32(uint32(k)),
					Size:             proto.Uint64(uint64(v.Size())),
					Collection:       proto.String(v.Collection),
					FileCount:        proto.Uint64(uint64(v.nm.FileCount())),
					DeleteCount:      proto.Uint64(uint64(v.nm.DeletedCount())),
					DeletedByteCount: proto.Uint64(v.nm.DeletedSize()),
					ReadOnly:         proto.Bool(v.readOnly),
					ReplicaPlacement: proto.Uint32(uint32(v.ReplicaPlacement.Byte())),
					Version:          proto.Uint32(uint32(v.Version())),
					Ttl:              proto.Uint32(v.Ttl.ToUint32()),
				}
				volumeMessages = append(volumeMessages, volumeMessage)
			} else {
				if v.exiredLongEnough(MAX_TTL_VOLUME_REMOVAL_DELAY) {
					s.DeleteVolume(location.volumes, v)
					glog.V(0).Infoln("volume", v.Id, "is deleted.")
				} else {
					glog.V(0).Infoln("volume", v.Id, "is expired.")
				}
			}
		}
	}

	joinMessage := &operation.JoinMessage{
		IsInit:         proto.Bool(!s.connected),
		Ip:             proto.String(s.Ip),
		Port:           proto.Uint32(uint32(s.Port)),
		PublicUrl:      proto.String(s.PublicUrl),
		MaxVolumeCount: proto.Uint32(uint32(maxVolumeCount)),
		MaxFileKey:     proto.Uint64(maxFileKey),
		DataCenter:     proto.String(s.dataCenter),
		Rack:           proto.String(s.rack),
		Volumes:        volumeMessages,
	}

	data, err := proto.Marshal(joinMessage)
	if err != nil {
		return "", "", err
	}

	joinUrl := "http://" + masterNode + "/dir/join"
	glog.V(4).Infof("Connecting to %s ...", joinUrl)

	jsonBlob, err := util.PostBytes(joinUrl, data)
	if err != nil {
		s.masterNodes.reset()
		return "", "", err
	}
	var ret operation.JoinResult
	if err := json.Unmarshal(jsonBlob, &ret); err != nil {
		glog.V(0).Infof("Failed to join %s with response: %s", joinUrl, string(jsonBlob))
		return masterNode, "", err
	}
	if ret.Error != "" {
		return masterNode, "", errors.New(ret.Error)
	}
	s.volumeSizeLimit = ret.VolumeSizeLimit
	secretKey = security.Secret(ret.SecretKey)
	s.connected = true
	return
}
func (s *Store) Close() {
	for _, location := range s.Locations {
		for _, v := range location.volumes {
			v.Close()
		}
	}
}
func (s *Store) Write(i VolumeId, n *Needle) (size uint32, err error) {
	if v := s.findVolume(i); v != nil {
		if v.readOnly {
			err = fmt.Errorf("Volume %d is read only", i)
			return
		}
		if MaxPossibleVolumeSize >= v.ContentSize()+uint64(size) {
			size, err = v.write(n)
		} else {
			err = fmt.Errorf("Volume Size Limit %d Exceeded! Current size is %d", s.volumeSizeLimit, v.ContentSize())
		}
		if s.volumeSizeLimit < v.ContentSize()+3*uint64(size) {
			glog.V(0).Infoln("volume", i, "size", v.ContentSize(), "will exceed limit", s.volumeSizeLimit)
			if _, _, e := s.SendHeartbeatToMaster(); e != nil {
				glog.V(0).Infoln("error when reporting size:", e)
			}
		}
		return
	}
	glog.V(0).Infoln("volume", i, "not found!")
	err = fmt.Errorf("Volume %d not found!", i)
	return
}
func (s *Store) Delete(i VolumeId, n *Needle) (uint32, error) {
	if v := s.findVolume(i); v != nil && !v.readOnly {
		return v.delete(n)
	}
	return 0, nil
}
func (s *Store) Read(i VolumeId, n *Needle) (int, error) {
	if v := s.findVolume(i); v != nil {
		return v.read(n)
	}
	return 0, fmt.Errorf("Volume %v not found!", i)
}
func (s *Store) GetVolume(i VolumeId) *Volume {
	return s.findVolume(i)
}

func (s *Store) HasVolume(i VolumeId) bool {
	v := s.findVolume(i)
	return v != nil
}
<commit_msg>Fix looking for master nodes.<commit_after>package storage

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/go/glog"
	"github.com/chrislusf/seaweedfs/go/operation"
	"github.com/chrislusf/seaweedfs/go/security"
	"github.com/chrislusf/seaweedfs/go/util"
	"github.com/golang/protobuf/proto"
)

const (
	MAX_TTL_VOLUME_REMOVAL_DELAY = 10 // 10 minutes
)

type DiskLocation struct {
	Directory      string
	MaxVolumeCount int
	volumes        map[VolumeId]*Volume
}

func (mn *DiskLocation) reset() {
}

type MasterNodes struct {
	nodes    []string
	lastNode int
}

func (mn *MasterNodes) String() string {
	return fmt.Sprintf("nodes:%v, lastNode:%d", mn.nodes, mn.lastNode)
}

func NewMasterNodes(bootstrapNode string) (mn *MasterNodes) {
	mn = &MasterNodes{nodes: []string{bootstrapNode}, lastNode: -1}
	return
}
func (mn *MasterNodes) reset() {
	glog.V(4).Infof("Resetting master nodes: %v", mn)
	if len(mn.nodes) > 1 && mn.lastNode > 0 {
		glog.V(0).Infof("Reset master %s from: %v", mn.nodes[mn.lastNode], mn.nodes)
		mn.lastNode = -mn.lastNode
	}
}
func (mn *MasterNodes) findMaster() (string, error) {
	if len(mn.nodes) == 0 {
		return "", errors.New("No master node found!")
	}
	if mn.lastNode < 0 {
		for _, m := range mn.nodes {
			glog.V(4).Infof("Listing masters on %s", m)
			if masters, e := operation.ListMasters(m); e == nil {
				if len(masters) == 0 {
					continue
				}
				mn.nodes = append(masters, m)
				mn.lastNode = rand.Intn(len(mn.nodes))
				glog.V(2).Infof("current master nodes is %v", mn)
				break
			} else {
				glog.V(4).Infof("Failed listing masters on %s: %v", m, e)
			}
		}
	}
	if mn.lastNode < 0 {
		return "", errors.New("No master node available!")
	}
	return mn.nodes[mn.lastNode], nil
}

/*
 * A VolumeServer contains one Store
 */
type Store struct {
	Ip              string
	Port            int
	PublicUrl       string
	Locations       []*DiskLocation
	dataCenter      string //optional informaton, overwriting master setting if exists
	rack            string //optional information, overwriting master setting if exists
	connected       bool
	volumeSizeLimit uint64 //read from the master
	masterNodes     *MasterNodes
}

func (s *Store) String() (str string) {
	str = fmt.Sprintf("Ip:%s, Port:%d, PublicUrl:%s, dataCenter:%s, rack:%s, connected:%v, volumeSizeLimit:%d, masterNodes:%s", s.Ip, s.Port, s.PublicUrl, s.dataCenter, s.rack, s.connected, s.volumeSizeLimit, s.masterNodes)
	return
}

func NewStore(port int, ip, publicUrl string, dirnames []string, maxVolumeCounts []int, needleMapKind NeedleMapType) (s *Store) {
	s = &Store{Port: port, Ip: ip, PublicUrl: publicUrl}
	s.Locations = make([]*DiskLocation, 0)
	for i := 0; i < len(dirnames); i++ {
		location := &DiskLocation{Directory: dirnames[i], MaxVolumeCount: maxVolumeCounts[i]}
		location.volumes = make(map[VolumeId]*Volume)
		location.loadExistingVolumes(needleMapKind)
		s.Locations = append(s.Locations, location)
	}
	return
}
func (s *Store) AddVolume(volumeListString string, collection string, needleMapKind NeedleMapType, replicaPlacement string, ttlString string) error {
	rt, e := NewReplicaPlacementFromString(replicaPlacement)
	if e != nil {
		return e
	}
	ttl, e := ReadTTL(ttlString)
	if e != nil {
		return e
	}
	for _, range_string := range strings.Split(volumeListString, ",") {
		if strings.Index(range_string, "-") < 0 {
			id_string := range_string
			id, err := NewVolumeId(id_string)
			if err != nil {
				return fmt.Errorf("Volume Id %s is not a valid unsigned integer!", id_string)
			}
			e = s.addVolume(VolumeId(id), collection, needleMapKind, rt, ttl)
		} else {
			pair := strings.Split(range_string, "-")
			start, start_err := strconv.ParseUint(pair[0], 10, 64)
			if start_err != nil {
				return fmt.Errorf("Volume Start Id %s is not a valid unsigned integer!", pair[0])
			}
			end, end_err := strconv.ParseUint(pair[1], 10, 64)
			if end_err != nil {
				return fmt.Errorf("Volume End Id %s is not a valid unsigned integer!", pair[1])
			}
			for id := start; id <= end; id++ {
				if err := s.addVolume(VolumeId(id), collection, needleMapKind, rt, ttl); err != nil {
					e = err
				}
			}
		}
	}
	return e
}
func (s *Store) DeleteCollection(collection string) (e error) {
	for _, location := range s.Locations {
		for k, v := range location.volumes {
			if v.Collection == collection {
				e = v.Destroy()
				if e != nil {
					return
				}
				delete(location.volumes, k)
			}
		}
	}
	return
}
func (s *Store) DeleteVolume(volumes map[VolumeId]*Volume, v *Volume) (e error) {
	e = v.Destroy()
	if e != nil {
		return
	}
	delete(volumes, v.Id)
	return
}
func (s *Store) findVolume(vid VolumeId) *Volume {
	for _, location := range s.Locations {
		if v, found := location.volumes[vid]; found {
			return v
		}
	}
	return nil
}
func (s *Store) findFreeLocation() (ret *DiskLocation) {
	max := 0
	for _, location := range s.Locations {
		currentFreeCount := location.MaxVolumeCount - len(location.volumes)
		if currentFreeCount > max {
			max = currentFreeCount
			ret = location
		}
	}
	return ret
}
func (s *Store) addVolume(vid VolumeId, collection string, needleMapKind NeedleMapType, replicaPlacement *ReplicaPlacement, ttl *TTL) error {
	if s.findVolume(vid) != nil {
		return fmt.Errorf("Volume Id %d already exists!", vid)
	}
	if location := s.findFreeLocation(); location != nil {
		glog.V(0).Infof("In dir %s adds volume:%v collection:%s replicaPlacement:%v ttl:%v",
			location.Directory, vid, collection, replicaPlacement, ttl)
		if volume, err := NewVolume(location.Directory, collection, vid, needleMapKind, replicaPlacement, ttl); err == nil {
			location.volumes[vid] = volume
			return nil
		} else {
			return err
		}
	}
	return fmt.Errorf("No more free space left")
}

func (l *DiskLocation) loadExistingVolumes(needleMapKind NeedleMapType) {
	if dirs, err := ioutil.ReadDir(l.Directory); err == nil {
		for _, dir := range dirs {
			name := dir.Name()
			if !dir.IsDir() && strings.HasSuffix(name, ".dat") {
				collection := ""
				base := name[:len(name)-len(".dat")]
				i := strings.Index(base, "_")
				if i > 0 {
					collection, base = base[0:i], base[i+1:]
				}
				if vid, err := NewVolumeId(base); err == nil {
					if l.volumes[vid] == nil {
						if v, e := NewVolume(l.Directory, collection, vid, needleMapKind, nil, nil); e == nil {
							l.volumes[vid] = v
							glog.V(0).Infof("data file %s, replicaPlacement=%s v=%d size=%d ttl=%s", l.Directory+"/"+name, v.ReplicaPlacement, v.Version(), v.Size(), v.Ttl.String())
						}
					}
				}
			}
		}
	}
	glog.V(0).Infoln("Store started on dir:", l.Directory, "with", len(l.volumes), "volumes", "max", l.MaxVolumeCount)
}
func (s *Store) Status() []*VolumeInfo {
	var stats []*VolumeInfo
	for _, location := range s.Locations {
		for k, v := range location.volumes {
			s := &VolumeInfo{
				Id:               VolumeId(k),
				Size:             v.ContentSize(),
				Collection:       v.Collection,
				ReplicaPlacement: v.ReplicaPlacement,
				Version:          v.Version(),
				FileCount:        v.nm.FileCount(),
				DeleteCount:      v.nm.DeletedCount(),
				DeletedByteCount: v.nm.DeletedSize(),
				ReadOnly:         v.readOnly,
				Ttl:              v.Ttl}
			stats = append(stats, s)
		}
	}
	sortVolumeInfos(stats)
	return stats
}

func (s *Store) SetDataCenter(dataCenter string) {
	s.dataCenter = dataCenter
}
func (s *Store) SetRack(rack string) {
	s.rack = rack
}

func (s *Store) SetBootstrapMaster(bootstrapMaster string) {
	s.masterNodes = NewMasterNodes(bootstrapMaster)
}
func (s *Store) SendHeartbeatToMaster() (masterNode string, secretKey security.Secret, e error) {
	masterNode, e = s.masterNodes.findMaster()
	if e != nil {
		return
	}
	var volumeMessages []*operation.VolumeInformationMessage
	maxVolumeCount := 0
	var maxFileKey uint64
	for _, location := range s.Locations {
		maxVolumeCount = maxVolumeCount + location.MaxVolumeCount
		for k, v := range location.volumes {
			if maxFileKey < v.nm.MaxFileKey() {
				maxFileKey = v.nm.MaxFileKey()
			}
			if !v.expired(s.volumeSizeLimit) {
				volumeMessage := &operation.VolumeInformationMessage{
					Id:               proto.Uint32(uint32(k)),
					Size:             proto.Uint64(uint64(v.Size())),
					Collection:       proto.String(v.Collection),
					FileCount:        proto.Uint64(uint64(v.nm.FileCount())),
					DeleteCount:      proto.Uint64(uint64(v.nm.DeletedCount())),
					DeletedByteCount: proto.Uint64(v.nm.DeletedSize()),
					ReadOnly:         proto.Bool(v.readOnly),
					ReplicaPlacement: proto.Uint32(uint32(v.ReplicaPlacement.Byte())),
					Version:          proto.Uint32(uint32(v.Version())),
					Ttl:              proto.Uint32(v.Ttl.ToUint32()),
				}
				volumeMessages = append(volumeMessages, volumeMessage)
			} else {
				if v.exiredLongEnough(MAX_TTL_VOLUME_REMOVAL_DELAY) {
					s.DeleteVolume(location.volumes, v)
					glog.V(0).Infoln("volume", v.Id, "is deleted.")
				} else {
					glog.V(0).Infoln("volume", v.Id, "is expired.")
				}
			}
		}
	}

	joinMessage := &operation.JoinMessage{
		IsInit:         proto.Bool(!s.connected),
		Ip:             proto.String(s.Ip),
		Port:           proto.Uint32(uint32(s.Port)),
		PublicUrl:      proto.String(s.PublicUrl),
		MaxVolumeCount: proto.Uint32(uint32(maxVolumeCount)),
		MaxFileKey:     proto.Uint64(maxFileKey),
		DataCenter:     proto.String(s.dataCenter),
		Rack:           proto.String(s.rack),
		Volumes:        volumeMessages,
	}

	data, err := proto.Marshal(joinMessage)
	if err != nil {
		return "", "", err
	}

	joinUrl := "http://" + masterNode + "/dir/join"
	glog.V(4).Infof("Connecting to %s ...", joinUrl)

	jsonBlob, err := util.PostBytes(joinUrl, data)
	if err != nil {
		s.masterNodes.reset()
		return "", "", err
	}
	var ret operation.JoinResult
	if err := json.Unmarshal(jsonBlob, &ret); err != nil {
		glog.V(0).Infof("Failed to join %s with response: %s", joinUrl, string(jsonBlob))
		s.masterNodes.reset()
		return masterNode, "", err
	}
	if ret.Error != "" {
		s.masterNodes.reset()
		return masterNode, "", errors.New(ret.Error)
	}
	s.volumeSizeLimit = ret.VolumeSizeLimit
	secretKey = security.Secret(ret.SecretKey)
	s.connected = true
	return
}
func (s *Store) Close() {
	for _, location := range s.Locations {
		for _, v := range location.volumes {
			v.Close()
		}
	}
}
func (s *Store) Write(i VolumeId, n *Needle) (size uint32, err error) {
	if v := s.findVolume(i); v != nil {
		if v.readOnly {
			err = fmt.Errorf("Volume %d is read only", i)
			return
		}
		if MaxPossibleVolumeSize >= v.ContentSize()+uint64(size) {
			size, err = v.write(n)
		} else {
			err = fmt.Errorf("Volume Size Limit %d Exceeded! Current size is %d", s.volumeSizeLimit, v.ContentSize())
		}
		if s.volumeSizeLimit < v.ContentSize()+3*uint64(size) {
			glog.V(0).Infoln("volume", i, "size", v.ContentSize(), "will exceed limit", s.volumeSizeLimit)
			if _, _, e := s.SendHeartbeatToMaster(); e != nil {
				glog.V(0).Infoln("error when reporting size:", e)
			}
		}
		return
	}
	glog.V(0).Infoln("volume", i, "not found!")
	err = fmt.Errorf("Volume %d not found!", i)
	return
}
func (s *Store) Delete(i VolumeId, n *Needle) (uint32, error) {
	if v := s.findVolume(i); v != nil && !v.readOnly {
		return v.delete(n)
	}
	return 0, nil
}
func (s *Store) Read(i VolumeId, n *Needle) (int, error) {
	if v := s.findVolume(i); v != nil {
		return v.read(n)
	}
	return 0, fmt.Errorf("Volume %v not found!", i)
}
func (s *Store) GetVolume(i VolumeId) *Volume {
	return s.findVolume(i)
}

func (s *Store) HasVolume(i VolumeId) bool {
	v := s.findVolume(i)
	return v != nil
}
<|endoftext|>


<commit_before>// ahocorasick.go: implementation of the Aho-Corasick string matching
// algorithm. Actually implemented as matching against []byte rather
// than the Go string type. Throughout this code []byte is referred to
// as a blice.
//
// http://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_string_matching_algorithm
//
// Copyright (c) 2013 CloudFlare, Inc.

package ahocorasick

import (
	"container/list"
)

// A node in the trie structure used to implement Aho-Corasick
type node struct {
	b []byte // The blice at this node

	output bool // True means this node represents a blice that should
	// be output when matching
	index int // index into original dictionary if output is true

	counter int // Set to the value of the Matcher.counter when a
	// match is output to prevent duplicates output

	// The use of fixed size arrays is space-inefficient but fast for
	// lookups.

	child [256]*node // A non-nil entry in this array means that the
	// index represents a byte value which can be
	// appended to the current node. Blices in the
	// trie are built up byte by byte through these
	// child node pointers.

	suffix *node // Pointer to the longest possible strict suffix of
	// this node

	fail *node // Pointer to the next node which is in the dictionary
	// which can be reached from here following suffixes. Called fail
	// because it is used to fallback in the trie when a match fails.
}

// Matcher is returned by NewMatcher and contains a list of blices to
// match against
type Matcher struct {
	counter int // Counts the number of matches done, and is used to
	// prevent output of multiple matches of the same string
	trie []node // preallocated block of memory containing all the
	// nodes
	extent int // offset into trie that is currently free
}

// finndBlice looks for a blice in the trie starting from the root and
// returns a pointer to the node representing the end of the blice. If
// the blice is not found it returns nil.
func (m *Matcher) findBlice(b []byte) *node {
	n := &m.trie[0]

	for n != nil && len(b) > 0 {
		n = n.child[int(b[0])]
		b = b[1:]
	}

	return n
}

// getFreeNode: gets a free node structure from the Matcher's trie
// pool and updates the extent to point to the next free node.
func (m *Matcher) getFreeNode() *node {
	m.extent += 1

	return &m.trie[m.extent-1]
}

// buildTrie builds the fundamental trie structure from a set of
// blices.
func (m *Matcher) buildTrie(dictionary [][]byte) {

	// Work out the maximum size for the trie (all dictionary entries
	// are distinct plus the root). This is used to preallocate memory
	// for it.

	max := 1
	for _, blice := range dictionary {
		max += len(blice)
	}
	m.trie = make([]node, max)

	// Calling this an ignoring its argument simply allocated
	// m.trie[0] which will be the root element

	m.getFreeNode()

	// This loop builds the nodes in the trie by following through
	// each dictionary entry building the children pointers.

	for i, blice := range dictionary {
		n := &m.trie[0]
		path := make([]byte, 0)
		for _, b := range blice {
			path = append(path, b)

			c := n.child[int(b)]

			if c == nil {
				c = m.getFreeNode()
				n.child[int(b)] = c
				c.b = make([]byte, len(path))
				copy(c.b, path)

				// Nodes directly under the root node will have the
				// root as their fail point as there are no suffixes
				// possible.

				if len(path) == 1 {
					c.fail = &m.trie[0]
				}

				c.suffix = &m.trie[0]
			}

			n = c
		}

		// The last value of n points to the node representing a
		// dictionary entry

		n.output = true
		n.index = i
	}

	l := new(list.List)
	l.PushBack(&m.trie[0])

	for l.Len() > 0 {
		n := l.Remove(l.Front()).(*node)

		for i := 0; i < 256; i++ {
			c := n.child[i]
			if c != nil {
				l.PushBack(c)

				for j := 1; j < len(c.b); j++ {
					c.fail = m.findBlice(c.b[j:])
					if c.fail != nil {
						break
					}
				}

				if c.fail == nil {
					c.fail = &m.trie[0]
				}

				for j := 1; j < len(c.b); j++ {
					s := m.findBlice(c.b[j:])
					if s != nil && s.output {
						c.suffix = s
						break
					}
				}
			}
		}
	}

	m.trie = m.trie[:m.extent]
}

// NewMatcher creates a new Matcher used to match against a set of
// blices
func NewMatcher(dictionary [][]byte) *Matcher {
	m := new(Matcher)

	m.buildTrie(dictionary)

	return m
}

// NewStringMatcher creates a new Matcher used to match against a set
// of strings (this is a helper to make initialization easy)
func NewStringMatcher(dictionary []string) *Matcher {
	m := new(Matcher)

	d := make([][]byte, 0)
	for _, s := range dictionary {
		d = append(d, []byte(s))
	}

	m.buildTrie(d)

	return m
}

// Match searches in for blices and returns all the blices found as
// indexes into the original dictionary
func (m *Matcher) Match(in []byte) []int {
	m.counter += 1
	hits := make([]int, 0)

	n := &m.trie[0]

	for _, b := range in {
		c := int(b)

		for n.child[c] == nil && n != &m.trie[0] {
			n = n.fail
		}

		if n.child[c] != nil {
			f := n.child[c]
			n = f

			if f.output && f.counter != m.counter {
				hits = append(hits, f.index)
				f.counter = m.counter
			}

			for f.suffix != &m.trie[0] {
				f = f.suffix
				if f.counter != m.counter {
					hits = append(hits, f.index)
					f.counter = m.counter
				} else {

					// There's no point working our way up the
					// suffixes if it's been done before for this call
					// to Match. The matches are already in hits.

					break
				}
			}
		}
	}

	return hits
}
<commit_msg>Allocate less in Match.<commit_after>// ahocorasick.go: implementation of the Aho-Corasick string matching
// algorithm. Actually implemented as matching against []byte rather
// than the Go string type. Throughout this code []byte is referred to
// as a blice.
//
// http://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_string_matching_algorithm
//
// Copyright (c) 2013 CloudFlare, Inc.

package ahocorasick

import (
	"container/list"
)

// A node in the trie structure used to implement Aho-Corasick
type node struct {
	b []byte // The blice at this node

	output bool // True means this node represents a blice that should
	// be output when matching
	index int // index into original dictionary if output is true

	counter int // Set to the value of the Matcher.counter when a
	// match is output to prevent duplicates output

	// The use of fixed size arrays is space-inefficient but fast for
	// lookups.

	child [256]*node // A non-nil entry in this array means that the
	// index represents a byte value which can be
	// appended to the current node. Blices in the
	// trie are built up byte by byte through these
	// child node pointers.

	suffix *node // Pointer to the longest possible strict suffix of
	// this node

	fail *node // Pointer to the next node which is in the dictionary
	// which can be reached from here following suffixes. Called fail
	// because it is used to fallback in the trie when a match fails.
}

// Matcher is returned by NewMatcher and contains a list of blices to
// match against
type Matcher struct {
	counter int // Counts the number of matches done, and is used to
	// prevent output of multiple matches of the same string
	trie []node // preallocated block of memory containing all the
	// nodes
	extent int // offset into trie that is currently free
}

// finndBlice looks for a blice in the trie starting from the root and
// returns a pointer to the node representing the end of the blice. If
// the blice is not found it returns nil.
func (m *Matcher) findBlice(b []byte) *node {
	n := &m.trie[0]

	for n != nil && len(b) > 0 {
		n = n.child[int(b[0])]
		b = b[1:]
	}

	return n
}

// getFreeNode: gets a free node structure from the Matcher's trie
// pool and updates the extent to point to the next free node.
func (m *Matcher) getFreeNode() *node {
	m.extent += 1

	return &m.trie[m.extent-1]
}

// buildTrie builds the fundamental trie structure from a set of
// blices.
func (m *Matcher) buildTrie(dictionary [][]byte) {

	// Work out the maximum size for the trie (all dictionary entries
	// are distinct plus the root). This is used to preallocate memory
	// for it.

	max := 1
	for _, blice := range dictionary {
		max += len(blice)
	}
	m.trie = make([]node, max)

	// Calling this an ignoring its argument simply allocated
	// m.trie[0] which will be the root element

	m.getFreeNode()

	// This loop builds the nodes in the trie by following through
	// each dictionary entry building the children pointers.

	for i, blice := range dictionary {
		n := &m.trie[0]
		path := make([]byte, 0)
		for _, b := range blice {
			path = append(path, b)

			c := n.child[int(b)]

			if c == nil {
				c = m.getFreeNode()
				n.child[int(b)] = c
				c.b = make([]byte, len(path))
				copy(c.b, path)

				// Nodes directly under the root node will have the
				// root as their fail point as there are no suffixes
				// possible.

				if len(path) == 1 {
					c.fail = &m.trie[0]
				}

				c.suffix = &m.trie[0]
			}

			n = c
		}

		// The last value of n points to the node representing a
		// dictionary entry

		n.output = true
		n.index = i
	}

	l := new(list.List)
	l.PushBack(&m.trie[0])

	for l.Len() > 0 {
		n := l.Remove(l.Front()).(*node)

		for i := 0; i < 256; i++ {
			c := n.child[i]
			if c != nil {
				l.PushBack(c)

				for j := 1; j < len(c.b); j++ {
					c.fail = m.findBlice(c.b[j:])
					if c.fail != nil {
						break
					}
				}

				if c.fail == nil {
					c.fail = &m.trie[0]
				}

				for j := 1; j < len(c.b); j++ {
					s := m.findBlice(c.b[j:])
					if s != nil && s.output {
						c.suffix = s
						break
					}
				}
			}
		}
	}

	m.trie = m.trie[:m.extent]
}

// NewMatcher creates a new Matcher used to match against a set of
// blices
func NewMatcher(dictionary [][]byte) *Matcher {
	m := new(Matcher)

	m.buildTrie(dictionary)

	return m
}

// NewStringMatcher creates a new Matcher used to match against a set
// of strings (this is a helper to make initialization easy)
func NewStringMatcher(dictionary []string) *Matcher {
	m := new(Matcher)

	d := make([][]byte, 0)
	for _, s := range dictionary {
		d = append(d, []byte(s))
	}

	m.buildTrie(d)

	return m
}

// Match searches in for blices and returns all the blices found as
// indexes into the original dictionary
func (m *Matcher) Match(in []byte) []int {
	m.counter += 1
	var hits []int

	n := &m.trie[0]

	for _, b := range in {
		c := int(b)

		for n.child[c] == nil && n != &m.trie[0] {
			n = n.fail
		}

		if n.child[c] != nil {
			f := n.child[c]
			n = f

			if f.output && f.counter != m.counter {
				hits = append(hits, f.index)
				f.counter = m.counter
			}

			for f.suffix != &m.trie[0] {
				f = f.suffix
				if f.counter != m.counter {
					hits = append(hits, f.index)
					f.counter = m.counter
				} else {

					// There's no point working our way up the
					// suffixes if it's been done before for this call
					// to Match. The matches are already in hits.

					break
				}
			}
		}
	}

	return hits
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 Google Inc.
//
// 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 adreno

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/gapis/perfetto"
	perfetto_service "github.com/google/gapid/gapis/perfetto/service"
	"github.com/google/gapid/gapis/service"
)

var (
	slicesQuery = "" +
		"SELECT s.context_id, s.render_target, s.frame_id, s.submission_id, s.hw_queue_id, s.command_buffer, s.ts, s.dur, s.name, depth, arg_set_id, track_id, t.name " +
		"FROM gpu_track t LEFT JOIN gpu_slice s " +
		"ON s.track_id = t.id WHERE t.scope = 'gpu_render_stage'"
	argsQueryFmt = "" +
		"SELECT key, string_value FROM args WHERE args.arg_set_id = %d"
	counterTracksQuery = "" +
		"SELECT id, name, unit, description FROM gpu_counter_track ORDER BY id"
	countersQueryFmt = "" +
		"SELECT ts, value FROM counter c WHERE c.track_id = %d ORDER BY ts"
)

func ProcessProfilingData(ctx context.Context, processor *perfetto.Processor, desc *device.GpuCounterDescriptor, handleMapping *map[uint64][]service.VulkanHandleMappingItem) (*service.ProfilingData, error) {
	slices, err := processGpuSlices(ctx, processor, handleMapping)
	if err != nil {
		log.Err(ctx, err, "Failed to get GPU slices")
	}
	counters, err := processCounters(ctx, processor, desc)
	if err != nil {
		log.Err(ctx, err, "Failed to get GPU counters")
	}
	return &service.ProfilingData{Slices: slices, Counters: counters}, nil
}

func processGpuSlices(ctx context.Context, processor *perfetto.Processor, handleMapping *map[uint64][]service.VulkanHandleMappingItem) (*service.ProfilingData_GpuSlices, error) {
	slicesQueryResult, err := processor.Query(slicesQuery)
	if err != nil {
		return nil, log.Errf(ctx, err, "SQL query failed: %v", slicesQuery)
	}

	trackIdCache := make(map[int64]bool)
	argsQueryCache := make(map[int64]*perfetto_service.QueryResult)
	slicesColumns := slicesQueryResult.GetColumns()
	numSliceRows := slicesQueryResult.GetNumRecords()
	slices := make([]*service.ProfilingData_GpuSlices_Slice, numSliceRows)
	var tracks []*service.ProfilingData_GpuSlices_Track
	// Grab all the column values. Depends on the order of columns selected in slicesQuery

	contextIds := slicesColumns[0].GetLongValues()
	for i, v := range contextIds {
		if m, ok := (*handleMapping)[uint64(v)]; ok {
			contextIds[i] = int64(m[0].TraceValue)
		} else {
			log.E(ctx, "Context Id not found: %v", v)
		}
	}

	renderTargets := slicesColumns[1].GetLongValues()
	for i, v := range renderTargets {
		if m, ok := (*handleMapping)[uint64(v)]; ok {
			renderTargets[i] = int64(m[0].TraceValue)
		} else {
			log.E(ctx, "Render Target not found: %v", v)
		}
	}

	commandBuffers := slicesColumns[5].GetLongValues()
	for i, v := range commandBuffers {
		if m, ok := (*handleMapping)[uint64(v)]; ok {
			commandBuffers[i] = int64(m[0].TraceValue)
		} else {
			log.E(ctx, "Command Buffer not found: %v", v)
		}
	}

	frameIds := slicesColumns[2].GetLongValues()
	submissionIds := slicesColumns[3].GetLongValues()
	hwQueueIds := slicesColumns[4].GetLongValues()
	timestamps := slicesColumns[6].GetLongValues()
	durations := slicesColumns[7].GetLongValues()
	names := slicesColumns[8].GetStringValues()
	depths := slicesColumns[9].GetLongValues()
	argSetIds := slicesColumns[10].GetLongValues()
	trackIds := slicesColumns[11].GetLongValues()
	trackNames := slicesColumns[12].GetStringValues()

	for i := uint64(0); i < numSliceRows; i++ {
		var argsQueryResult *perfetto_service.QueryResult
		var ok bool
		if argsQueryResult, ok = argsQueryCache[argSetIds[i]]; !ok {
			argsQuery := fmt.Sprintf(argsQueryFmt, argSetIds[i])
			argsQueryResult, err = processor.Query(argsQuery)
			if err != nil {
				log.W(ctx, "SQL query failed: %v", argsQuery)
			}
			argsQueryCache[argSetIds[i]] = argsQueryResult
		}
		argsColumns := argsQueryResult.GetColumns()
		numArgsRows := argsQueryResult.GetNumRecords()
		var extras []*service.ProfilingData_GpuSlices_Slice_Extra
		for j := uint64(0); j < numArgsRows; j++ {
			keys := argsColumns[0].GetStringValues()
			values := argsColumns[1].GetStringValues()
			extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
				Name:  keys[j],
				Value: &service.ProfilingData_GpuSlices_Slice_Extra_StringValue{StringValue: values[j]},
			})
		}
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "contextId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(contextIds[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "renderTarget",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(renderTargets[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "commandBuffer",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(commandBuffers[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "frameId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(frameIds[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "submissionId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(submissionIds[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "hwQueueId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(hwQueueIds[i])},
		})

		slices[i] = &service.ProfilingData_GpuSlices_Slice{
			Ts:      uint64(timestamps[i]),
			Dur:     uint64(durations[i]),
			Label:   names[i],
			Depth:   int32(depths[i]),
			Extras:  extras,
			TrackId: int32(trackIds[i]),
		}

		if _, ok := trackIdCache[trackIds[i]]; !ok {
			trackIdCache[trackIds[i]] = true
			tracks = append(tracks, &service.ProfilingData_GpuSlices_Track{
				Id:   int32(trackIds[i]),
				Name: trackNames[i],
			})
		}
	}

	return &service.ProfilingData_GpuSlices{
		Slices: slices,
		Tracks: tracks,
	}, nil
}

func processCounters(ctx context.Context, processor *perfetto.Processor, desc *device.GpuCounterDescriptor) ([]*service.ProfilingData_Counter, error) {
	counterTracksQueryResult, err := processor.Query(counterTracksQuery)
	if err != nil {
		return nil, log.Errf(ctx, err, "SQL query failed: %v", counterTracksQuery)
	}
	// t.id, name, unit, description, ts, value
	tracksColumns := counterTracksQueryResult.GetColumns()
	numTracksRows := counterTracksQueryResult.GetNumRecords()
	counters := make([]*service.ProfilingData_Counter, numTracksRows)
	// Grab all the column values. Depends on the order of columns selected in countersQuery
	trackIds := tracksColumns[0].GetLongValues()
	names := tracksColumns[1].GetStringValues()
	units := tracksColumns[2].GetStringValues()
	descriptions := tracksColumns[3].GetStringValues()

	for i := uint64(0); i < numTracksRows; i++ {
		countersQuery := fmt.Sprintf(countersQueryFmt, trackIds[i])
		countersQueryResult, err := processor.Query(countersQuery)
		countersColumns := countersQueryResult.GetColumns()
		if err != nil {
			return nil, log.Errf(ctx, err, "SQL query failed: %v", counterTracksQuery)
		}
		timestampsLong := countersColumns[0].GetLongValues()
		timestamps := make([]uint64, len(timestampsLong))
		for i, t := range timestampsLong {
			timestamps[i] = uint64(t)
		}
		values := countersColumns[1].GetDoubleValues()
		// TODO(apbodnar) Populate the `default` field once the trace processor supports it (b/147432390)
		counters[i] = &service.ProfilingData_Counter{
			Id:          uint32(trackIds[i]),
			Name:        names[i],
			Unit:        units[i],
			Description: descriptions[i],
			Timestamps:  timestamps,
			Values:      values,
		}
	}
	return counters, nil
}
<commit_msg>Consider if handle types match when replacing the replay handles (#72)<commit_after>// Copyright (C) 2019 Google Inc.
//
// 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 adreno

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/gapis/perfetto"
	perfetto_service "github.com/google/gapid/gapis/perfetto/service"
	"github.com/google/gapid/gapis/service"
)

var (
	slicesQuery = "" +
		"SELECT s.context_id, s.render_target, s.frame_id, s.submission_id, s.hw_queue_id, s.command_buffer, s.ts, s.dur, s.name, depth, arg_set_id, track_id, t.name " +
		"FROM gpu_track t LEFT JOIN gpu_slice s " +
		"ON s.track_id = t.id WHERE t.scope = 'gpu_render_stage'"
	argsQueryFmt = "" +
		"SELECT key, string_value FROM args WHERE args.arg_set_id = %d"
	counterTracksQuery = "" +
		"SELECT id, name, unit, description FROM gpu_counter_track ORDER BY id"
	countersQueryFmt = "" +
		"SELECT ts, value FROM counter c WHERE c.track_id = %d ORDER BY ts"
)

func ProcessProfilingData(ctx context.Context, processor *perfetto.Processor, desc *device.GpuCounterDescriptor, handleMapping *map[uint64][]service.VulkanHandleMappingItem) (*service.ProfilingData, error) {
	slices, err := processGpuSlices(ctx, processor, handleMapping)
	if err != nil {
		log.Err(ctx, err, "Failed to get GPU slices")
	}
	counters, err := processCounters(ctx, processor, desc)
	if err != nil {
		log.Err(ctx, err, "Failed to get GPU counters")
	}
	return &service.ProfilingData{Slices: slices, Counters: counters}, nil
}

func extractTraceHandles(ctx context.Context, replayHandles *[]int64, replayHandleType string, handleMapping *map[uint64][]service.VulkanHandleMappingItem) {
	for i, v := range *replayHandles {
		handles, ok := (*handleMapping)[uint64(v)]
		if !ok {
			log.E(ctx, "%v not found in replay: %v", replayHandleType, v)
			continue
		}

		found := false
		for _, handle := range handles {
			if handle.HandleType == replayHandleType {
				(*replayHandles)[i] = int64(handle.TraceValue)
				found = true
				break
			}
		}

		if !found {
			log.E(ctx, "Incorrect Handle type for %v: %v", replayHandleType, v)
		}
	}
}

func processGpuSlices(ctx context.Context, processor *perfetto.Processor, capture *path.Capture, handleMapping *map[uint64][]service.VulkanHandleMappingItem, subCommandIndicesMap *map[api.CommandSubmissionKey][]uint64) (*service.ProfilingData_GpuSlices, error) {
	slicesQueryResult, err := processor.Query(slicesQuery)
	if err != nil {
		return nil, log.Errf(ctx, err, "SQL query failed: %v", slicesQuery)
	}

	trackIdCache := make(map[int64]bool)
	argsQueryCache := make(map[int64]*perfetto_service.QueryResult)
	slicesColumns := slicesQueryResult.GetColumns()
	numSliceRows := slicesQueryResult.GetNumRecords()
	slices := make([]*service.ProfilingData_GpuSlices_Slice, numSliceRows)
	var tracks []*service.ProfilingData_GpuSlices_Track
	// Grab all the column values. Depends on the order of columns selected in slicesQuery

	contextIds := slicesColumns[0].GetLongValues()
	extractTraceHandles(ctx, &contextIds, "VkDevice", handleMapping)

	renderTargets := slicesColumns[1].GetLongValues()
	extractTraceHandles(ctx, &renderTargets, "VkFramebuffer", handleMapping)

	commandBuffers := slicesColumns[5].GetLongValues()
	extractTraceHandles(ctx, &commandBuffers, "VkCommandBuffer", handleMapping)

	frameIds := slicesColumns[2].GetLongValues()
	submissionIds := slicesColumns[3].GetLongValues()
	hwQueueIds := slicesColumns[4].GetLongValues()
	timestamps := slicesColumns[6].GetLongValues()
	durations := slicesColumns[7].GetLongValues()
	names := slicesColumns[8].GetStringValues()
	depths := slicesColumns[9].GetLongValues()
	argSetIds := slicesColumns[10].GetLongValues()
	trackIds := slicesColumns[11].GetLongValues()
	trackNames := slicesColumns[12].GetStringValues()

	for i := uint64(0); i < numSliceRows; i++ {
		var argsQueryResult *perfetto_service.QueryResult
		var ok bool
		if argsQueryResult, ok = argsQueryCache[argSetIds[i]]; !ok {
			argsQuery := fmt.Sprintf(argsQueryFmt, argSetIds[i])
			argsQueryResult, err = processor.Query(argsQuery)
			if err != nil {
				log.W(ctx, "SQL query failed: %v", argsQuery)
			}
			argsQueryCache[argSetIds[i]] = argsQueryResult
		}
		argsColumns := argsQueryResult.GetColumns()
		numArgsRows := argsQueryResult.GetNumRecords()
		var extras []*service.ProfilingData_GpuSlices_Slice_Extra
		for j := uint64(0); j < numArgsRows; j++ {
			keys := argsColumns[0].GetStringValues()
			values := argsColumns[1].GetStringValues()
			extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
				Name:  keys[j],
				Value: &service.ProfilingData_GpuSlices_Slice_Extra_StringValue{StringValue: values[j]},
			})
		}
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "contextId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(contextIds[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "renderTarget",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(renderTargets[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "commandBuffer",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(commandBuffers[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "frameId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(frameIds[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "submissionId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(submissionIds[i])},
		})
		extras = append(extras, &service.ProfilingData_GpuSlices_Slice_Extra{
			Name:  "hwQueueId",
			Value: &service.ProfilingData_GpuSlices_Slice_Extra_IntValue{IntValue: uint64(hwQueueIds[i])},
		})

		slices[i] = &service.ProfilingData_GpuSlices_Slice{
			Ts:      uint64(timestamps[i]),
			Dur:     uint64(durations[i]),
			Label:   names[i],
			Depth:   int32(depths[i]),
			Extras:  extras,
			TrackId: int32(trackIds[i]),
		}

		if _, ok := trackIdCache[trackIds[i]]; !ok {
			trackIdCache[trackIds[i]] = true
			tracks = append(tracks, &service.ProfilingData_GpuSlices_Track{
				Id:   int32(trackIds[i]),
				Name: trackNames[i],
			})
		}
	}

	return &service.ProfilingData_GpuSlices{
		Slices: slices,
		Tracks: tracks,
	}, nil
}

func processCounters(ctx context.Context, processor *perfetto.Processor, desc *device.GpuCounterDescriptor) ([]*service.ProfilingData_Counter, error) {
	counterTracksQueryResult, err := processor.Query(counterTracksQuery)
	if err != nil {
		return nil, log.Errf(ctx, err, "SQL query failed: %v", counterTracksQuery)
	}
	// t.id, name, unit, description, ts, value
	tracksColumns := counterTracksQueryResult.GetColumns()
	numTracksRows := counterTracksQueryResult.GetNumRecords()
	counters := make([]*service.ProfilingData_Counter, numTracksRows)
	// Grab all the column values. Depends on the order of columns selected in countersQuery
	trackIds := tracksColumns[0].GetLongValues()
	names := tracksColumns[1].GetStringValues()
	units := tracksColumns[2].GetStringValues()
	descriptions := tracksColumns[3].GetStringValues()

	for i := uint64(0); i < numTracksRows; i++ {
		countersQuery := fmt.Sprintf(countersQueryFmt, trackIds[i])
		countersQueryResult, err := processor.Query(countersQuery)
		countersColumns := countersQueryResult.GetColumns()
		if err != nil {
			return nil, log.Errf(ctx, err, "SQL query failed: %v", counterTracksQuery)
		}
		timestampsLong := countersColumns[0].GetLongValues()
		timestamps := make([]uint64, len(timestampsLong))
		for i, t := range timestampsLong {
			timestamps[i] = uint64(t)
		}
		values := countersColumns[1].GetDoubleValues()
		// TODO(apbodnar) Populate the `default` field once the trace processor supports it (b/147432390)
		counters[i] = &service.ProfilingData_Counter{
			Id:          uint32(trackIds[i]),
			Name:        names[i],
			Unit:        units[i],
			Description: descriptions[i],
			Timestamps:  timestamps,
			Values:      values,
		}
	}
	return counters, nil
}
<|endoftext|>


<commit_before>// Package traceapp implements the Appdash web UI.
//
// The web UI can be effectively launched using the appdash command (see
// cmd/appdash) or via embedding this package within your app.
//
// Templates and other resources needed by this package to render the UI are
// built into the program using go-bindata, so you still get to have single
// binary deployment.
//
// For an example of embedding the Appdash web UI within your own application
// via the traceapp package, see the examples/cmd/webapp example.
package traceapp

import (
	"encoding/json"
	htmpl "html/template"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"sort"
	"sync"

	"github.com/elazarl/go-bindata-assetfs"
	"github.com/gorilla/mux"

	"sourcegraph.com/sourcegraph/appdash"
	static "sourcegraph.com/sourcegraph/appdash-data"
)

// App is an HTTP application handler that also exposes methods for
// constructing URL routes.
type App struct {
	*Router

	Store   appdash.Store
	Queryer appdash.Queryer

	tmplLock sync.Mutex
	tmpls    map[string]*htmpl.Template
}

// New creates a new application handler. If r is nil, a new router is
// created.
func New(r *Router) *App {
	if r == nil {
		r = NewRouter(nil)
	}

	app := &App{Router: r}

	r.r.Get(RootRoute).Handler(handlerFunc(app.serveRoot))
	r.r.Get(TraceRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceSpanRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceProfileRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceSpanProfileRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceUploadRoute).Handler(handlerFunc(app.serveTraceUpload))
	r.r.Get(TracesRoute).Handler(handlerFunc(app.serveTraces))

	// Static file serving.
	r.r.Get(StaticRoute).Handler(http.StripPrefix("/static/", http.FileServer(&assetfs.AssetFS{
		Asset:    static.Asset,
		AssetDir: static.AssetDir,
		Prefix:   "",
	})))
	return app
}

// ServeHTTP implements http.Handler.
func (a *App) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	a.Router.r.ServeHTTP(w, r)
}

func (a *App) serveRoot(w http.ResponseWriter, r *http.Request) error {
	return a.renderTemplate(w, r, "root.html", http.StatusOK, &struct {
		TemplateCommon
	}{})
}

func (a *App) serveTrace(w http.ResponseWriter, r *http.Request) error {
	v := mux.Vars(r)

	traceID, err := appdash.ParseID(v["Trace"])
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return err
	}

	trace, err := a.Store.Trace(traceID)
	if err != nil {
		if err == appdash.ErrTraceNotFound {
			http.Error(w, err.Error(), http.StatusNotFound)
			return nil
		}
		return err
	}

	// Get sub-span if the Span route var is present.
	if spanIDStr := v["Span"]; spanIDStr != "" {
		spanID, err := appdash.ParseID(spanIDStr)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return err
		}
		trace = trace.FindSpan(spanID)
		if trace == nil {
			http.Error(w, err.Error(), http.StatusNotFound)
			return nil
		}
	}

	// We could use a separate handler for this, but as we need the above to
	// determine the correct trace (or therein sub-trace), we just handle any
	// JSON profile requests here.
	if path.Base(r.URL.Path) == "profile" {
		return a.profile(trace, w)
	}

	visData, err := a.d3timeline(trace)
	if err != nil {
		return err
	}

	// Determine the profile URL.
	var profile *url.URL
	if trace.ID.Parent == 0 {
		profile, err = a.Router.URLToTraceProfile(trace.Span.ID.Trace)
	} else {
		profile, err = a.Router.URLToTraceSpanProfile(trace.Span.ID.Trace, trace.Span.ID.Span)
	}
	if err != nil {
		return err
	}

	return a.renderTemplate(w, r, "trace.html", http.StatusOK, &struct {
		TemplateCommon
		Trace      *appdash.Trace
		VisData    []timelineItem
		ProfileURL string
	}{
		Trace:      trace,
		VisData:    visData,
		ProfileURL: profile.String(),
	})
}

func (a *App) serveTraces(w http.ResponseWriter, r *http.Request) error {
	traces, err := a.Queryer.Traces()
	if err != nil {
		return err
	}

	// Sort the traces by ID to ensure that the display order doesn't change upon
	// multiple page reloads if Queryer.Traces is e.g. backed by a map (which has
	// a random iteration order).
	sort.Sort(tracesByID(traces))

	return a.renderTemplate(w, r, "traces.html", http.StatusOK, &struct {
		TemplateCommon
		Traces []*appdash.Trace
	}{
		Traces: traces,
	})
}

func (a *App) serveTraceUpload(w http.ResponseWriter, r *http.Request) error {
	// Read the uploaded JSON trace data.
	defer r.Body.Close()
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	// Unmarshal the trace.
	var traces []*appdash.Trace
	err = json.Unmarshal(data, &traces)
	if err != nil {
		return err
	}

	// Collect the unmarshaled traces, ignoring any previously existing ones (i.e.
	// ones that would collide / be merged together).
	for _, trace := range traces {
		_, err = a.Store.Trace(trace.Span.ID.Trace)
		if err != appdash.ErrTraceNotFound {
			// The trace collides with an existing trace, ignore it.
			continue
		}

		// Collect the trace (store it for later viewing).
		if err = collectTrace(a.Store, trace); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>traceapp: let errors in HTTP handlers bubble up properly.<commit_after>// Package traceapp implements the Appdash web UI.
//
// The web UI can be effectively launched using the appdash command (see
// cmd/appdash) or via embedding this package within your app.
//
// Templates and other resources needed by this package to render the UI are
// built into the program using go-bindata, so you still get to have single
// binary deployment.
//
// For an example of embedding the Appdash web UI within your own application
// via the traceapp package, see the examples/cmd/webapp example.
package traceapp

import (
	"encoding/json"
	"errors"
	htmpl "html/template"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"sort"
	"sync"

	"github.com/elazarl/go-bindata-assetfs"
	"github.com/gorilla/mux"

	"sourcegraph.com/sourcegraph/appdash"
	static "sourcegraph.com/sourcegraph/appdash-data"
)

// App is an HTTP application handler that also exposes methods for
// constructing URL routes.
type App struct {
	*Router

	Store   appdash.Store
	Queryer appdash.Queryer

	tmplLock sync.Mutex
	tmpls    map[string]*htmpl.Template
}

// New creates a new application handler. If r is nil, a new router is
// created.
func New(r *Router) *App {
	if r == nil {
		r = NewRouter(nil)
	}

	app := &App{Router: r}

	r.r.Get(RootRoute).Handler(handlerFunc(app.serveRoot))
	r.r.Get(TraceRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceSpanRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceProfileRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceSpanProfileRoute).Handler(handlerFunc(app.serveTrace))
	r.r.Get(TraceUploadRoute).Handler(handlerFunc(app.serveTraceUpload))
	r.r.Get(TracesRoute).Handler(handlerFunc(app.serveTraces))

	// Static file serving.
	r.r.Get(StaticRoute).Handler(http.StripPrefix("/static/", http.FileServer(&assetfs.AssetFS{
		Asset:    static.Asset,
		AssetDir: static.AssetDir,
		Prefix:   "",
	})))
	return app
}

// ServeHTTP implements http.Handler.
func (a *App) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	a.Router.r.ServeHTTP(w, r)
}

func (a *App) serveRoot(w http.ResponseWriter, r *http.Request) error {
	return a.renderTemplate(w, r, "root.html", http.StatusOK, &struct {
		TemplateCommon
	}{})
}

func (a *App) serveTrace(w http.ResponseWriter, r *http.Request) error {
	v := mux.Vars(r)

	traceID, err := appdash.ParseID(v["Trace"])
	if err != nil {
		return err
	}

	trace, err := a.Store.Trace(traceID)
	if err != nil {
		return err
	}

	// Get sub-span if the Span route var is present.
	if spanIDStr := v["Span"]; spanIDStr != "" {
		spanID, err := appdash.ParseID(spanIDStr)
		if err != nil {
			return err
		}
		trace = trace.FindSpan(spanID)
		if trace == nil {
			return errors.New("could not find the specified trace span")
		}
	}

	// We could use a separate handler for this, but as we need the above to
	// determine the correct trace (or therein sub-trace), we just handle any
	// JSON profile requests here.
	if path.Base(r.URL.Path) == "profile" {
		return a.profile(trace, w)
	}

	visData, err := a.d3timeline(trace)
	if err != nil {
		return err
	}

	// Determine the profile URL.
	var profile *url.URL
	if trace.ID.Parent == 0 {
		profile, err = a.Router.URLToTraceProfile(trace.Span.ID.Trace)
	} else {
		profile, err = a.Router.URLToTraceSpanProfile(trace.Span.ID.Trace, trace.Span.ID.Span)
	}
	if err != nil {
		return err
	}

	return a.renderTemplate(w, r, "trace.html", http.StatusOK, &struct {
		TemplateCommon
		Trace      *appdash.Trace
		VisData    []timelineItem
		ProfileURL string
	}{
		Trace:      trace,
		VisData:    visData,
		ProfileURL: profile.String(),
	})
}

func (a *App) serveTraces(w http.ResponseWriter, r *http.Request) error {
	traces, err := a.Queryer.Traces()
	if err != nil {
		return err
	}

	// Sort the traces by ID to ensure that the display order doesn't change upon
	// multiple page reloads if Queryer.Traces is e.g. backed by a map (which has
	// a random iteration order).
	sort.Sort(tracesByID(traces))

	return a.renderTemplate(w, r, "traces.html", http.StatusOK, &struct {
		TemplateCommon
		Traces []*appdash.Trace
	}{
		Traces: traces,
	})
}

func (a *App) serveTraceUpload(w http.ResponseWriter, r *http.Request) error {
	// Read the uploaded JSON trace data.
	defer r.Body.Close()
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	// Unmarshal the trace.
	var traces []*appdash.Trace
	err = json.Unmarshal(data, &traces)
	if err != nil {
		return err
	}

	// Collect the unmarshaled traces, ignoring any previously existing ones (i.e.
	// ones that would collide / be merged together).
	for _, trace := range traces {
		_, err = a.Store.Trace(trace.Span.ID.Trace)
		if err != appdash.ErrTraceNotFound {
			// The trace collides with an existing trace, ignore it.
			continue
		}

		// Collect the trace (store it for later viewing).
		if err = collectTrace(a.Store, trace); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>