<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	_ "github.com/couchbaselabs/go_n1ql"
	"log"
	"os"
)

var serverURL = flag.String("server", "http://localhost:9000",
	"couchbase server URL")

func main() {

	flag.Parse()
	n1ql, err := sql.Open("n1ql", *serverURL)
	if err != nil {
		log.Fatal(err)
	}

	err = n1ql.Ping()
	if err != nil {
		log.Fatal(err)
	}

	// Set query parameters
	os.Setenv("n1ql_timeout", "10s")
	ac := []byte(`[{"user": "admin:Administrator", "pass": "asdasd"}]`)
	os.Setenv("n1ql_creds", string(ac))

	result, err := n1ql.Exec("Create primary index on `beer-sample`")
	if err != nil {
		log.Fatal(err)
	}

	name := "brewery"
	rows, err := n1ql.Query("select * from `beer-sample` where type = ?", name)

	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()
	rowsReturned := 0
	for rows.Next() {
		var contacts string
		if err := rows.Scan(&contacts); err != nil {
			log.Fatal(err)
		}
		rowsReturned++
	}

	log.Printf("Rows returned %d : \n", rowsReturned)
	if err := rows.Err(); err != nil {
		log.Fatal(err)
	}

	stmt, err := n1ql.Prepare("Upsert INTO default values (?,?)")
	if err != nil {
		log.Fatal(err)
	}
	// Map Values need to be marshaled
	value, _ := json.Marshal(map[string]interface{}{"name": "irish", "type": "contact"})

	var rowsAffected int64
	for i := 0; i < 10000; i++ {
		key := fmt.Sprintf("irish%d", i)
		value := make(map[string]interface{})
		value["name"] = key
		value["type"] = "contact"
		valueBytes, _ := json.Marshal(value)
		result, err = stmt.Exec(key, valueBytes)

		if err != nil {
			fmt.Errorf(" Failed here %v", err)
			continue
		}

		ra, err := result.RowsAffected()
		if err != nil {
			log.Fatal(err)
		}
		rowsAffected += ra
	}

	log.Printf("Total Rows Affected %d", rowsAffected)

	stmt.Close()
	result, err = stmt.Exec("test", "this shouldn't work")
	if err == nil {
		log.Fatal("Statement not closed")
	}

	result, err = n1ql.Exec("delete from default  use keys ? ", "irish")
	if err != nil {
		log.Fatal(err)
	}

	keys := make([]string, 0)
	for i := 0; i < 100000; i++ {
		keys = append(keys, fmt.Sprintf("irish%d", i))
	}

	value, _ = json.Marshal(keys)
	result, err = n1ql.Exec("delete from default use keys ?", value)
	if err != nil {
		log.Fatal(err)
	}

}
<commit_msg>MB-16327.Check for errors in the query response<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	_ "github.com/couchbaselabs/go_n1ql"
	"log"
	"os"
)

var serverURL = flag.String("server", "localhost:8093",
	"couchbase server URL")

func main() {

	flag.Parse()
	n1ql, err := sql.Open("n1ql", *serverURL)
	if err != nil {
		log.Fatal(err)
	}

	err = n1ql.Ping()
	if err != nil {
		log.Fatal(err)
	}

	// Set query parameters
	os.Setenv("n1ql_timeout", "10s")
	ac := []byte(`[{"user": "admin:Administrator", "pass": "aaasdasd"}]`)
	os.Setenv("n1ql_creds", string(ac))

	result, err := n1ql.Exec("Create primary index on `beer-sample`")
	if err != nil {
		log.Printf(" heeerre")
		log.Fatal(err)
	}

	name := "brewery"
	rows, err := n1ql.Query("select * from `beer-sample` where type = ?", name)

	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()
	rowsReturned := 0
	for rows.Next() {
		var contacts string
		if err := rows.Scan(&contacts); err != nil {
			log.Fatal(err)
		}
		rowsReturned++
	}

	log.Printf("Rows returned %d : \n", rowsReturned)
	if err := rows.Err(); err != nil {
		log.Fatal(err)
	}

	stmt, err := n1ql.Prepare("Upsert INTO default values (?,?)")
	if err != nil {
		log.Fatal(err)
	}
	// Map Values need to be marshaled
	value, _ := json.Marshal(map[string]interface{}{"name": "irish", "type": "contact"})

	var rowsAffected int64
	for i := 0; i < 10000; i++ {
		key := fmt.Sprintf("irish%d", i)
		value := make(map[string]interface{})
		value["name"] = key
		value["type"] = "contact"
		valueBytes, _ := json.Marshal(value)
		result, err = stmt.Exec(key, valueBytes)

		if err != nil {
			fmt.Errorf(" Failed here %v", err)
			continue
		}

		ra, err := result.RowsAffected()
		if err != nil {
			log.Fatal(err)
		}
		rowsAffected += ra
	}

	log.Printf("Total Rows Affected %d", rowsAffected)

	stmt.Close()
	result, err = stmt.Exec("test", "this shouldn't work")
	if err == nil {
		log.Fatal("Statement not closed")
	}

	result, err = n1ql.Exec("delete from default  use keys ? ", "irish")
	if err != nil {
		log.Fatal(err)
	}

	keys := make([]string, 0)
	for i := 0; i < 100000; i++ {
		keys = append(keys, fmt.Sprintf("irish%d", i))
	}

	value, _ = json.Marshal(keys)
	result, err = n1ql.Exec("delete from default use keys ?", value)
	if err != nil {
		log.Fatal(err)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"os/exec"
	"runtime/pprof"
	"strings"

	"github.com/octavore/delta/delta"
	"github.com/octavore/delta/delta/formatter"
)

func main() {
	open := flag.Bool("open", false, "open the file in the gui")
	html := flag.Bool("html", false, "print out html")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	flag.Parse()
	if flag.NArg() < 2 {
		flag.PrintDefaults()
		return
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	pathFrom := flag.Arg(0)
	pathTo := flag.Arg(1)
	pathBase := os.Getenv("BASE")
	if *open {
		openDiff(pathBase, pathFrom, pathTo)
		return
	}
	printDiff(pathFrom, pathTo, *html)
}

// openDiffs diffs the given files and writes the result to a tempfile,
// then opens it in the gui.
func openDiff(pathBase, pathFrom, pathTo string) {
	d, err := diff(pathFrom, pathTo)
	if err != nil {
		os.Stderr.WriteString(err.Error())
		return
	}
	f, err := ioutil.TempFile("", "delta-diff")
	if err != nil {
		os.Stderr.WriteString(err.Error())
		return
	}
	io.WriteString(f, formatter.HTML(d))

	dir, _ := os.Getwd()
	u, _ := url.Parse("delta://openset")
	v := url.Values{}

	tmpFrom := strings.HasPrefix(pathFrom, os.TempDir())
	tmpTo := strings.HasPrefix(pathTo, os.TempDir())
	if tmpFrom && !tmpTo {
		pathFrom = pathTo
	} else if !tmpFrom && tmpTo {
		pathTo = pathFrom
	}

	v.Add("wd", dir)
	v.Add("base", pathBase)
	v.Add("left", pathFrom)
	v.Add("right", pathTo)
	v.Add("diff", f.Name())
	u.RawQuery = v.Encode()
	exec.Command("open", u.String()).Run()
}

// diff reads in files in pathFrom and pathTo, and returns a diff
func diff(pathFrom, pathTo string) (*delta.DiffSolution, error) {
	from, err := ioutil.ReadFile(pathFrom)
	if err != nil {
		return nil, fmt.Errorf("error reading %q: %v", pathFrom, err)
	}
	to, err := ioutil.ReadFile(pathTo)
	if err != nil {
		return nil, fmt.Errorf("error reading %q: %v", pathTo, err)
	}
	return delta.DiffHistogram(string(from), string(to)), nil
}

func printDiff(pathFrom, pathTo string, html bool) {
	d, err := diff(pathFrom, pathTo)
	if err != nil {
		os.Stderr.WriteString(err.Error())
		return
	}
	if html {
		fmt.Println(formatter.HTML(d))
		return
	}
	fmt.Println(formatter.ColoredText(d))
}
<commit_msg>Add version to delta.<commit_after>package main

/*
	delta is a command-line diff utility.
	Usage:
		`delta <file1> <file2>`
*/

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"os/exec"
	"runtime/pprof"
	"strings"

	"github.com/octavore/delta/delta"
	"github.com/octavore/delta/delta/formatter"
)

const VERSION = "0.1.0"

func main() {
	open := flag.Bool("open", false, "open the file in the gui")
	html := flag.Bool("html", false, "print out html")
	version := flag.Bool("version", false, "display delta version")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	flag.Parse()
	if *version {
		fmt.Println("delta", VERSION)
		return
	}
	if flag.NArg() < 2 {
		flag.PrintDefaults()
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	pathFrom := flag.Arg(0)
	pathTo := flag.Arg(1)
	pathBase := os.Getenv("BASE")
	if *open {
		openDiff(pathBase, pathFrom, pathTo)
		return
	}
	printDiff(pathFrom, pathTo, *html)
}

// openDiffs diffs the given files and writes the result to a tempfile,
// then opens it in the gui.
func openDiff(pathBase, pathFrom, pathTo string) {
	d, err := diff(pathFrom, pathTo)
	if err != nil {
		os.Stderr.WriteString(err.Error())
		return
	}
	f, err := ioutil.TempFile("", "delta-diff")
	if err != nil {
		os.Stderr.WriteString(err.Error())
		return
	}
	io.WriteString(f, formatter.HTML(d))

	dir, _ := os.Getwd()
	u, _ := url.Parse("delta://openset")
	v := url.Values{}

	tmpFrom := strings.HasPrefix(pathFrom, os.TempDir())
	tmpTo := strings.HasPrefix(pathTo, os.TempDir())
	if tmpFrom && !tmpTo {
		pathFrom = pathTo
	} else if !tmpFrom && tmpTo {
		pathTo = pathFrom
	}

	v.Add("wd", dir)
	v.Add("base", pathBase)
	v.Add("left", pathFrom)
	v.Add("right", pathTo)
	v.Add("v", VERSION)
	v.Add("diff", f.Name())
	u.RawQuery = v.Encode()
	exec.Command("open", u.String()).Run()
}

// diff reads in files in pathFrom and pathTo, and returns a diff
func diff(pathFrom, pathTo string) (*delta.DiffSolution, error) {
	from, err := ioutil.ReadFile(pathFrom)
	if err != nil {
		return nil, fmt.Errorf("error reading %q: %v", pathFrom, err)
	}
	to, err := ioutil.ReadFile(pathTo)
	if err != nil {
		return nil, fmt.Errorf("error reading %q: %v", pathTo, err)
	}
	return delta.DiffHistogram(string(from), string(to)), nil
}

func printDiff(pathFrom, pathTo string, html bool) {
	d, err := diff(pathFrom, pathTo)
	if err != nil {
		os.Stderr.WriteString(err.Error())
		return
	}
	if html {
		fmt.Println(formatter.HTML(d))
		return
	}
	fmt.Println(formatter.ColoredText(d))
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"path"
	"sync"

	"github.com/OpenBazaar/openbazaar-go/repo"
	"github.com/OpenBazaar/spvwallet"
	_ "github.com/mutecomm/go-sqlcipher"
	"github.com/op/go-logging"
)

var log = logging.MustGetLogger("db")

type SQLiteDatastore struct {
	config          repo.Config
	followers       repo.Followers
	following       repo.Following
	offlineMessages repo.OfflineMessages
	pointers        repo.Pointers
	keys            spvwallet.Keys
	state           spvwallet.State
	stxos           spvwallet.Stxos
	txns            spvwallet.Txns
	utxos           spvwallet.Utxos
	watchedScripts  spvwallet.WatchedScripts
	settings        repo.Settings
	inventory       repo.Inventory
	purchases       repo.Purchases
	sales           repo.Sales
	cases           repo.Cases
	chat            repo.Chat
	db              *sql.DB
	lock            *sync.Mutex
}

func Create(repoPath, password string, testnet bool) (*SQLiteDatastore, error) {
	var dbPath string
	if testnet {
		dbPath = path.Join(repoPath, "datastore", "testnet.db")
	} else {
		dbPath = path.Join(repoPath, "datastore", "mainnet.db")
	}
	conn, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, err
	}
	if password != "" {
		p := "pragma key = '" + password + "';"
		conn.Exec(p)
	}

	l := new(sync.Mutex)
	sqliteDB := &SQLiteDatastore{
		config: &ConfigDB{
			db:   conn,
			lock: l,
			path: dbPath,
		},
		followers: &FollowerDB{
			db:   conn,
			lock: l,
		},
		following: &FollowingDB{
			db:   conn,
			lock: l,
		},
		offlineMessages: &OfflineMessagesDB{
			db:   conn,
			lock: l,
		},
		pointers: &PointersDB{
			db:   conn,
			lock: l,
		},
		keys: &KeysDB{
			db:   conn,
			lock: l,
		},
		state: &StateDB{
			db:   conn,
			lock: l,
		},
		stxos: &StxoDB{
			db:   conn,
			lock: l,
		},
		txns: &TxnsDB{
			db:   conn,
			lock: l,
		},
		utxos: &UtxoDB{
			db:   conn,
			lock: l,
		},
		settings: &SettingsDB{
			db:   conn,
			lock: l,
		},
		inventory: &InventoryDB{
			db:   conn,
			lock: l,
		},
		purchases: &PurchasesDB{
			db:   conn,
			lock: l,
		},
		sales: &SalesDB{
			db:   conn,
			lock: l,
		},
		watchedScripts: &WatchedScriptsDB{
			db:   conn,
			lock: l,
		},
		cases: &CasesDB{
			db:   conn,
			lock: l,
		},
		chat: &ChatDB{
			db:   conn,
			lock: l,
		},
		db:   conn,
		lock: l,
	}

	return sqliteDB, nil
}

func (d *SQLiteDatastore) Close() {
	d.db.Close()
}

func (d *SQLiteDatastore) Config() repo.Config {
	return d.config
}

func (d *SQLiteDatastore) Followers() repo.Followers {
	return d.followers
}

func (d *SQLiteDatastore) Following() repo.Following {
	return d.following
}

func (d *SQLiteDatastore) OfflineMessages() repo.OfflineMessages {
	return d.offlineMessages
}

func (d *SQLiteDatastore) Pointers() repo.Pointers {
	return d.pointers
}

func (d *SQLiteDatastore) Keys() spvwallet.Keys {
	return d.keys
}

func (d *SQLiteDatastore) State() spvwallet.State {
	return d.state
}

func (d *SQLiteDatastore) Stxos() spvwallet.Stxos {
	return d.stxos
}

func (d *SQLiteDatastore) Txns() spvwallet.Txns {
	return d.txns
}

func (d *SQLiteDatastore) Utxos() spvwallet.Utxos {
	return d.utxos
}

func (d *SQLiteDatastore) Settings() repo.Settings {
	return d.settings
}

func (d *SQLiteDatastore) Inventory() repo.Inventory {
	return d.inventory
}

func (d *SQLiteDatastore) Purchases() repo.Purchases {
	return d.purchases
}

func (d *SQLiteDatastore) Sales() repo.Sales {
	return d.sales
}

func (d *SQLiteDatastore) WatchedScripts() spvwallet.WatchedScripts {
	return d.watchedScripts
}

func (d *SQLiteDatastore) Cases() repo.Cases {
	return d.cases
}

func (d *SQLiteDatastore) Chat() repo.Chat {
	return d.chat
}

func (d *SQLiteDatastore) Copy(dbPath string, password string) error {
	d.lock.Lock()
	defer d.lock.Unlock()
	var cp string
	stmt := "select name from sqlite_master where type='table'"
	rows, err := d.db.Query(stmt)
	if err != nil {
		log.Error(err)
		return err
	}
	var tables []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return err
		}
		tables = append(tables, name)
	}
	if password == "" {
		cp = `attach database '` + dbPath + `' as plaintext key '';`
		for _, name := range tables {
			cp = cp + "insert into plaintext." + name + " select * from main." + name + ";"
		}
	} else {
		cp = `attach database '` + dbPath + `' as encrypted key '` + password + `';`
		for _, name := range tables {
			cp = cp + "insert into encrypted." + name + " select * from main." + name + ";"
		}
	}

	_, err = d.db.Exec(cp)
	if err != nil {
		return err
	}

	return nil
}

func initDatabaseTables(db *sql.DB, password string) error {
	var sqlStmt string
	if password != "" {
		sqlStmt = "PRAGMA key = '" + password + "';"
	}
	sqlStmt += `
	PRAGMA user_version = 0;
	create table config (key text primary key not null, value blob);
	create table followers (peerID text primary key not null);
	create table following (peerID text primary key not null);
	create table offlinemessages (url text primary key not null, timestamp integer);
	create table pointers (pointerID text primary key not null, key text, address text, purpose integer, timestamp integer);
	create table keys (scriptPubKey text primary key not null, purpose integer, keyIndex integer, used integer);
	create table utxos (outpoint text primary key not null, value integer, height integer, scriptPubKey text, freeze int);
	create table stxos (outpoint text primary key not null, value integer, height integer, scriptPubKey text, spendHeight integer, spendTxid text);
	create table txns (txid text primary key not null, tx blob);
	create table state (key text primary key not null, value text);
	create table inventory (slug text primary key not null, count integer);
	create table purchases (orderID text primary key not null, contract blob, state integer, read integer, date integer, total integer, thumbnail text, vendorID text, vendorBlockchainID text, title text, shippingName text, shippingAddress text, paymentAddr text, funded integer, transactions blob);
	create table sales (orderID text primary key not null, contract blob, state integer, read integer, date integer, total integer, thumbnail text, buyerID text, buyerBlockchainID text, title text, shippingName text, shippingAddress text, paymentAddr text, funded integer, transactions blob);
	create table watchedscripts (scriptPubKey text primary key not null);
	create table cases (caseID text primary key not null, buyerContract blob, vendorContract blob, buyerValidationErrors blob, vendorValidationErrors blob, buyerPayoutAddress text, vendorPayoutAddress text, buyerOutpoints blob, vendorOutpoints blob, state integer, read integer, date integer, buyerOpened integer, claim text, disputeResolution blob);
	create table chat (peerID text, subject text, message text, read integer, timestamp integer, outgoing integer);
	create index chat_peerID ON chat(peerId);
	create index chat_subject ON chat(subject);
	create index chat_read ON chat(read);
	create index chat_timestamp ON chat(timestamp);
	`
	_, err := db.Exec(sqlStmt)
	if err != nil {
		return err
	}
	return nil
}

type ConfigDB struct {
	db   *sql.DB
	lock *sync.Mutex
	path string
}

func (c *ConfigDB) Init(mnemonic string, identityKey []byte, password string) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	if err := initDatabaseTables(c.db, password); err != nil {
		return err
	}
	tx, err := c.db.Begin()
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare("insert into config(key, value) values(?,?)")
	if err != nil {
		return err
	}
	defer stmt.Close()
	_, err = stmt.Exec("mnemonic", mnemonic)
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = stmt.Exec("identityKey", identityKey)
	if err != nil {
		tx.Rollback()
		return err
	}
	tx.Commit()
	return nil

}

func (c *ConfigDB) GetMnemonic() (string, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	stmt, err := c.db.Prepare("select value from config where key=?")
	defer stmt.Close()
	var mnemonic string
	err = stmt.QueryRow("mnemonic").Scan(&mnemonic)
	if err != nil {
		log.Fatal(err)
	}
	return mnemonic, nil
}

func (c *ConfigDB) GetIdentityKey() ([]byte, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	stmt, err := c.db.Prepare("select value from config where key=?")
	defer stmt.Close()
	var identityKey []byte
	err = stmt.QueryRow("identityKey").Scan(&identityKey)
	if err != nil {
		log.Fatal(err)
	}
	return identityKey, nil
}

func (c *ConfigDB) IsEncrypted() bool {
	pwdCheck := "select count(*) from sqlite_master;"
	_, err := c.db.Exec(pwdCheck) // Fails if wrong password is entered
	if err != nil {
		return true
	}
	return false
}
<commit_msg>Create index in single line<commit_after>package db

import (
	"database/sql"
	"path"
	"sync"

	"github.com/OpenBazaar/openbazaar-go/repo"
	"github.com/OpenBazaar/spvwallet"
	_ "github.com/mutecomm/go-sqlcipher"
	"github.com/op/go-logging"
)

var log = logging.MustGetLogger("db")

type SQLiteDatastore struct {
	config          repo.Config
	followers       repo.Followers
	following       repo.Following
	offlineMessages repo.OfflineMessages
	pointers        repo.Pointers
	keys            spvwallet.Keys
	state           spvwallet.State
	stxos           spvwallet.Stxos
	txns            spvwallet.Txns
	utxos           spvwallet.Utxos
	watchedScripts  spvwallet.WatchedScripts
	settings        repo.Settings
	inventory       repo.Inventory
	purchases       repo.Purchases
	sales           repo.Sales
	cases           repo.Cases
	chat            repo.Chat
	db              *sql.DB
	lock            *sync.Mutex
}

func Create(repoPath, password string, testnet bool) (*SQLiteDatastore, error) {
	var dbPath string
	if testnet {
		dbPath = path.Join(repoPath, "datastore", "testnet.db")
	} else {
		dbPath = path.Join(repoPath, "datastore", "mainnet.db")
	}
	conn, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, err
	}
	if password != "" {
		p := "pragma key = '" + password + "';"
		conn.Exec(p)
	}

	l := new(sync.Mutex)
	sqliteDB := &SQLiteDatastore{
		config: &ConfigDB{
			db:   conn,
			lock: l,
			path: dbPath,
		},
		followers: &FollowerDB{
			db:   conn,
			lock: l,
		},
		following: &FollowingDB{
			db:   conn,
			lock: l,
		},
		offlineMessages: &OfflineMessagesDB{
			db:   conn,
			lock: l,
		},
		pointers: &PointersDB{
			db:   conn,
			lock: l,
		},
		keys: &KeysDB{
			db:   conn,
			lock: l,
		},
		state: &StateDB{
			db:   conn,
			lock: l,
		},
		stxos: &StxoDB{
			db:   conn,
			lock: l,
		},
		txns: &TxnsDB{
			db:   conn,
			lock: l,
		},
		utxos: &UtxoDB{
			db:   conn,
			lock: l,
		},
		settings: &SettingsDB{
			db:   conn,
			lock: l,
		},
		inventory: &InventoryDB{
			db:   conn,
			lock: l,
		},
		purchases: &PurchasesDB{
			db:   conn,
			lock: l,
		},
		sales: &SalesDB{
			db:   conn,
			lock: l,
		},
		watchedScripts: &WatchedScriptsDB{
			db:   conn,
			lock: l,
		},
		cases: &CasesDB{
			db:   conn,
			lock: l,
		},
		chat: &ChatDB{
			db:   conn,
			lock: l,
		},
		db:   conn,
		lock: l,
	}

	return sqliteDB, nil
}

func (d *SQLiteDatastore) Close() {
	d.db.Close()
}

func (d *SQLiteDatastore) Config() repo.Config {
	return d.config
}

func (d *SQLiteDatastore) Followers() repo.Followers {
	return d.followers
}

func (d *SQLiteDatastore) Following() repo.Following {
	return d.following
}

func (d *SQLiteDatastore) OfflineMessages() repo.OfflineMessages {
	return d.offlineMessages
}

func (d *SQLiteDatastore) Pointers() repo.Pointers {
	return d.pointers
}

func (d *SQLiteDatastore) Keys() spvwallet.Keys {
	return d.keys
}

func (d *SQLiteDatastore) State() spvwallet.State {
	return d.state
}

func (d *SQLiteDatastore) Stxos() spvwallet.Stxos {
	return d.stxos
}

func (d *SQLiteDatastore) Txns() spvwallet.Txns {
	return d.txns
}

func (d *SQLiteDatastore) Utxos() spvwallet.Utxos {
	return d.utxos
}

func (d *SQLiteDatastore) Settings() repo.Settings {
	return d.settings
}

func (d *SQLiteDatastore) Inventory() repo.Inventory {
	return d.inventory
}

func (d *SQLiteDatastore) Purchases() repo.Purchases {
	return d.purchases
}

func (d *SQLiteDatastore) Sales() repo.Sales {
	return d.sales
}

func (d *SQLiteDatastore) WatchedScripts() spvwallet.WatchedScripts {
	return d.watchedScripts
}

func (d *SQLiteDatastore) Cases() repo.Cases {
	return d.cases
}

func (d *SQLiteDatastore) Chat() repo.Chat {
	return d.chat
}

func (d *SQLiteDatastore) Copy(dbPath string, password string) error {
	d.lock.Lock()
	defer d.lock.Unlock()
	var cp string
	stmt := "select name from sqlite_master where type='table'"
	rows, err := d.db.Query(stmt)
	if err != nil {
		log.Error(err)
		return err
	}
	var tables []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return err
		}
		tables = append(tables, name)
	}
	if password == "" {
		cp = `attach database '` + dbPath + `' as plaintext key '';`
		for _, name := range tables {
			cp = cp + "insert into plaintext." + name + " select * from main." + name + ";"
		}
	} else {
		cp = `attach database '` + dbPath + `' as encrypted key '` + password + `';`
		for _, name := range tables {
			cp = cp + "insert into encrypted." + name + " select * from main." + name + ";"
		}
	}

	_, err = d.db.Exec(cp)
	if err != nil {
		return err
	}

	return nil
}

func initDatabaseTables(db *sql.DB, password string) error {
	var sqlStmt string
	if password != "" {
		sqlStmt = "PRAGMA key = '" + password + "';"
	}
	sqlStmt += `
	PRAGMA user_version = 0;
	create table config (key text primary key not null, value blob);
	create table followers (peerID text primary key not null);
	create table following (peerID text primary key not null);
	create table offlinemessages (url text primary key not null, timestamp integer);
	create table pointers (pointerID text primary key not null, key text, address text, purpose integer, timestamp integer);
	create table keys (scriptPubKey text primary key not null, purpose integer, keyIndex integer, used integer);
	create table utxos (outpoint text primary key not null, value integer, height integer, scriptPubKey text, freeze int);
	create table stxos (outpoint text primary key not null, value integer, height integer, scriptPubKey text, spendHeight integer, spendTxid text);
	create table txns (txid text primary key not null, tx blob);
	create table state (key text primary key not null, value text);
	create table inventory (slug text primary key not null, count integer);
	create table purchases (orderID text primary key not null, contract blob, state integer, read integer, date integer, total integer, thumbnail text, vendorID text, vendorBlockchainID text, title text, shippingName text, shippingAddress text, paymentAddr text, funded integer, transactions blob);
	create table sales (orderID text primary key not null, contract blob, state integer, read integer, date integer, total integer, thumbnail text, buyerID text, buyerBlockchainID text, title text, shippingName text, shippingAddress text, paymentAddr text, funded integer, transactions blob);
	create table watchedscripts (scriptPubKey text primary key not null);
	create table cases (caseID text primary key not null, buyerContract blob, vendorContract blob, buyerValidationErrors blob, vendorValidationErrors blob, buyerPayoutAddress text, vendorPayoutAddress text, buyerOutpoints blob, vendorOutpoints blob, state integer, read integer, date integer, buyerOpened integer, claim text, disputeResolution blob);
	create table chat (peerID text, subject text, message text, read integer, timestamp integer, outgoing integer);
	create index index_chat ON chat(peerID, subject, read, timestamp);
	`
	_, err := db.Exec(sqlStmt)
	if err != nil {
		return err
	}
	return nil
}

type ConfigDB struct {
	db   *sql.DB
	lock *sync.Mutex
	path string
}

func (c *ConfigDB) Init(mnemonic string, identityKey []byte, password string) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	if err := initDatabaseTables(c.db, password); err != nil {
		return err
	}
	tx, err := c.db.Begin()
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare("insert into config(key, value) values(?,?)")
	if err != nil {
		return err
	}
	defer stmt.Close()
	_, err = stmt.Exec("mnemonic", mnemonic)
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = stmt.Exec("identityKey", identityKey)
	if err != nil {
		tx.Rollback()
		return err
	}
	tx.Commit()
	return nil

}

func (c *ConfigDB) GetMnemonic() (string, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	stmt, err := c.db.Prepare("select value from config where key=?")
	defer stmt.Close()
	var mnemonic string
	err = stmt.QueryRow("mnemonic").Scan(&mnemonic)
	if err != nil {
		log.Fatal(err)
	}
	return mnemonic, nil
}

func (c *ConfigDB) GetIdentityKey() ([]byte, error) {
	c.lock.Lock()
	defer c.lock.Unlock()
	stmt, err := c.db.Prepare("select value from config where key=?")
	defer stmt.Close()
	var identityKey []byte
	err = stmt.QueryRow("identityKey").Scan(&identityKey)
	if err != nil {
		log.Fatal(err)
	}
	return identityKey, nil
}

func (c *ConfigDB) IsEncrypted() bool {
	pwdCheck := "select count(*) from sqlite_master;"
	_, err := c.db.Exec(pwdCheck) // Fails if wrong password is entered
	if err != nil {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"fmt"

	"github.com/coreos/etcd/clientv3"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
)

var (
	grantPermissionPrefix bool
	permFromKey           bool
)

// NewRoleCommand returns the cobra command for "role".
func NewRoleCommand() *cobra.Command {
	ac := &cobra.Command{
		Use:   "role <subcommand>",
		Short: "Role related commands",
	}

	ac.AddCommand(newRoleAddCommand())
	ac.AddCommand(newRoleDeleteCommand())
	ac.AddCommand(newRoleGetCommand())
	ac.AddCommand(newRoleListCommand())
	ac.AddCommand(newRoleGrantPermissionCommand())
	ac.AddCommand(newRoleRevokePermissionCommand())

	return ac
}

func newRoleAddCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "add <role name>",
		Short: "Adds a new role",
		Run:   roleAddCommandFunc,
	}
}

func newRoleDeleteCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "delete <role name>",
		Short: "Deletes a role",
		Run:   roleDeleteCommandFunc,
	}
}

func newRoleGetCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "get <role name>",
		Short: "Gets detailed information of a role",
		Run:   roleGetCommandFunc,
	}
}

func newRoleListCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "Lists all roles",
		Run:   roleListCommandFunc,
	}
}

func newRoleGrantPermissionCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "grant-permission [options] <role name> <permission type> <key> [endkey]",
		Short: "Grants a key to a role",
		Run:   roleGrantPermissionCommandFunc,
	}

	cmd.Flags().BoolVar(&grantPermissionPrefix, "prefix", false, "grant a prefix permission")
	cmd.Flags().BoolVar(&permFromKey, "from-key", false, "grant a permission of keys that are greater than or equal to the given key using byte compare")

	return cmd
}

func newRoleRevokePermissionCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "revoke-permission <role name> <key> [endkey]",
		Short: "Revokes a key from a role",
		Run:   roleRevokePermissionCommandFunc,
	}

	cmd.Flags().BoolVar(&permFromKey, "from-key", false, "grant a permission of keys that are greater than or equal to the given key using byte compare")

	return cmd
}

// roleAddCommandFunc executes the "role add" command.
func roleAddCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role add command requires role name as its argument."))
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleAdd(context.TODO(), args[0])
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleAdd(args[0], *resp)
}

// roleDeleteCommandFunc executes the "role delete" command.
func roleDeleteCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role delete command requires role name as its argument."))
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleDelete(context.TODO(), args[0])
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleDelete(args[0], *resp)
}

// roleGetCommandFunc executes the "role get" command.
func roleGetCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role get command requires role name as its argument."))
	}

	name := args[0]
	resp, err := mustClientFromCmd(cmd).Auth.RoleGet(context.TODO(), name)
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleGet(name, *resp)
}

// roleListCommandFunc executes the "role list" command.
func roleListCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 0 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role list command requires no arguments."))
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleList(context.TODO())
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleList(*resp)
}

// roleGrantPermissionCommandFunc executes the "role grant-permission" command.
func roleGrantPermissionCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) < 3 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role grant command requires role name, permission type, and key [endkey] as its argument."))
	}

	perm, err := clientv3.StrToPermissionType(args[1])
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	rangeEnd := ""
	if 4 <= len(args) {
		if grantPermissionPrefix {
			ExitWithError(ExitBadArgs, fmt.Errorf("don't pass both of --prefix option and range end to grant permission command"))
		}

		if permFromKey {
			ExitWithError(ExitBadArgs, fmt.Errorf("don't pass both of --from-key option and range end to grant permission command"))
		}

		rangeEnd = args[3]
	} else if grantPermissionPrefix {
		if permFromKey {
			ExitWithError(ExitBadArgs, fmt.Errorf("don't pass both of --from-key option and --prefix option to grant permission command"))
		}

		rangeEnd = clientv3.GetPrefixRangeEnd(args[2])
	} else if permFromKey {
		rangeEnd = "\x00"
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleGrantPermission(context.TODO(), args[0], args[2], rangeEnd, perm)
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleGrantPermission(args[0], *resp)
}

// roleRevokePermissionCommandFunc executes the "role revoke-permission" command.
func roleRevokePermissionCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) < 2 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role revoke-permission command requires role name and key [endkey] as its argument."))
	}

	rangeEnd := ""
	if 3 <= len(args) {
		rangeEnd = args[2]
	} else if permFromKey {
		rangeEnd = "\x00"
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleRevokePermission(context.TODO(), args[0], args[1], rangeEnd)
	if err != nil {
		ExitWithError(ExitError, err)
	}
	display.RoleRevokePermission(args[0], args[1], rangeEnd, *resp)
}
<commit_msg>ctlv3: add --prefix support to role revoke-permission, cleanup role flag handling<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"fmt"

	"github.com/coreos/etcd/clientv3"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
)

var (
	rolePermPrefix  bool
	rolePermFromKey bool
)

// NewRoleCommand returns the cobra command for "role".
func NewRoleCommand() *cobra.Command {
	ac := &cobra.Command{
		Use:   "role <subcommand>",
		Short: "Role related commands",
	}

	ac.AddCommand(newRoleAddCommand())
	ac.AddCommand(newRoleDeleteCommand())
	ac.AddCommand(newRoleGetCommand())
	ac.AddCommand(newRoleListCommand())
	ac.AddCommand(newRoleGrantPermissionCommand())
	ac.AddCommand(newRoleRevokePermissionCommand())

	return ac
}

func newRoleAddCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "add <role name>",
		Short: "Adds a new role",
		Run:   roleAddCommandFunc,
	}
}

func newRoleDeleteCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "delete <role name>",
		Short: "Deletes a role",
		Run:   roleDeleteCommandFunc,
	}
}

func newRoleGetCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "get <role name>",
		Short: "Gets detailed information of a role",
		Run:   roleGetCommandFunc,
	}
}

func newRoleListCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "Lists all roles",
		Run:   roleListCommandFunc,
	}
}

func newRoleGrantPermissionCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "grant-permission [options] <role name> <permission type> <key> [endkey]",
		Short: "Grants a key to a role",
		Run:   roleGrantPermissionCommandFunc,
	}

	cmd.Flags().BoolVar(&rolePermPrefix, "prefix", false, "grant a prefix permission")
	cmd.Flags().BoolVar(&rolePermFromKey, "from-key", false, "grant a permission of keys that are greater than or equal to the given key using byte compare")

	return cmd
}

func newRoleRevokePermissionCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "revoke-permission <role name> <key> [endkey]",
		Short: "Revokes a key from a role",
		Run:   roleRevokePermissionCommandFunc,
	}

	cmd.Flags().BoolVar(&rolePermPrefix, "prefix", false, "revoke a prefix permission")
	cmd.Flags().BoolVar(&rolePermFromKey, "from-key", false, "revoke a permission of keys that are greater than or equal to the given key using byte compare")

	return cmd
}

// roleAddCommandFunc executes the "role add" command.
func roleAddCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role add command requires role name as its argument."))
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleAdd(context.TODO(), args[0])
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleAdd(args[0], *resp)
}

// roleDeleteCommandFunc executes the "role delete" command.
func roleDeleteCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role delete command requires role name as its argument."))
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleDelete(context.TODO(), args[0])
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleDelete(args[0], *resp)
}

// roleGetCommandFunc executes the "role get" command.
func roleGetCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role get command requires role name as its argument."))
	}

	name := args[0]
	resp, err := mustClientFromCmd(cmd).Auth.RoleGet(context.TODO(), name)
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleGet(name, *resp)
}

// roleListCommandFunc executes the "role list" command.
func roleListCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 0 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role list command requires no arguments."))
	}

	resp, err := mustClientFromCmd(cmd).Auth.RoleList(context.TODO())
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleList(*resp)
}

// roleGrantPermissionCommandFunc executes the "role grant-permission" command.
func roleGrantPermissionCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) < 3 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role grant command requires role name, permission type, and key [endkey] as its argument."))
	}

	perm, err := clientv3.StrToPermissionType(args[1])
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	rangeEnd, rerr := rangeEndFromPermFlags(args[2:])
	if rerr != nil {
		ExitWithError(ExitBadArgs, rerr)
	}
	resp, err := mustClientFromCmd(cmd).Auth.RoleGrantPermission(context.TODO(), args[0], args[2], rangeEnd, perm)
	if err != nil {
		ExitWithError(ExitError, err)
	}

	display.RoleGrantPermission(args[0], *resp)
}

// roleRevokePermissionCommandFunc executes the "role revoke-permission" command.
func roleRevokePermissionCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) < 2 {
		ExitWithError(ExitBadArgs, fmt.Errorf("role revoke-permission command requires role name and key [endkey] as its argument."))
	}

	rangeEnd, rerr := rangeEndFromPermFlags(args[1:])
	if rerr != nil {
		ExitWithError(ExitBadArgs, rerr)
	}
	resp, err := mustClientFromCmd(cmd).Auth.RoleRevokePermission(context.TODO(), args[0], args[1], rangeEnd)
	if err != nil {
		ExitWithError(ExitError, err)
	}
	display.RoleRevokePermission(args[0], args[1], rangeEnd, *resp)
}

func rangeEndFromPermFlags(args []string) (string, error) {
	if len(args) == 1 {
		if rolePermPrefix {
			if rolePermFromKey {
				return "", fmt.Errorf("--from-key and --prefix flags are mutually exclusive")
			}
			return clientv3.GetPrefixRangeEnd(args[0]), nil
		}
		if rolePermFromKey {
			return "\x00", nil
		}
		// single key case
		return "", nil
	}
	if rolePermPrefix {
		return "", fmt.Errorf("unexpected endkey argument with --prefix flag")
	}
	if rolePermFromKey {
		return "", fmt.Errorf("unexpected endkey argument with --from-key flag")
	}
	return args[1], nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package networking

// The ports we setup on our services.
const (
	// ServiceHTTPPort is the port that we setup our Serving and Activator K8s services for
	// HTTP/1 endpoints.
	ServiceHTTPPort = 80

	// ServiceHTTP2Port is the port that we setup our Serving and Activator K8s services for
	// HTTP/2 endpoints.
	ServiceHTTP2Port = 81

	// BackendHTTPPort is the backend, i.e. `targetPort` that we setup for HTTP services.
	BackendHTTPPort = 8012

	// BackendHTTP2Port is the backend, i.e. `targetPort` that we setup for HTTP services.
	BackendHTTP2Port = 8013

	// RequestQueueAdminPort specifies the port number for
	// health check and lifecyle hooks for queue-proxy.
	RequestQueueAdminPort = 8022

	// RequestQueueMetricsPort specifies the port number for metrics emitted
	// by queue-proxy.
	RequestQueueMetricsPort = 9090

	// ServicePortNameHTTP1 is the name of the external port of the service for HTTP/1.1
	ServicePortNameHTTP1 = "http"
	// ServicePortNameH2C is the name of the external port of the service for HTTP/2
	ServicePortNameH2C = "http2"
)

// ServicePortName returns the port for the app level protocol.
func ServicePortName(proto ProtocolType) string {
	if proto == ProtocolH2C {
		return ServicePortNameH2C
	}
	return ServicePortNameHTTP1
}

// ServicePort chooses the service (load balancer) port for the public service.
func ServicePort(proto ProtocolType) int {
	if proto == ProtocolH2C {
		return ServiceHTTP2Port
	}
	return ServiceHTTPPort
}
<commit_msg>Enable queue proxy metrics scraping by Prometheus. (#4111)<commit_after>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package networking

// The ports we setup on our services.
const (
	// ServiceHTTPPort is the port that we setup our Serving and Activator K8s services for
	// HTTP/1 endpoints.
	ServiceHTTPPort = 80

	// ServiceHTTP2Port is the port that we setup our Serving and Activator K8s services for
	// HTTP/2 endpoints.
	ServiceHTTP2Port = 81

	// BackendHTTPPort is the backend, i.e. `targetPort` that we setup for HTTP services.
	BackendHTTPPort = 8012

	// BackendHTTP2Port is the backend, i.e. `targetPort` that we setup for HTTP services.
	BackendHTTP2Port = 8013

	// QueueAdminPort specifies the port number for
	// health check and lifecycle hooks for queue-proxy.
	QueueAdminPort = 8022

	// AutoscalingQueueMetricsPort specifies the port number for metrics emitted
	// by queue-proxy for autoscaler.
	AutoscalingQueueMetricsPort = 9090

	// UserQueueMetricsPort specifies the port number for metrics emitted
	// by queue-proxy for end user.
	UserQueueMetricsPort = 9091

	// ServicePortNameHTTP1 is the name of the external port of the service for HTTP/1.1
	ServicePortNameHTTP1 = "http"

	// ServicePortNameH2C is the name of the external port of the service for HTTP/2
	ServicePortNameH2C = "http2"
)

// ServicePortName returns the port for the app level protocol.
func ServicePortName(proto ProtocolType) string {
	if proto == ProtocolH2C {
		return ServicePortNameH2C
	}
	return ServicePortNameHTTP1
}

// ServicePort chooses the service (load balancer) port for the public service.
func ServicePort(proto ProtocolType) int {
	if proto == ProtocolH2C {
		return ServiceHTTP2Port
	}
	return ServiceHTTPPort
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/chzyer/flagly"
	"github.com/chzyer/flow"
	"github.com/chzyer/logex"
	"github.com/chzyer/madq/go/bench"
	"github.com/chzyer/madq/go/debug"
)

type Madq struct {
	Bench *bench.Config `flagly:"handler"`
	Debug *debug.Config `flagly:"handler"`
}

func main() {
	madq := new(Madq)
	f := flow.New()

	flagly.Run(madq, f)

	if err := f.Wait(); err != nil {
		logex.Fatal(err)
	}
}
<commit_msg>[madq] add MAXPROCS opt<commit_after>package main

import (
	"runtime"

	"github.com/chzyer/flagly"
	"github.com/chzyer/flow"
	"github.com/chzyer/logex"
	"github.com/chzyer/madq/go/bench"
	"github.com/chzyer/madq/go/debug"
)

type Madq struct {
	CPU   int
	Bench *bench.Config `flagly:"handler"`
	Debug *debug.Config `flagly:"handler"`
}

func (m *Madq) FlaglyEnter() {
	runtime.GOMAXPROCS(m.CPU)
}

func main() {
	madq := new(Madq)
	f := flow.New()

	flagly.Run(madq, f)

	if err := f.Wait(); err != nil {
		logex.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package apiclient

import (
	"github.com/golang/protobuf/proto"

	"encoding/json"
	"fmt"
	"reflect"
)

// NewAnyValue creates and returns a new AnyValue in which the type
// information and the serialized form of a proto.Message are stored.
func NewAnyValue(pb proto.Message) (av *AnyValue, err error) {
	var value []byte
	typ := proto.MessageName(pb)
	value, err = proto.Marshal(pb)
	if err != nil {
		return
	}
	av = &AnyValue{
		TypeUrl: &typ,
		Value:   value,
	}
	return
}

// GetProtoMessage deserializes and returns the proto.Message stored
// in the AnyValue
func (av *AnyValue) GetProtoMessage() (pb proto.Message, err error) {
	typ := proto.MessageType(av.GetTypeUrl())
	if typ == nil {
		return nil, fmt.Errorf("%s has not been registered with Protobuf library", av.TypeUrl)
	}
	pb = reflect.New(typ.Elem()).Interface().(proto.Message)
	if err = proto.Unmarshal(av.Value, pb); err != nil {
		return nil, err
	}
	return
}

// MarshalJSON implements a JSON marshaller that marshals the value
// stored in av.Value.
func (av *AnyValue) MarshalJSON() ([]byte, error) {
	pb, err := av.ToProtoMessage()
	if err != nil {
		return nil, err
	}
	return json.Marshal(pb)
}
<commit_msg>ToProtoMessage to GetProtoMessage<commit_after>package apiclient

import (
	"github.com/golang/protobuf/proto"

	"encoding/json"
	"fmt"
	"reflect"
)

// NewAnyValue creates and returns a new AnyValue in which the type
// information and the serialized form of a proto.Message are stored.
func NewAnyValue(pb proto.Message) (av *AnyValue, err error) {
	var value []byte
	typ := proto.MessageName(pb)
	value, err = proto.Marshal(pb)
	if err != nil {
		return
	}
	av = &AnyValue{
		TypeUrl: &typ,
		Value:   value,
	}
	return
}

// GetProtoMessage deserializes and returns the proto.Message stored
// in the AnyValue
func (av *AnyValue) GetProtoMessage() (pb proto.Message, err error) {
	typ := proto.MessageType(av.GetTypeUrl())
	if typ == nil {
		return nil, fmt.Errorf("%s has not been registered with Protobuf library", av.TypeUrl)
	}
	pb = reflect.New(typ.Elem()).Interface().(proto.Message)
	if err = proto.Unmarshal(av.Value, pb); err != nil {
		return nil, err
	}
	return
}

// MarshalJSON implements a JSON marshaller that marshals the value
// stored in av.Value.
func (av *AnyValue) MarshalJSON() ([]byte, error) {
	pb, err := av.GetProtoMessage()
	if err != nil {
		return nil, err
	}
	return json.Marshal(pb)
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 graphx

import (
	"reflect"
	"strings"
	"testing"

	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime"
	v1 "github.com/apache/beam/sdks/go/pkg/beam/core/runtime/graphx/v1"
)

func TestEncodeType(t *testing.T) {
	t.Run("NoUnexportedFields", func(t *testing.T) {
		type MyAwesomeType struct {
			ExportedField string
		}
		rt := reflect.TypeOf((*MyAwesomeType)(nil)).Elem()

		pbT, err := encodeType(rt)
		if err != nil {
			t.Fatalf("got error = %v, want nil", err)
		}
		if got, want := pbT.Kind, v1.Type_STRUCT; got != want {
			t.Fatalf("got pbT.Kin == %v, want %v", got, want)
		}
	})
	t.Run("UnregisteredWithUnexportedField", func(t *testing.T) {
		type MyProblematicType struct {
			unexportedField string
		}
		rt := reflect.TypeOf((*MyProblematicType)(nil)).Elem()
		pbT, err := encodeType(rt)
		if err == nil {
			t.Fatalf("got type = %v, nil, want unexported field error", pbT)
		}
		if !strings.Contains(err.Error(), "type has unexported field: unexportedField") {
			t.Errorf("expected error about unexported field, got %q", err.Error())
		}
	})
	t.Run("RegisteredWithUnexportedField", func(t *testing.T) {
		type MyRegisteredType struct {
			unexportedField string
		}
		rt := reflect.TypeOf((*MyRegisteredType)(nil)).Elem()
		runtime.RegisterType(rt)
		pbT, err := encodeType(rt)
		if err != nil {
			t.Fatalf("got error = %v, want nil", err)
		}
		if got, want := pbT.Kind, v1.Type_EXTERNAL; got != want {
			t.Fatalf("got pbT.Kin == %v, want %v", got, want)
		}
	})
}
<commit_msg>!fixup typos<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 graphx

import (
	"reflect"
	"strings"
	"testing"

	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime"
	v1 "github.com/apache/beam/sdks/go/pkg/beam/core/runtime/graphx/v1"
)

func TestEncodeType(t *testing.T) {
	t.Run("NoUnexportedFields", func(t *testing.T) {
		type MyAwesomeType struct {
			ExportedField string
		}
		rt := reflect.TypeOf((*MyAwesomeType)(nil)).Elem()

		pbT, err := encodeType(rt)
		if err != nil {
			t.Fatalf("got error = %v, want nil", err)
		}
		if got, want := pbT.Kind, v1.Type_STRUCT; got != want {
			t.Fatalf("got pbT.Kind == %v, want %v", got, want)
		}
	})
	t.Run("UnregisteredWithUnexportedField", func(t *testing.T) {
		type MyProblematicType struct {
			unexportedField string
		}
		rt := reflect.TypeOf((*MyProblematicType)(nil)).Elem()
		pbT, err := encodeType(rt)
		if err == nil {
			t.Fatalf("got type = %v, nil, want unexported field error", pbT)
		}
		if !strings.Contains(err.Error(), "type has unexported field: unexportedField") {
			t.Errorf("expected error about unexported field, got %q", err.Error())
		}
	})
	t.Run("RegisteredWithUnexportedField", func(t *testing.T) {
		type MyRegisteredType struct {
			unexportedField string
		}
		rt := reflect.TypeOf((*MyRegisteredType)(nil)).Elem()
		runtime.RegisterType(rt)
		pbT, err := encodeType(rt)
		if err != nil {
			t.Fatalf("got error = %v, want nil", err)
		}
		if got, want := pbT.Kind, v1.Type_EXTERNAL; got != want {
			t.Fatalf("got pbT.Kind == %v, want %v", got, want)
		}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"html/template"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
)

type installer interface {
	Install(projectDirectory, dockerComposeFileName, projectName string) error
}

type systemdInstaller struct {
	systemdDirectory string
	commandExecutor  Executor
}

func (installer *systemdInstaller) Install(projectDirectory, dockerComposeFileName, projectName string) error {

	serviceName := getServiceName(projectName)
	serviceViewModel := serviceDefinition{
		ProjectName:       projectName,
		ProjectDirectory:  projectDirectory,
		DockerComposeFile: dockerComposeFileName,
	}

	if err := createSystemdService(serviceViewModel, installer.systemdDirectory); err != nil {
		return errors.Wrap(err, fmt.Sprintf("Failed to create a systemd service for project %q (Directory: %q, Docker Compose File: %q)", projectName, projectDirectory, dockerComposeFileName))
	}

	reloadError := installer.commandExecutor.Run("systemctl", "daemon-reload")
	if reloadError != nil {
		return errors.Wrap(reloadError, "Failed to reload the systemd configuration")
	}

	enableError := installer.commandExecutor.Run("systemctl", "enable", serviceName)
	if enableError != nil {
		return errors.Wrap(enableError, fmt.Sprintf("Failed to enable %q", serviceName))
	}

	startError := installer.commandExecutor.Run("systemctl", "start", serviceName)
	if startError != nil {
		return errors.Wrap(startError, fmt.Sprintf("Failed to start %q", serviceName))
	}

	return nil
}

func createSystemdService(service serviceDefinition, targetDirectory string) error {
	serviceFilePath := filepath.Join(targetDirectory, getServiceName(service.ProjectName))
	file, err := os.OpenFile(serviceFilePath, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0664)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("Failed to open the systemd service file: %q", serviceFilePath))
	}

	defer file.Close()

	serviceTemplate, err := template.New("systemdservice").Parse(serviceTemplate)
	if err != nil {
		return errors.Wrap(err, "Failed to parse systemd service template")
	}

	serviceTemplate.Execute(file, service)

	return nil
}

func getServiceName(projectName string) string {
	return fmt.Sprintf("%s.service", projectName)
}

const serviceTemplate = `[Unit]
Description={{ .ProjectName }} Service
After=network.service docker.service
Requires=docker.service

[Service]
Restart=always
RestartSec=10
TimeoutSec=300
WorkingDirectory={{ .ProjectDirectory }}
ExecStartPre=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" pull
ExecStart=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" up
ExecStop=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" stop
ExecStopPost=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" down

[Install]
WantedBy=docker.service
`

type serviceDefinition struct {
	ProjectName       string
	ProjectDirectory  string
	DockerComposeFile string
}
<commit_msg>Cleanup<commit_after>package main

import (
	"fmt"
	"html/template"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
)

type installer interface {
	Install(projectDirectory, dockerComposeFileName, projectName string) error
}

type systemdInstaller struct {
	systemdDirectory string
	commandExecutor  Executor
}

func (installer *systemdInstaller) Install(projectDirectory, dockerComposeFileName, projectName string) error {

	serviceName := getServiceName(projectName)
	serviceViewModel := serviceDefinition{
		ProjectName:       projectName,
		ProjectDirectory:  projectDirectory,
		DockerComposeFile: dockerComposeFileName,
	}

	if err := createSystemdService(serviceViewModel, installer.systemdDirectory); err != nil {
		return errors.Wrap(err, fmt.Sprintf("Failed to create a systemd service for project %q (Directory: %q, Docker Compose File: %q)", projectName, projectDirectory, dockerComposeFileName))
	}

	reloadError := installer.commandExecutor.Run("systemctl", "daemon-reload")
	if reloadError != nil {
		return errors.Wrap(reloadError, "Failed to reload the systemd configuration")
	}

	enableError := installer.commandExecutor.Run("systemctl", "enable", serviceName)
	if enableError != nil {
		return errors.Wrap(enableError, fmt.Sprintf("Failed to enable %q", serviceName))
	}

	startError := installer.commandExecutor.Run("systemctl", "start", serviceName)
	if startError != nil {
		return errors.Wrap(startError, fmt.Sprintf("Failed to start %q", serviceName))
	}

	return nil
}

func createSystemdService(service serviceDefinition, targetDirectory string) error {
	serviceFilePath := filepath.Join(targetDirectory, getServiceName(service.ProjectName))
	file, err := os.OpenFile(serviceFilePath, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0664)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("Failed to open the systemd service file: %q", serviceFilePath))
	}

	defer file.Close()

	serviceTemplate, err := template.New("systemdservice").Parse(serviceTemplate)
	if err != nil {
		return errors.Wrap(err, "Failed to parse systemd service template")
	}

	serviceTemplate.Execute(file, service)

	return nil
}

func getServiceName(projectName string) string {
	return fmt.Sprintf("%s.service", projectName)
}

const serviceTemplate = `[Unit]
Description={{ .ProjectName }} Service
After=docker.service
Requires=docker.service

[Service]
Type=notify
NotifyAccess=all
Restart=always
RestartSec=10
TimeoutSec=300
WorkingDirectory={{ .ProjectDirectory }}
ExecStartPre=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" pull
ExecStart=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" up
ExecStop=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" stop
ExecStopPost=/usr/bin/env docker-compose -p "{{ .ProjectName }}" -f "{{ .DockerComposeFile }}" down

[Install]
WantedBy=multi-user.target
`

type serviceDefinition struct {
	ProjectName       string
	ProjectDirectory  string
	DockerComposeFile string
}
<|endoftext|>


<commit_before>// Copyright 2015 Square 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 query

import (
	"errors"
	"fmt"
	"math"
	"sync/atomic"

	"github.com/square/metrics/api"
	"github.com/square/metrics/query/aggregate"
)

func init() {
	// Arithmetic operators
	MustRegister(MakeOperatorMetricFunction("+", func(x float64, y float64) float64 { return x + y }))
	MustRegister(MakeOperatorMetricFunction("-", func(x float64, y float64) float64 { return x - y }))
	MustRegister(MakeOperatorMetricFunction("*", func(x float64, y float64) float64 { return x * y }))
	MustRegister(MakeOperatorMetricFunction("/", func(x float64, y float64) float64 { return x / y }))
	// Aggregates
	MustRegister(MakeAggregateMetricFunction("aggregate.max", aggregate.AggregateMap["aggregate.max"]))
	MustRegister(MakeAggregateMetricFunction("aggregate.min", aggregate.AggregateMap["aggregate.min"]))
	MustRegister(MakeAggregateMetricFunction("aggregate.mean", aggregate.AggregateMap["aggregate.mean"]))
	MustRegister(MakeAggregateMetricFunction("aggregate.sum", aggregate.AggregateMap["aggregate.sum"]))
	// Transformations
	MustRegister(MakeTransformMetricFunction("transform.derivative", 0, transformTable["transform.derivative"]))
	MustRegister(MakeTransformMetricFunction("transform.integral", 0, transformTable["transform.integral"]))
	MustRegister(MakeTransformMetricFunction("transform.rate", 0, transformTable["transform.rate"]))
	MustRegister(MakeTransformMetricFunction("transform.cumulative", 0, transformTable["transform.cumulative"]))
	MustRegister(MakeTransformMetricFunction("transform.moving_average", 0, transformTable["transform.moving_average"]))
	MustRegister(MakeTransformMetricFunction("transform.default", 0, transformTable["transform.default"]))
	MustRegister(MakeTransformMetricFunction("transform.abs", 0, transformTable["transform.abs"]))
	// Timeshift
	MustRegister(TimeshiftFunction)
	// Filter
	meanSummary := func(xs []float64) float64 {
		count := 0
		sum := 0.0
		for _, x := range xs {
			if !math.IsNaN(x) {
				count++
				sum += x
			}
		}
		return sum / count
	}
	MustRegister(MakeFilterMetricFunction("filter.highest_mean", meanSummary, false))
	MustRegister(MakeFilterMetricFunction("filter.lowest_mean", meanSummary, true))
	maxSummary := func(xs []float64) float64 {
		high := math.NaN()
		for _, x := range xs {
			if math.IsNaN(high) || x > high {
				high = x
			}
		}
		return high
	}
	minSummary := func(xs []float64) float64 {
		low := math.NaN()
		for _, x := range xs {
			if math.IsNaN(low) || x < low {
				low = x
			}
		}
		return low
	}
	MustRegister(MakeFilterMetricFunction("filter.highest_max", maxSummary, false))
	MustRegister(MakeFilterMetricFunction("filter.lowest_max", maxSummary, true))
	MustRegister(MakeFilterMetricFunction("filter.highest_min", minSummary, false))
	MustRegister(MakeFilterMetricFunction("filter.lowest_min", minSummary, true))
}

// EvaluationContext is the central piece of logic, providing
// helper funcions & varaibles to evaluate a given piece of
// metrics query.
// * Contains Backend object, which can be used to fetch data
// from the backend system.s
// * Contains current timerange being queried for - this can be
// changed by say, application of time shift function.
type EvaluationContext struct {
	MultiBackend api.MultiBackend // Backend to fetch data from
	API          api.API          // Api to obtain metadata from
	Timerange    api.Timerange    // Timerange to fetch data from
	SampleMethod api.SampleMethod // SampleMethod to use when up/downsampling to match the requested resolution
	Predicate    api.Predicate    // Predicate to apply to TagSets prior to fetching
	FetchLimit   fetchCounter     // A limit on the number of fetches which may be performed
}

// fetchCounter is used to count the number of fetches remaining in a thread-safe manner.
type fetchCounter struct {
	count *int32
}

func NewFetchCounter(n int) fetchCounter {
	n32 := int32(n)
	return fetchCounter{
		count: &n32,
	}
}

// Consume decrements the internal counter and returns whether the result is at least 0.
// It does so in a threadsafe manner.
func (c fetchCounter) Consume(n int) bool {
	return atomic.AddInt32(c.count, -int32(n)) >= 0
}

// Expression is a piece of code, which can be evaluated in a given
// EvaluationContext. EvaluationContext must never be changed in an Evalute().
//
// The contract of Expressions is that leaf nodes must sample a resulting
// timeseries according to the resolution specified in its EvaluationContext's
// Timerange. Internal nodes may assume that results from evaluating child
// Expressions correspond to the timerange in the current EvaluationContext.
type Expression interface {
	// Evaluate the given expression.
	Evaluate(context EvaluationContext) (value, error)
}

// Implementations
// ===============

// Generates a Timeseries from the encapsulated scalar.
func (expr scalarExpression) Evaluate(context EvaluationContext) (value, error) {
	return scalarValue(expr.value), nil
}

func (expr stringExpression) Evaluate(context EvaluationContext) (value, error) {
	return stringValue(expr.value), nil
}

func (expr *metricFetchExpression) Evaluate(context EvaluationContext) (value, error) {
	// Merge predicates appropriately
	var predicate api.Predicate
	if context.Predicate == nil && expr.predicate == nil {
		predicate = api.TruePredicate
	} else if context.Predicate == nil {
		predicate = expr.predicate
	} else if expr.predicate == nil {
		predicate = context.Predicate
	} else {
		predicate = &andPredicate{[]api.Predicate{expr.predicate, context.Predicate}}
	}

	metricTagSets, err := context.API.GetAllTags(api.MetricKey(expr.metricName))
	if err != nil {
		return nil, err
	}
	filtered := applyPredicates(metricTagSets, predicate)

	ok := context.FetchLimit.Consume(len(filtered))

	if !ok {
		return nil, errors.New("fetch limit exceeded: too many series to fetch")
	}

	metrics := make([]api.TaggedMetric, len(filtered))
	for i := range metrics {
		metrics[i] = api.TaggedMetric{api.MetricKey(expr.metricName), filtered[i]}
	}

	serieslist, err := context.MultiBackend.FetchMultipleSeries(metrics, context.SampleMethod, context.Timerange, context.API)

	if err != nil {
		return nil, err
	}

	return seriesListValue(serieslist), nil
}

func (expr *functionExpression) Evaluate(context EvaluationContext) (value, error) {
	fun, ok := GetFunction(expr.functionName)
	if !ok {
		return nil, fmt.Errorf("no such function %s", expr.functionName)
	}

	return fun.Evaluate(context, expr.arguments, expr.groupBy)
}

// Auxiliary functions
// ===================

func applyPredicates(tagSets []api.TagSet, predicate api.Predicate) []api.TagSet {
	output := []api.TagSet{}
	for _, ts := range tagSets {
		if predicate.Apply(ts) {
			output = append(output, ts)
		}
	}
	return output
}

// evaluateExpressions evaluates all provided Expressions in the
// EvaluationContext. If any evaluations error, evaluateExpressions will
// propagate that error. The resulting SeriesLists will be in an order
// corresponding to the provided Expresesions.
func evaluateExpressions(context EvaluationContext, expressions []Expression) ([]value, error) {
	if len(expressions) == 0 {
		return []value{}, nil
	}
	results := make([]value, len(expressions))
	for i, expr := range expressions {
		result, err := expr.Evaluate(context)
		if err != nil {
			return nil, err
		}
		results[i] = result
	}
	return results, nil
}
<commit_msg>match types in division<commit_after>// Copyright 2015 Square 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 query

import (
	"errors"
	"fmt"
	"math"
	"sync/atomic"

	"github.com/square/metrics/api"
	"github.com/square/metrics/query/aggregate"
)

func init() {
	// Arithmetic operators
	MustRegister(MakeOperatorMetricFunction("+", func(x float64, y float64) float64 { return x + y }))
	MustRegister(MakeOperatorMetricFunction("-", func(x float64, y float64) float64 { return x - y }))
	MustRegister(MakeOperatorMetricFunction("*", func(x float64, y float64) float64 { return x * y }))
	MustRegister(MakeOperatorMetricFunction("/", func(x float64, y float64) float64 { return x / y }))
	// Aggregates
	MustRegister(MakeAggregateMetricFunction("aggregate.max", aggregate.AggregateMap["aggregate.max"]))
	MustRegister(MakeAggregateMetricFunction("aggregate.min", aggregate.AggregateMap["aggregate.min"]))
	MustRegister(MakeAggregateMetricFunction("aggregate.mean", aggregate.AggregateMap["aggregate.mean"]))
	MustRegister(MakeAggregateMetricFunction("aggregate.sum", aggregate.AggregateMap["aggregate.sum"]))
	// Transformations
	MustRegister(MakeTransformMetricFunction("transform.derivative", 0, transformTable["transform.derivative"]))
	MustRegister(MakeTransformMetricFunction("transform.integral", 0, transformTable["transform.integral"]))
	MustRegister(MakeTransformMetricFunction("transform.rate", 0, transformTable["transform.rate"]))
	MustRegister(MakeTransformMetricFunction("transform.cumulative", 0, transformTable["transform.cumulative"]))
	MustRegister(MakeTransformMetricFunction("transform.moving_average", 0, transformTable["transform.moving_average"]))
	MustRegister(MakeTransformMetricFunction("transform.default", 0, transformTable["transform.default"]))
	MustRegister(MakeTransformMetricFunction("transform.abs", 0, transformTable["transform.abs"]))
	// Timeshift
	MustRegister(TimeshiftFunction)
	// Filter
	meanSummary := func(xs []float64) float64 {
		count := 0
		sum := 0.0
		for _, x := range xs {
			if !math.IsNaN(x) {
				count++
				sum += x
			}
		}
		return sum / float64(count)
	}
	MustRegister(MakeFilterMetricFunction("filter.highest_mean", meanSummary, false))
	MustRegister(MakeFilterMetricFunction("filter.lowest_mean", meanSummary, true))
	maxSummary := func(xs []float64) float64 {
		high := math.NaN()
		for _, x := range xs {
			if math.IsNaN(high) || x > high {
				high = x
			}
		}
		return high
	}
	minSummary := func(xs []float64) float64 {
		low := math.NaN()
		for _, x := range xs {
			if math.IsNaN(low) || x < low {
				low = x
			}
		}
		return low
	}
	MustRegister(MakeFilterMetricFunction("filter.highest_max", maxSummary, false))
	MustRegister(MakeFilterMetricFunction("filter.lowest_max", maxSummary, true))
	MustRegister(MakeFilterMetricFunction("filter.highest_min", minSummary, false))
	MustRegister(MakeFilterMetricFunction("filter.lowest_min", minSummary, true))
}

// EvaluationContext is the central piece of logic, providing
// helper funcions & varaibles to evaluate a given piece of
// metrics query.
// * Contains Backend object, which can be used to fetch data
// from the backend system.s
// * Contains current timerange being queried for - this can be
// changed by say, application of time shift function.
type EvaluationContext struct {
	MultiBackend api.MultiBackend // Backend to fetch data from
	API          api.API          // Api to obtain metadata from
	Timerange    api.Timerange    // Timerange to fetch data from
	SampleMethod api.SampleMethod // SampleMethod to use when up/downsampling to match the requested resolution
	Predicate    api.Predicate    // Predicate to apply to TagSets prior to fetching
	FetchLimit   fetchCounter     // A limit on the number of fetches which may be performed
}

// fetchCounter is used to count the number of fetches remaining in a thread-safe manner.
type fetchCounter struct {
	count *int32
}

func NewFetchCounter(n int) fetchCounter {
	n32 := int32(n)
	return fetchCounter{
		count: &n32,
	}
}

// Consume decrements the internal counter and returns whether the result is at least 0.
// It does so in a threadsafe manner.
func (c fetchCounter) Consume(n int) bool {
	return atomic.AddInt32(c.count, -int32(n)) >= 0
}

// Expression is a piece of code, which can be evaluated in a given
// EvaluationContext. EvaluationContext must never be changed in an Evalute().
//
// The contract of Expressions is that leaf nodes must sample a resulting
// timeseries according to the resolution specified in its EvaluationContext's
// Timerange. Internal nodes may assume that results from evaluating child
// Expressions correspond to the timerange in the current EvaluationContext.
type Expression interface {
	// Evaluate the given expression.
	Evaluate(context EvaluationContext) (value, error)
}

// Implementations
// ===============

// Generates a Timeseries from the encapsulated scalar.
func (expr scalarExpression) Evaluate(context EvaluationContext) (value, error) {
	return scalarValue(expr.value), nil
}

func (expr stringExpression) Evaluate(context EvaluationContext) (value, error) {
	return stringValue(expr.value), nil
}

func (expr *metricFetchExpression) Evaluate(context EvaluationContext) (value, error) {
	// Merge predicates appropriately
	var predicate api.Predicate
	if context.Predicate == nil && expr.predicate == nil {
		predicate = api.TruePredicate
	} else if context.Predicate == nil {
		predicate = expr.predicate
	} else if expr.predicate == nil {
		predicate = context.Predicate
	} else {
		predicate = &andPredicate{[]api.Predicate{expr.predicate, context.Predicate}}
	}

	metricTagSets, err := context.API.GetAllTags(api.MetricKey(expr.metricName))
	if err != nil {
		return nil, err
	}
	filtered := applyPredicates(metricTagSets, predicate)

	ok := context.FetchLimit.Consume(len(filtered))

	if !ok {
		return nil, errors.New("fetch limit exceeded: too many series to fetch")
	}

	metrics := make([]api.TaggedMetric, len(filtered))
	for i := range metrics {
		metrics[i] = api.TaggedMetric{api.MetricKey(expr.metricName), filtered[i]}
	}

	serieslist, err := context.MultiBackend.FetchMultipleSeries(metrics, context.SampleMethod, context.Timerange, context.API)

	if err != nil {
		return nil, err
	}

	return seriesListValue(serieslist), nil
}

func (expr *functionExpression) Evaluate(context EvaluationContext) (value, error) {
	fun, ok := GetFunction(expr.functionName)
	if !ok {
		return nil, fmt.Errorf("no such function %s", expr.functionName)
	}

	return fun.Evaluate(context, expr.arguments, expr.groupBy)
}

// Auxiliary functions
// ===================

func applyPredicates(tagSets []api.TagSet, predicate api.Predicate) []api.TagSet {
	output := []api.TagSet{}
	for _, ts := range tagSets {
		if predicate.Apply(ts) {
			output = append(output, ts)
		}
	}
	return output
}

// evaluateExpressions evaluates all provided Expressions in the
// EvaluationContext. If any evaluations error, evaluateExpressions will
// propagate that error. The resulting SeriesLists will be in an order
// corresponding to the provided Expresesions.
func evaluateExpressions(context EvaluationContext, expressions []Expression) ([]value, error) {
	if len(expressions) == 0 {
		return []value{}, nil
	}
	results := make([]value, len(expressions))
	for i, expr := range expressions {
		result, err := expr.Evaluate(context)
		if err != nil {
			return nil, err
		}
		results[i] = result
	}
	return results, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package kops

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"

	"k8s.io/kops/tests/e2e/pkg/util"
)

// DownloadKops will download the kops binary from the version marker URL
// Returning the path to the local kops binary and the URL to use for KOPS_BASE_URL
// Example markerURL: https://storage.googleapis.com/kops-ci/bin/latest-ci-updown-green.txt
func DownloadKops(markerURL, downloadPath string) (string, string, error) {
	var b bytes.Buffer
	if err := util.HTTPGETWithHeaders(markerURL, nil, &b); err != nil {
		return "", "", err
	}
	kopsBaseURL := b.String()

	var kopsFile *os.File
	if downloadPath == "" {
		tmp, err := ioutil.TempFile("", "kops")
		if err != nil {
			return "", "", err
		}
		kopsFile = tmp
	} else {
		tmp, err := os.Create(downloadPath)
		if err != nil {
			return "", "", err
		}
		kopsFile = tmp
	}

	kopsURL := fmt.Sprintf("%v/%v/%v/kops", kopsBaseURL, runtime.GOOS, runtime.GOARCH)
	if err := util.HTTPGETWithHeaders(kopsURL, nil, kopsFile); err != nil {
		return "", "", err
	}
	if err := kopsFile.Close(); err != nil {
		return "", "", err
	}
	if err := os.Chmod(kopsFile.Name(), 0755); err != nil {
		return "", "", err
	}
	return kopsFile.Name(), kopsBaseURL, nil
}
<commit_msg>Trim space on kops version markers<commit_after>/*
Copyright 2021 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package kops

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"strings"

	"k8s.io/kops/tests/e2e/pkg/util"
)

// DownloadKops will download the kops binary from the version marker URL
// Returning the path to the local kops binary and the URL to use for KOPS_BASE_URL
// Example markerURL: https://storage.googleapis.com/kops-ci/bin/latest-ci-updown-green.txt
func DownloadKops(markerURL, downloadPath string) (string, string, error) {
	var b bytes.Buffer
	if err := util.HTTPGETWithHeaders(markerURL, nil, &b); err != nil {
		return "", "", err
	}
	kopsBaseURL := strings.TrimSpace(b.String())

	var kopsFile *os.File
	if downloadPath == "" {
		tmp, err := ioutil.TempFile("", "kops")
		if err != nil {
			return "", "", err
		}
		kopsFile = tmp
	} else {
		tmp, err := os.Create(downloadPath)
		if err != nil {
			return "", "", err
		}
		kopsFile = tmp
	}

	kopsURL := fmt.Sprintf("%v/%v/%v/kops", kopsBaseURL, runtime.GOOS, runtime.GOARCH)
	if err := util.HTTPGETWithHeaders(kopsURL, nil, kopsFile); err != nil {
		return "", "", err
	}
	if err := kopsFile.Close(); err != nil {
		return "", "", err
	}
	if err := os.Chmod(kopsFile.Name(), 0755); err != nil {
		return "", "", err
	}
	return kopsFile.Name(), kopsBaseURL, nil
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"net"
	"net/http"
	"net/url"
	"sync"
	"testing"
	"time"

	"github.com/rafrombrc/gomock/gomock"
	"golang.org/x/net/websocket"
)

// newMockReadyNotifier wraps l in a mockReadyNotifier.
func newMockReadyNotifier(l Locator) *mockReadyNotifier {
	return &mockReadyNotifier{Locator: l, readyChan: make(chan bool)}
}

// mockReadyNotifier implements the Locator and ReadyNotifier interfaces. This
// is used to test blocking incoming WebSocket connections until the locator is
// ready.
type mockReadyNotifier struct {
	Locator
	readyChan chan bool
}

// SignalReady signals that the locator is ready.
func (m *mockReadyNotifier) SignalReady() { close(m.readyChan) }

// ReadyNotify returns a channel that is closed by SignalReady.
func (m *mockReadyNotifier) ReadyNotify() <-chan bool { return m.readyChan }

func TestLocatorReadyNotify(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	uaid := "fce61180716a40ed8e79bf5ff0ba34bc"

	mckLogger := NewMockLogger(mockCtrl)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
		gomock.Any()).AnyTimes()

	var (
		// routerPipes maps fake peer addresses to their respective pipes. Used
		// by dialRouter to connect to peers.
		routerPipes = make(map[netAddr]*pipeListener)

		// contacts is a list of peer URLs for the locator.
		contacts []string
	)

	// Fake listener for the sender's router, used to test self-routing.
	sndRouterAddr := netAddr{"tcp", "snd-router.example.com:3000"}
	sndRouterPipe := newPipeListener()
	defer sndRouterPipe.Close()
	routerPipes[sndRouterAddr] = sndRouterPipe
	contacts = append(contacts, "http://snd-router.example.com:3000")

	// Fake listener for the receiver's router, used to test routing updates
	// to different hosts.
	recvRouterAddr := netAddr{"tcp", "recv-router.example.com:3000"}
	recvRouterPipe := newPipeListener()
	defer recvRouterPipe.Close()
	routerPipes[recvRouterAddr] = recvRouterPipe
	contacts = append(contacts, "http://recv-router.example.com:3000")

	// Fake listener for the receiver's WebSocket handler, used to accept a
	// WebSocket client connection.
	socketHandlerPipe := newPipeListener()
	defer socketHandlerPipe.Close()

	// Fake locator.
	mckLocator := NewMockLocator(mockCtrl)
	mckLocator.EXPECT().Contacts(uaid).Return(contacts, nil).Times(2)

	// Fake dialer to connect to each peer's routing listener.
	dialRouter := func(network, address string) (net.Conn, error) {
		if pipe, ok := routerPipes[netAddr{network, address}]; ok {
			return pipe.Dial(network, address)
		}
		return nil, &netErr{temporary: false, timeout: false}
	}
	// Configures a fake router for the app.
	setRouter := func(app *Application, listener net.Listener) {
		r := NewBroadcastRouter()
		r.setApp(app)
		r.setClientOptions(10, 3*time.Second, 3*time.Second) // Defaults.
		r.setClientTransport(&http.Transport{Dial: dialRouter})
		r.listenWithConfig(listenerConfig{listener: listener})
		r.maxDataLen = 4096
		r.server = newServeWaiter(&http.Server{Handler: r.ServeMux()})
		app.SetRouter(r)
	}

	// sndApp is the server broadcasting the update. The locator returns the
	// addresses of the sender and receiver to test self-routing.
	sndApp := NewApplication()
	sndApp.SetLogger(mckLogger)
	sndStat := NewMockStatistician(mockCtrl)
	sndApp.SetMetrics(sndStat)
	sndApp.SetLocator(mckLocator)
	// Set up a fake router for the sender.
	setRouter(sndApp, sndRouterPipe)

	// recvApp is the server receiving the update.
	recvApp := NewApplication()
	recvApp.SetLogger(mckLogger)
	recvStat := NewMockStatistician(mockCtrl)
	recvApp.SetMetrics(recvStat)
	recvStore := NewMockStore(mockCtrl)
	recvApp.SetStore(recvStore)
	// Wrap the fake locator in a type that implements ReadyNotifier.
	recvLocator := newMockReadyNotifier(mckLocator)
	recvApp.SetLocator(recvLocator)
	// Set up a fake WebSocket handler for the receiver.
	recvSocketHandler := NewSocketHandler()
	recvSocketHandler.setApp(recvApp)
	recvSocketHandler.listenWithConfig(listenerConfig{
		listener: socketHandlerPipe})
	recvSocketHandler.server = newServeWaiter(&http.Server{Handler: recvSocketHandler.ServeMux()})
	recvApp.SetSocketHandler(recvSocketHandler)
	// Set up a fake router for the receiver.
	setRouter(recvApp, recvRouterPipe)

	chid := "2b7c5c27d6224bfeaf1c158c3c57fca3"
	version := int64(2)
	data := "I'm a little teapot, short and stout."

	var wg sync.WaitGroup // Waits for the client to close.
	wg.Add(1)
	dialChan := make(chan bool) // Signals when the client connects.
	timeout := closeAfter(2 * time.Second)

	go func() {
		defer wg.Done()
		origin := &url.URL{Scheme: "ws", Host: "recv-conn.example.com"}
		ws, err := dialSocketListener(socketHandlerPipe, &websocket.Config{
			Location: origin,
			Origin:   origin,
			Version:  websocket.ProtocolVersionHybi13,
		})
		if err != nil {
			t.Errorf("Error dialing host: %s", err)
			return
		}
		defer ws.Close()
		err = websocket.JSON.Send(ws, struct {
			Type       string   `json:"messageType"`
			DeviceID   string   `json:"uaid"`
			ChannelIDs []string `json:"channelIDs"`
		}{"hello", uaid, []string{}})
		if err != nil {
			t.Errorf("Error writing handshake request: %s", err)
			return
		}
		helloReply := new(HelloReply)
		if err = websocket.JSON.Receive(ws, helloReply); err != nil {
			t.Errorf("Error reading handshake reply: %s", err)
			return
		}
		select {
		case dialChan <- true:
		case <-timeout:
			t.Errorf("Timed out waiting for router")
			return
		}
		flushReply := new(FlushReply)
		if err = websocket.JSON.Receive(ws, flushReply); err != nil {
			t.Errorf("Error reading routed update: %s", err)
			return
		}
		ok := false
		expected := Update{chid, uint64(version), data}
		for _, update := range flushReply.Updates {
			if ok = update == expected; ok {
				break
			}
		}
		if !ok {
			t.Errorf("Missing update %#v in %#v", expected, flushReply.Updates)
			return
		}
	}()

	// Start the handlers.
	errChan := make(chan error, 3)
	go sndApp.Router().Start(errChan)
	go recvApp.SocketHandler().Start(errChan)
	go recvApp.Router().Start(errChan)

	// First and second routing attempts to self.
	sndStat.EXPECT().Increment("updates.routed.unknown").Times(2)

	// Initial routing attempt to peer.
	recvStat.EXPECT().Increment("updates.routed.unknown")
	// Client connects to peer.
	recvStat.EXPECT().Increment("client.socket.connect")
	recvStore.EXPECT().CanStore(0).Return(true)
	recvStat.EXPECT().Increment("updates.client.hello")
	recvStore.EXPECT().FetchAll(uaid, gomock.Any()).Return(nil, nil, nil)
	recvStat.EXPECT().Timer("client.flush", gomock.Any())
	// Second routing attempt to peer.
	recvStat.EXPECT().Increment("updates.routed.incoming")
	recvStat.EXPECT().Increment("updates.sent")
	recvStat.EXPECT().Timer("client.flush", gomock.Any())
	recvStat.EXPECT().Increment("updates.routed.received")
	recvStat.EXPECT().Timer("client.socket.lifespan", gomock.Any())
	recvStat.EXPECT().Increment("client.socket.disconnect")

	// Initial routing attempt should fail; the WebSocket listener shouldn't
	// accept client connections before the locator is ready.
	delivered, err := sndApp.Router().Route(nil, uaid, chid, version, timeNow(),
		"disconnected", data)
	if err != nil {
		t.Errorf("Error routing to disconnected client: %s", err)
	} else if delivered {
		t.Error("Should not route to disconnected client")
	}
	// Signal the locator is ready, then wait for the client to connect.
	recvLocator.SignalReady()
	select {
	case <-dialChan:
	case <-time.After(5 * time.Second):
		t.Fatalf("Timed out waiting for the client to connect")
	}
	// Routing should succeed once the client is connected.
	delivered, err = sndApp.Router().Route(nil, uaid, chid, version, timeNow(),
		"connected", data)
	if err != nil {
		t.Errorf("Error routing to connected client: %s", err)
	} else if !delivered {
		t.Error("Should route to connected client")
	}

	mckLocator.EXPECT().Close().Times(2)
	if err := recvApp.Close(); err != nil {
		t.Errorf("Error closing peer: %s", err)
	}
	wg.Wait()
	if err := sndApp.Close(); err != nil {
		t.Errorf("Error closing self: %s", err)
	}
	// Wait for the handlers to stop.
	for i := 0; i < 3; i++ {
		<-errChan
	}
}
<commit_msg>Add missing `mckStore.Close()` expectation.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"net"
	"net/http"
	"net/url"
	"sync"
	"testing"
	"time"

	"github.com/rafrombrc/gomock/gomock"
	"golang.org/x/net/websocket"
)

// newMockReadyNotifier wraps l in a mockReadyNotifier.
func newMockReadyNotifier(l Locator) *mockReadyNotifier {
	return &mockReadyNotifier{Locator: l, readyChan: make(chan bool)}
}

// mockReadyNotifier implements the Locator and ReadyNotifier interfaces. This
// is used to test blocking incoming WebSocket connections until the locator is
// ready.
type mockReadyNotifier struct {
	Locator
	readyChan chan bool
}

// SignalReady signals that the locator is ready.
func (m *mockReadyNotifier) SignalReady() { close(m.readyChan) }

// ReadyNotify returns a channel that is closed by SignalReady.
func (m *mockReadyNotifier) ReadyNotify() <-chan bool { return m.readyChan }

func TestLocatorReadyNotify(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	uaid := "fce61180716a40ed8e79bf5ff0ba34bc"

	mckLogger := NewMockLogger(mockCtrl)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
		gomock.Any()).AnyTimes()

	var (
		// routerPipes maps fake peer addresses to their respective pipes. Used
		// by dialRouter to connect to peers.
		routerPipes = make(map[netAddr]*pipeListener)

		// contacts is a list of peer URLs for the locator.
		contacts []string
	)

	// Fake listener for the sender's router, used to test self-routing.
	sndRouterAddr := netAddr{"tcp", "snd-router.example.com:3000"}
	sndRouterPipe := newPipeListener()
	defer sndRouterPipe.Close()
	routerPipes[sndRouterAddr] = sndRouterPipe
	contacts = append(contacts, "http://snd-router.example.com:3000")

	// Fake listener for the receiver's router, used to test routing updates
	// to different hosts.
	recvRouterAddr := netAddr{"tcp", "recv-router.example.com:3000"}
	recvRouterPipe := newPipeListener()
	defer recvRouterPipe.Close()
	routerPipes[recvRouterAddr] = recvRouterPipe
	contacts = append(contacts, "http://recv-router.example.com:3000")

	// Fake listener for the receiver's WebSocket handler, used to accept a
	// WebSocket client connection.
	socketHandlerPipe := newPipeListener()
	defer socketHandlerPipe.Close()

	// Fake locator.
	mckLocator := NewMockLocator(mockCtrl)
	mckLocator.EXPECT().Contacts(uaid).Return(contacts, nil).Times(2)

	// Fake dialer to connect to each peer's routing listener.
	dialRouter := func(network, address string) (net.Conn, error) {
		if pipe, ok := routerPipes[netAddr{network, address}]; ok {
			return pipe.Dial(network, address)
		}
		return nil, &netErr{temporary: false, timeout: false}
	}
	// Configures a fake router for the app.
	setRouter := func(app *Application, listener net.Listener) {
		r := NewBroadcastRouter()
		r.setApp(app)
		r.setClientOptions(10, 3*time.Second, 3*time.Second) // Defaults.
		r.setClientTransport(&http.Transport{Dial: dialRouter})
		r.listenWithConfig(listenerConfig{listener: listener})
		r.maxDataLen = 4096
		r.server = newServeWaiter(&http.Server{Handler: r.ServeMux()})
		app.SetRouter(r)
	}

	// sndApp is the server broadcasting the update. The locator returns the
	// addresses of the sender and receiver to test self-routing.
	sndApp := NewApplication()
	sndApp.SetLogger(mckLogger)
	sndStat := NewMockStatistician(mockCtrl)
	sndApp.SetMetrics(sndStat)
	sndStore := NewMockStore(mockCtrl)
	sndApp.SetStore(sndStore)
	sndApp.SetLocator(mckLocator)
	// Set up a fake router for the sender.
	setRouter(sndApp, sndRouterPipe)

	// recvApp is the server receiving the update.
	recvApp := NewApplication()
	recvApp.SetLogger(mckLogger)
	recvStat := NewMockStatistician(mockCtrl)
	recvApp.SetMetrics(recvStat)
	recvStore := NewMockStore(mockCtrl)
	recvApp.SetStore(recvStore)
	// Wrap the fake locator in a type that implements ReadyNotifier.
	recvLocator := newMockReadyNotifier(mckLocator)
	recvApp.SetLocator(recvLocator)
	// Set up a fake WebSocket handler for the receiver.
	recvSocketHandler := NewSocketHandler()
	recvSocketHandler.setApp(recvApp)
	recvSocketHandler.listenWithConfig(listenerConfig{
		listener: socketHandlerPipe})
	recvSocketHandler.server = newServeWaiter(&http.Server{Handler: recvSocketHandler.ServeMux()})
	recvApp.SetSocketHandler(recvSocketHandler)
	// Set up a fake router for the receiver.
	setRouter(recvApp, recvRouterPipe)

	chid := "2b7c5c27d6224bfeaf1c158c3c57fca3"
	version := int64(2)
	data := "I'm a little teapot, short and stout."

	var wg sync.WaitGroup // Waits for the client to close.
	wg.Add(1)
	dialChan := make(chan bool) // Signals when the client connects.
	timeout := closeAfter(2 * time.Second)

	go func() {
		defer wg.Done()
		origin := &url.URL{Scheme: "ws", Host: "recv-conn.example.com"}
		ws, err := dialSocketListener(socketHandlerPipe, &websocket.Config{
			Location: origin,
			Origin:   origin,
			Version:  websocket.ProtocolVersionHybi13,
		})
		if err != nil {
			t.Errorf("Error dialing host: %s", err)
			return
		}
		defer ws.Close()
		err = websocket.JSON.Send(ws, struct {
			Type       string   `json:"messageType"`
			DeviceID   string   `json:"uaid"`
			ChannelIDs []string `json:"channelIDs"`
		}{"hello", uaid, []string{}})
		if err != nil {
			t.Errorf("Error writing handshake request: %s", err)
			return
		}
		helloReply := new(HelloReply)
		if err = websocket.JSON.Receive(ws, helloReply); err != nil {
			t.Errorf("Error reading handshake reply: %s", err)
			return
		}
		select {
		case dialChan <- true:
		case <-timeout:
			t.Errorf("Timed out waiting for router")
			return
		}
		flushReply := new(FlushReply)
		if err = websocket.JSON.Receive(ws, flushReply); err != nil {
			t.Errorf("Error reading routed update: %s", err)
			return
		}
		ok := false
		expected := Update{chid, uint64(version), data}
		for _, update := range flushReply.Updates {
			if ok = update == expected; ok {
				break
			}
		}
		if !ok {
			t.Errorf("Missing update %#v in %#v", expected, flushReply.Updates)
			return
		}
	}()

	// Start the handlers.
	errChan := make(chan error, 3)
	go sndApp.Router().Start(errChan)
	go recvApp.SocketHandler().Start(errChan)
	go recvApp.Router().Start(errChan)

	// First and second routing attempts to self.
	sndStat.EXPECT().Increment("updates.routed.unknown").Times(2)

	// Initial routing attempt to peer.
	recvStat.EXPECT().Increment("updates.routed.unknown")
	// Client connects to peer.
	recvStat.EXPECT().Increment("client.socket.connect")
	recvStore.EXPECT().CanStore(0).Return(true)
	recvStat.EXPECT().Increment("updates.client.hello")
	recvStore.EXPECT().FetchAll(uaid, gomock.Any()).Return(nil, nil, nil)
	recvStat.EXPECT().Timer("client.flush", gomock.Any())
	// Second routing attempt to peer.
	recvStat.EXPECT().Increment("updates.routed.incoming")
	recvStat.EXPECT().Increment("updates.sent")
	recvStat.EXPECT().Timer("client.flush", gomock.Any())
	recvStat.EXPECT().Increment("updates.routed.received")
	recvStat.EXPECT().Timer("client.socket.lifespan", gomock.Any())
	recvStat.EXPECT().Increment("client.socket.disconnect")

	// Initial routing attempt should fail; the WebSocket listener shouldn't
	// accept client connections before the locator is ready.
	delivered, err := sndApp.Router().Route(nil, uaid, chid, version, timeNow(),
		"disconnected", data)
	if err != nil {
		t.Errorf("Error routing to disconnected client: %s", err)
	} else if delivered {
		t.Error("Should not route to disconnected client")
	}
	// Signal the locator is ready, then wait for the client to connect.
	recvLocator.SignalReady()
	select {
	case <-dialChan:
	case <-time.After(5 * time.Second):
		t.Fatalf("Timed out waiting for the client to connect")
	}
	// Routing should succeed once the client is connected.
	delivered, err = sndApp.Router().Route(nil, uaid, chid, version, timeNow(),
		"connected", data)
	if err != nil {
		t.Errorf("Error routing to connected client: %s", err)
	} else if !delivered {
		t.Error("Should route to connected client")
	}

	gomock.InOrder(
		mckLocator.EXPECT().Close(),
		recvStore.EXPECT().Close(),
	)
	if err := recvApp.Close(); err != nil {
		t.Errorf("Error closing peer: %s", err)
	}
	wg.Wait()
	gomock.InOrder(
		mckLocator.EXPECT().Close(),
		sndStore.EXPECT().Close(),
	)
	if err := sndApp.Close(); err != nil {
		t.Errorf("Error closing self: %s", err)
	}
	// Wait for the handlers to stop.
	for i := 0; i < 3; i++ {
		<-errChan
	}
}
<|endoftext|>