<commit_before>package rmq

import (
	"errors"
	"strings"
	"sync"
	"time"
)

//TestRedisClient is a mock for redis
type TestRedisClient struct {
	store sync.Map
	ttl   sync.Map
}

var lock sync.Mutex

//NewTestRedisClient returns a NewTestRedisClient
func NewTestRedisClient() *TestRedisClient {
	return &TestRedisClient{}
}

// Set sets key to hold the string value.
// If key already holds a value, it is overwritten, regardless of its type.
// Any previous time to live associated with the key is discarded on successful SET operation.
func (client *TestRedisClient) Set(key string, value string, expiration time.Duration) error {

	lock.Lock()
	defer lock.Unlock()

	client.store.Store(key, value)
	//Delete any previous time to live associated with the key
	client.ttl.Delete(key)

	//0.0 expiration means that the value won't expire
	if expiration.Seconds() != 0.0 {
		//Store the unix time at which we should delete this
		client.ttl.Store(key, time.Now().Add(expiration).Unix())
	}

	return nil
}

// Get the value of key.
// If the key does not exist or isn't a string
// the special value nil is returned.
func (client *TestRedisClient) Get(key string) (string, error) {

	value, found := client.store.Load(key)

	if found {
		if stringValue, casted := value.(string); casted {
			return stringValue, nil
		}
	}

	return "nil", nil
}

//Del removes the specified key. A key is ignored if it does not exist.
func (client *TestRedisClient) Del(key string) (affected int64, err error) {

	_, found := client.store.Load(key)
	client.store.Delete(key)
	client.ttl.Delete(key)

	if found {
		return 1, nil
	}
	return 0, nil

}

// TTL returns the remaining time to live of a key that has a timeout.
// This introspection capability allows a Redis client to check how many seconds a given key will continue to be part of the dataset.
// In Redis 2.6 or older the command returns -1 if the key does not exist or if the key exist but has no associated expire.
// Starting with Redis 2.8 the return value in case of error changed:
// The command returns -2 if the key does not exist.
// The command returns -1 if the key exists but has no associated expire.
func (client *TestRedisClient) TTL(key string) (ttl time.Duration, err error) {

	//Lookup the expiration map
	expiration, found := client.ttl.Load(key)

	//Found an expiration time
	if found {

		//It was there, but it expired; removing it now
		if expiration.(int64) < time.Now().Unix() {
			client.ttl.Delete(key)
			return -2, nil
		}

		ttl = time.Duration(expiration.(int64) - time.Now().Unix())
		return ttl, nil
	}

	//Lookup the store in case this key exists but don't have an expiration
	//date
	_, found = client.store.Load(key)

	//The key was in store but didn't have an expiration associated
	//to it.
	if found {
		return -1, nil
	}

	return -2, nil
}

// LPush inserts the specified value at the head of the list stored at key.
// If key does not exist, it is created as empty list before performing the push operations.
// When key holds a value that is not a list, an error is returned.
// It is possible to push multiple elements using a single command call just specifying multiple arguments
// at the end of the command. Elements are inserted one after the other to the head of the list,
// from the leftmost element to the rightmost element.
func (client *TestRedisClient) LPush(key string, value ...string) (total int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	list, err := client.findList(key)

	if err != nil {
		return 0, nil
	}

	client.storeList(key, append(value, list...))
	return int64(len(list)) + 1, nil
}

//LLen returns the length of the list stored at key.
//If key does not exist, it is interpreted as an empty list and 0 is returned.
//An error is returned when the value stored at key is not a list.
func (client *TestRedisClient) LLen(key string) (affected int64, err error) {
	list, err := client.findList(key)

	if err != nil {
		return 0, nil
	}
	return int64(len(list)), nil
}

// LRem removes the first count occurrences of elements equal to
// value from the list stored at key. The count argument influences
// the operation in the following ways:
// count > 0: Remove elements equal to value moving from head to tail.
// count < 0: Remove elements equal to value moving from tail to head.
// count = 0: Remove all elements equal to value. For example,
// LREM list -2 "hello" will remove the last two occurrences of "hello" in
// the list stored at list. Note that non-existing keys are treated like empty
// lists, so when key does not exist, the command will always return 0.
func (client *TestRedisClient) LRem(key string, count int64, value string) (affected int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	list, err := client.findList(key)

	//Wasn't a list, or is empty
	if err != nil || len(list) == 0 {
		return 0, nil
	}

	//Create a list that have the capacity to store
	//the old one
	//This will be much more performant in case of
	//very long list
	newList := make([]string, 0, len(list))

	if err != nil {
		return 0, nil
	}

	//left to right removal of count elements
	if count >= 0 {

		//All the elements are to be removed.
		//Set count to max possible elements
		if count == 0 {
			count = int64(len(list))
		}
		//left to right traversal
		for index := 0; index < len(list); index++ {

			//isn't what we look for or we found enough element already
			if strings.Compare(list[index], value) != 0 || affected > count {
				newList = append(newList, list[index])
			} else {
				affected++
			}
		}
		//right to left removal of count elements
	} else if count < 0 {

		//right to left traversal
		for index := len(list) - 1; index >= 0; index-- {

			//isn't what we look for or we found enough element already
			if strings.Compare(list[index], value) != 0 || affected > count {
				//prepend instead of append to keep the order
				newList = append([]string{list[index]}, newList...)
			} else {
				affected++
			}
		}
	}

	//store the updated list
	client.storeList(key, newList)

	return affected, nil
}

// LTrim trims an existing list so that it will contain only the specified range of elements specified.
// Both start and stop are zero-based indexes, where 0 is the first element of the list (the head),
// 1 the next element and so on. For example: LTRIM foobar 0 2 will modify the list stored
// at foobar so that only the first three elements of the list will remain.
// start and end can also be negative numbers indicating offsets from the end of the list,
// where -1 is the last element of the list, -2 the penultimate element and so on.
// Out of range indexes will not produce an error: if start is larger than the end of the list,
// or start > end, the result will be an empty list (which causes key to be removed).
// If end is larger than the end of the list, Redis will treat it like the last element of the list
func (client *TestRedisClient) LTrim(key string, start, stop int64) error {

	lock.Lock()
	defer lock.Unlock()

	list, err := client.findList(key)

	//Wasn't a list, or is empty
	if err != nil || len(list) == 0 {
		return nil
	}

	if start < 0 {
		start += int64(len(list))
	}
	if stop < 0 {
		stop += int64(len(list))
	}

	//invalid values cause the remove of the key
	if start > stop {
		client.store.Delete(key)
		return nil
	}

	client.storeList(key, list[start:stop])
	return nil
}

// RPopLPush atomically returns and removes the last element (tail) of the list stored at source,
// and pushes the element at the first element (head) of the list stored at destination.
// For example: consider source holding the list a,b,c, and destination holding the list x,y,z.
// Executing RPOPLPUSH results in source holding a,b and destination holding c,x,y,z.
// If source does not exist, the value nil is returned and no operation is performed.
// If source and destination are the same, the operation is equivalent to removing the
// last element from the list and pushing it as first element of the list,
// so it can be considered as a list rotation command.
func (client *TestRedisClient) RPopLPush(source, destination string) (value string, err error) {

	lock.Lock()
	defer lock.Unlock()

	sourceList, sourceErr := client.findList(source)
	destList, destErr := client.findList(destination)

	//One of the two isn't a list
	if sourceErr != nil || destErr != nil {
		return "", ErrorNotFound
	}
	//we have nothing to move
	if len(sourceList) == 0 {
		return "", ErrorNotFound
	}

	//Remove the last element of source (tail)
	client.storeList(source, sourceList[0:len(sourceList)-1])
	//Put the last element of source (tail) and prepend it to dest
	client.storeList(destination, append([]string{sourceList[len(sourceList)-1]}, destList...))

	return sourceList[len(sourceList)-1], nil
}

// SAdd adds the specified members to the set stored at key.
// Specified members that are already a member of this set are ignored.
// If key does not exist, a new set is created before adding the specified members.
// An error is returned when the value stored at key is not a set.
func (client *TestRedisClient) SAdd(key, value string) (total int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	set, err := client.findSet(key)
	if err != nil {
		return 0, err
	}

	set[value] = struct{}{}
	client.storeSet(key, set)
	return int64(len(set)), nil
}

// SMembers returns all the members of the set value stored at key.
// This has the same effect as running SINTER with one argument key.
func (client *TestRedisClient) SMembers(key string) (members []string, err error) {
	set, err := client.findSet(key)
	if err != nil {
		return members, nil
	}

	members = make([]string, 0, len(set))
	for k := range set {
		members = append(members, k)
	}

	return members, nil
}

// SRem removes the specified members from the set stored at key.
// Specified members that are not a member of this set are ignored.
// If key does not exist, it is treated as an empty set and this command returns 0.
// An error is returned when the value stored at key is not a set.
func (client *TestRedisClient) SRem(key, value string) (affected int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	set, err := client.findSet(key)
	if err != nil || len(set) == 0 {
		return 0, nil
	}

	if _, found := set[value]; found != false {
		delete(set, value)
		return 1, nil
	}

	return 0, nil
}

// FlushDb delete all the keys of the currently selected DB. This command never fails.
func (client *TestRedisClient) FlushDb() error {
	client.store = *new(sync.Map)
	client.ttl = *new(sync.Map)
	return nil
}

//storeSet stores a set
func (client *TestRedisClient) storeSet(key string, set map[string]struct{}) {
	client.store.Store(key, set)
}

//findSet finds a set
func (client *TestRedisClient) findSet(key string) (map[string]struct{}, error) {
	//Lookup the store for the list
	storedValue, found := client.store.Load(key)
	if found {
		//list are stored as pointer to []string
		set, casted := storedValue.(map[string]struct{})

		if casted {
			return set, nil
		}

		return nil, errors.New("Stored value wasn't a set")
	}

	//return an empty set if not found
	return make(map[string]struct{}), nil
}

//storeList is an helper function so others don't have to deal with pointers
func (client *TestRedisClient) storeList(key string, list []string) {
	client.store.Store(key, &list)
}

//findList returns the list stored at key.
//if key doesn't exist, an empty list is returned
//an error is returned when the value at key isn't a list
func (client *TestRedisClient) findList(key string) ([]string, error) {
	//Lookup the store for the list
	storedValue, found := client.store.Load(key)
	if found {

		//list are stored as pointer to []string
		list, casted := storedValue.(*[]string)

		//Successful cass from interface{} to *[]string
		//Preprend the new key
		if casted {

			//This mock use sync.Map to be thread safe.
			//sync.Map only accepts interface{} as values and
			//in order to store an array as interface{}, you need
			//to use a pointer to it.
			//We could return the pointer instead of the value
			//and gain some performances here. Returning the pointer,
			//however, will open us up to race conditions.
			return *list, nil
		}

		return nil, errors.New("Stored value wasn't a list")
	}

	//return an empty list if not found
	return []string{}, nil
}
<commit_msg>format code<commit_after>package rmq

import (
	"errors"
	"strings"
	"sync"
	"time"
)

// TestRedisClient is a mock for redis
type TestRedisClient struct {
	store sync.Map
	ttl   sync.Map
}

var lock sync.Mutex

// NewTestRedisClient returns a NewTestRedisClient
func NewTestRedisClient() *TestRedisClient {
	return &TestRedisClient{}
}

// Set sets key to hold the string value.
// If key already holds a value, it is overwritten, regardless of its type.
// Any previous time to live associated with the key is discarded on successful SET operation.
func (client *TestRedisClient) Set(key string, value string, expiration time.Duration) error {

	lock.Lock()
	defer lock.Unlock()

	client.store.Store(key, value)
	//Delete any previous time to live associated with the key
	client.ttl.Delete(key)

	//0.0 expiration means that the value won't expire
	if expiration.Seconds() != 0.0 {
		//Store the unix time at which we should delete this
		client.ttl.Store(key, time.Now().Add(expiration).Unix())
	}

	return nil
}

// Get the value of key.
// If the key does not exist or isn't a string
// the special value nil is returned.
func (client *TestRedisClient) Get(key string) (string, error) {

	value, found := client.store.Load(key)

	if found {
		if stringValue, casted := value.(string); casted {
			return stringValue, nil
		}
	}

	return "nil", nil
}

// Del removes the specified key. A key is ignored if it does not exist.
func (client *TestRedisClient) Del(key string) (affected int64, err error) {

	_, found := client.store.Load(key)
	client.store.Delete(key)
	client.ttl.Delete(key)

	if found {
		return 1, nil
	}
	return 0, nil

}

// TTL returns the remaining time to live of a key that has a timeout.
// This introspection capability allows a Redis client to check how many seconds a given key will continue to be part of the dataset.
// In Redis 2.6 or older the command returns -1 if the key does not exist or if the key exist but has no associated expire.
// Starting with Redis 2.8 the return value in case of error changed:
// The command returns -2 if the key does not exist.
// The command returns -1 if the key exists but has no associated expire.
func (client *TestRedisClient) TTL(key string) (ttl time.Duration, err error) {

	//Lookup the expiration map
	expiration, found := client.ttl.Load(key)

	//Found an expiration time
	if found {

		//It was there, but it expired; removing it now
		if expiration.(int64) < time.Now().Unix() {
			client.ttl.Delete(key)
			return -2, nil
		}

		ttl = time.Duration(expiration.(int64) - time.Now().Unix())
		return ttl, nil
	}

	//Lookup the store in case this key exists but don't have an expiration
	//date
	_, found = client.store.Load(key)

	//The key was in store but didn't have an expiration associated
	//to it.
	if found {
		return -1, nil
	}

	return -2, nil
}

// LPush inserts the specified value at the head of the list stored at key.
// If key does not exist, it is created as empty list before performing the push operations.
// When key holds a value that is not a list, an error is returned.
// It is possible to push multiple elements using a single command call just specifying multiple arguments
// at the end of the command. Elements are inserted one after the other to the head of the list,
// from the leftmost element to the rightmost element.
func (client *TestRedisClient) LPush(key string, value ...string) (total int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	list, err := client.findList(key)

	if err != nil {
		return 0, nil
	}

	client.storeList(key, append(value, list...))
	return int64(len(list)) + 1, nil
}

// LLen returns the length of the list stored at key.
// If key does not exist, it is interpreted as an empty list and 0 is returned.
// An error is returned when the value stored at key is not a list.
func (client *TestRedisClient) LLen(key string) (affected int64, err error) {
	list, err := client.findList(key)

	if err != nil {
		return 0, nil
	}
	return int64(len(list)), nil
}

// LRem removes the first count occurrences of elements equal to
// value from the list stored at key. The count argument influences
// the operation in the following ways:
// count > 0: Remove elements equal to value moving from head to tail.
// count < 0: Remove elements equal to value moving from tail to head.
// count = 0: Remove all elements equal to value. For example,
// LREM list -2 "hello" will remove the last two occurrences of "hello" in
// the list stored at list. Note that non-existing keys are treated like empty
// lists, so when key does not exist, the command will always return 0.
func (client *TestRedisClient) LRem(key string, count int64, value string) (affected int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	list, err := client.findList(key)

	//Wasn't a list, or is empty
	if err != nil || len(list) == 0 {
		return 0, nil
	}

	//Create a list that have the capacity to store
	//the old one
	//This will be much more performant in case of
	//very long list
	newList := make([]string, 0, len(list))

	if err != nil {
		return 0, nil
	}

	//left to right removal of count elements
	if count >= 0 {

		//All the elements are to be removed.
		//Set count to max possible elements
		if count == 0 {
			count = int64(len(list))
		}
		//left to right traversal
		for index := 0; index < len(list); index++ {

			//isn't what we look for or we found enough element already
			if strings.Compare(list[index], value) != 0 || affected > count {
				newList = append(newList, list[index])
			} else {
				affected++
			}
		}
		//right to left removal of count elements
	} else if count < 0 {

		//right to left traversal
		for index := len(list) - 1; index >= 0; index-- {

			//isn't what we look for or we found enough element already
			if strings.Compare(list[index], value) != 0 || affected > count {
				//prepend instead of append to keep the order
				newList = append([]string{list[index]}, newList...)
			} else {
				affected++
			}
		}
	}

	//store the updated list
	client.storeList(key, newList)

	return affected, nil
}

// LTrim trims an existing list so that it will contain only the specified range of elements specified.
// Both start and stop are zero-based indexes, where 0 is the first element of the list (the head),
// 1 the next element and so on. For example: LTRIM foobar 0 2 will modify the list stored
// at foobar so that only the first three elements of the list will remain.
// start and end can also be negative numbers indicating offsets from the end of the list,
// where -1 is the last element of the list, -2 the penultimate element and so on.
// Out of range indexes will not produce an error: if start is larger than the end of the list,
// or start > end, the result will be an empty list (which causes key to be removed).
// If end is larger than the end of the list, Redis will treat it like the last element of the list
func (client *TestRedisClient) LTrim(key string, start, stop int64) error {

	lock.Lock()
	defer lock.Unlock()

	list, err := client.findList(key)

	//Wasn't a list, or is empty
	if err != nil || len(list) == 0 {
		return nil
	}

	if start < 0 {
		start += int64(len(list))
	}
	if stop < 0 {
		stop += int64(len(list))
	}

	//invalid values cause the remove of the key
	if start > stop {
		client.store.Delete(key)
		return nil
	}

	client.storeList(key, list[start:stop])
	return nil
}

// RPopLPush atomically returns and removes the last element (tail) of the list stored at source,
// and pushes the element at the first element (head) of the list stored at destination.
// For example: consider source holding the list a,b,c, and destination holding the list x,y,z.
// Executing RPOPLPUSH results in source holding a,b and destination holding c,x,y,z.
// If source does not exist, the value nil is returned and no operation is performed.
// If source and destination are the same, the operation is equivalent to removing the
// last element from the list and pushing it as first element of the list,
// so it can be considered as a list rotation command.
func (client *TestRedisClient) RPopLPush(source, destination string) (value string, err error) {

	lock.Lock()
	defer lock.Unlock()

	sourceList, sourceErr := client.findList(source)
	destList, destErr := client.findList(destination)

	//One of the two isn't a list
	if sourceErr != nil || destErr != nil {
		return "", ErrorNotFound
	}
	//we have nothing to move
	if len(sourceList) == 0 {
		return "", ErrorNotFound
	}

	//Remove the last element of source (tail)
	client.storeList(source, sourceList[0:len(sourceList)-1])
	//Put the last element of source (tail) and prepend it to dest
	client.storeList(destination, append([]string{sourceList[len(sourceList)-1]}, destList...))

	return sourceList[len(sourceList)-1], nil
}

// SAdd adds the specified members to the set stored at key.
// Specified members that are already a member of this set are ignored.
// If key does not exist, a new set is created before adding the specified members.
// An error is returned when the value stored at key is not a set.
func (client *TestRedisClient) SAdd(key, value string) (total int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	set, err := client.findSet(key)
	if err != nil {
		return 0, err
	}

	set[value] = struct{}{}
	client.storeSet(key, set)
	return int64(len(set)), nil
}

// SMembers returns all the members of the set value stored at key.
// This has the same effect as running SINTER with one argument key.
func (client *TestRedisClient) SMembers(key string) (members []string, err error) {
	set, err := client.findSet(key)
	if err != nil {
		return members, nil
	}

	members = make([]string, 0, len(set))
	for k := range set {
		members = append(members, k)
	}

	return members, nil
}

// SRem removes the specified members from the set stored at key.
// Specified members that are not a member of this set are ignored.
// If key does not exist, it is treated as an empty set and this command returns 0.
// An error is returned when the value stored at key is not a set.
func (client *TestRedisClient) SRem(key, value string) (affected int64, err error) {

	lock.Lock()
	defer lock.Unlock()

	set, err := client.findSet(key)
	if err != nil || len(set) == 0 {
		return 0, nil
	}

	if _, found := set[value]; found != false {
		delete(set, value)
		return 1, nil
	}

	return 0, nil
}

// FlushDb delete all the keys of the currently selected DB. This command never fails.
func (client *TestRedisClient) FlushDb() error {
	client.store = *new(sync.Map)
	client.ttl = *new(sync.Map)
	return nil
}

// storeSet stores a set
func (client *TestRedisClient) storeSet(key string, set map[string]struct{}) {
	client.store.Store(key, set)
}

// findSet finds a set
func (client *TestRedisClient) findSet(key string) (map[string]struct{}, error) {
	//Lookup the store for the list
	storedValue, found := client.store.Load(key)
	if found {
		//list are stored as pointer to []string
		set, casted := storedValue.(map[string]struct{})

		if casted {
			return set, nil
		}

		return nil, errors.New("Stored value wasn't a set")
	}

	//return an empty set if not found
	return make(map[string]struct{}), nil
}

// storeList is an helper function so others don't have to deal with pointers
func (client *TestRedisClient) storeList(key string, list []string) {
	client.store.Store(key, &list)
}

// findList returns the list stored at key.
// if key doesn't exist, an empty list is returned
// an error is returned when the value at key isn't a list
func (client *TestRedisClient) findList(key string) ([]string, error) {
	//Lookup the store for the list
	storedValue, found := client.store.Load(key)
	if found {

		//list are stored as pointer to []string
		list, casted := storedValue.(*[]string)

		//Successful cass from interface{} to *[]string
		//Preprend the new key
		if casted {

			//This mock use sync.Map to be thread safe.
			//sync.Map only accepts interface{} as values and
			//in order to store an array as interface{}, you need
			//to use a pointer to it.
			//We could return the pointer instead of the value
			//and gain some performances here. Returning the pointer,
			//however, will open us up to race conditions.
			return *list, nil
		}

		return nil, errors.New("Stored value wasn't a list")
	}

	//return an empty list if not found
	return []string{}, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 NATS Messaging System
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/nats-io/gnatsd/server"
	. "github.com/nats-io/nats-top/util"
)

func usage() {
	log.Fatalf("Usage: nats-top [-s server] [-m monitor] [-n num_connections] [-d delay_secs]\n")
}

var (
	host  = flag.String("s", "127.0.0.1", "The nats server host")
	port  = flag.Int("m", 8333, "The nats server monitoring port")
	conns = flag.Int("n", 1024, "Num of connections")
	delay = flag.Int("d", 1, "Delay in monitoring interval in seconds")
	sort  = flag.String("sort", "pending_size", "Value for which to sort connections")
)

func init() {
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
}

func main() {
	opts := make(map[string]interface{})
	opts["host"] = *host
	opts["port"] = *port
	opts["conns"] = *conns
	opts["delay"] = *delay
	opts["sort"] = *sort

	if opts["host"] == nil || opts["port"] == nil {
		log.Fatalf("Please specify the monitoring port for NATS.")
		usage()
	}

	// Smoke test the server once before starting
	_, err := Request("/varz", opts)
	if err != nil {
		log.Fatalf("ERROR: %v", err)
		os.Exit(1)
	}

	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	clearScreen()
	go StartSimpleUI(opts)

	select {
	case <-sigch:
		clearScreen()
		os.Exit(1)
	}
}

func clearScreen() {
	fmt.Print("\033[2J\033[1;1H")
}

func StartSimpleUI(opts map[string]interface{}) {
	var pollTime time.Time

	var inMsgsDelta int64
	var outMsgsDelta int64
	var inBytesDelta int64
	var outBytesDelta int64

	var inMsgsLastVal int64
	var outMsgsLastVal int64
	var inBytesLastVal int64
	var outBytesLastVal int64

	var inMsgsRate float64
	var outMsgsRate float64
	var inBytesRate float64
	var outBytesRate float64

	first := true
	pollTime = time.Now()
	for {
		wg := &sync.WaitGroup{}
		wg.Add(2)

		// Periodically poll for the varz, connz and routez
		var varz *server.Varz
		go func() {
			var err error
			defer wg.Done()

			result, err := Request("/varz", opts)
			if err != nil {
				log.Fatalf("Could not get /varz: %v", err)
			}

			if varzVal, ok := result.(*server.Varz); ok {
				varz = varzVal
			}
		}()

		var connz *server.Connz
		go func() {
			var err error
			defer wg.Done()

			result, err := Request("/connz", opts)
			if err != nil {
				log.Fatalf("Could not get /connz: %v", err)
			}

			if connzVal, ok := result.(*server.Connz); ok {
				connz = connzVal
			}
		}()
		wg.Wait()

		cpu := varz.CPU
		numConns := connz.NumConns
		memVal := varz.Mem

		// Periodic snapshot to get per sec metrics
		inMsgsVal := varz.InMsgs
		outMsgsVal := varz.OutMsgs
		inBytesVal := varz.InBytes
		outBytesVal := varz.OutBytes

		inMsgsDelta = inMsgsVal - inMsgsLastVal
		outMsgsDelta = outMsgsVal - outMsgsLastVal
		inBytesDelta = inBytesVal - inBytesLastVal
		outBytesDelta = outBytesVal - outBytesLastVal

		inMsgsLastVal = inMsgsVal
		outMsgsLastVal = outMsgsVal
		inBytesLastVal = inBytesVal
		outBytesLastVal = outBytesVal

		now := time.Now()
		tdelta := pollTime.Sub(now)
		pollTime = now

		// Calculate rates but the first time
		if !first {
			inMsgsRate = float64(inMsgsDelta) - tdelta.Seconds()
			outMsgsRate = float64(outMsgsDelta) - tdelta.Seconds()
			inBytesRate = float64(inBytesDelta) - tdelta.Seconds()
			outBytesRate = float64(outBytesDelta) - tdelta.Seconds()
		}

		mem := Psize(memVal)
		inMsgs := Psize(inMsgsVal)
		outMsgs := Psize(outMsgsVal)
		inBytes := Psize(inBytesVal)
		outBytes := Psize(outBytesVal)

		info := "\nServer:\n  Load: CPU: %.1f%%  Memory: %s\n"
		info += "  In:   Msgs: %s  Bytes: %s  Msgs/Sec: %.1f  Bytes/Sec: %.1f\n"
		info += "  Out:  Msgs: %s  Bytes: %s  Msgs/Sec: %.1f  Bytes/Sec: %.1f"

		text := fmt.Sprintf(info, cpu, mem,
			inMsgs, inBytes, inMsgsRate, inBytesRate,
			outMsgs, outBytes, outMsgsRate, outBytesRate)
		text += fmt.Sprintf("\n\nConnections: %d\n", numConns)

		connHeader := "  %-20s %-8s %-6s  %-10s  %-10s  %-10s  %-10s  %-10s\n"
		connRows := fmt.Sprintf(connHeader, "HOST", "CID", "SUBS", "PENDING", "MSGS_TO", "MSGS_FROM", "BYTES_TO", "BYTES_FROM")
		text += connRows

		connValues := "  %-20s %-8d %-6d  %-10d  %-10s  %-10s  %-10s  %-10s\n"
		for _, conn := range connz.Conns {
			host := fmt.Sprintf("%s:%d", conn.IP, conn.Port)
			connLine := fmt.Sprintf(connValues, host, conn.Cid, conn.NumSubs, conn.Pending,
				Psize(conn.OutMsgs), Psize(conn.InMsgs), Psize(conn.OutBytes), Psize(conn.InBytes))
			text += connLine
		}
		fmt.Print(text)

		if first {
			first = false
		}

		if val, ok := opts["delay"].(int); ok {
			time.Sleep(time.Duration(val) * time.Second)
			clearScreen()
		} else {
			log.Fatalf("error: could not use %s as a refreshing interval", opts["delay"])
			break
		}
	}
}
<commit_msg>Add placeholder for lang and version<commit_after>// Copyright (c) 2015 NATS Messaging System
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/nats-io/gnatsd/server"
	. "github.com/nats-io/nats-top/util"
)

func usage() {
	log.Fatalf("Usage: nats-top [-s server] [-m monitor] [-n num_connections] [-d delay_secs]\n")
}

var (
	host  = flag.String("s", "127.0.0.1", "The nats server host")
	port  = flag.Int("m", 8333, "The nats server monitoring port")
	conns = flag.Int("n", 1024, "Num of connections")
	delay = flag.Int("d", 1, "Delay in monitoring interval in seconds")
	sort  = flag.String("sort", "pending_size", "Value for which to sort connections")
)

func init() {
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
}

func main() {
	opts := make(map[string]interface{})
	opts["host"] = *host
	opts["port"] = *port
	opts["conns"] = *conns
	opts["delay"] = *delay
	opts["sort"] = *sort

	if opts["host"] == nil || opts["port"] == nil {
		log.Fatalf("Please specify the monitoring port for NATS.")
		usage()
	}

	// Smoke test the server once before starting
	_, err := Request("/varz", opts)
	if err != nil {
		log.Fatalf("ERROR: %v", err)
		os.Exit(1)
	}

	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	clearScreen()
	go StartSimpleUI(opts)

	select {
	case <-sigch:
		clearScreen()
		os.Exit(1)
	}
}

func clearScreen() {
	fmt.Print("\033[2J\033[1;1H")
}

func StartSimpleUI(opts map[string]interface{}) {
	var pollTime time.Time

	var inMsgsDelta int64
	var outMsgsDelta int64
	var inBytesDelta int64
	var outBytesDelta int64

	var inMsgsLastVal int64
	var outMsgsLastVal int64
	var inBytesLastVal int64
	var outBytesLastVal int64

	var inMsgsRate float64
	var outMsgsRate float64
	var inBytesRate float64
	var outBytesRate float64

	first := true
	pollTime = time.Now()
	for {
		wg := &sync.WaitGroup{}
		wg.Add(2)

		// Periodically poll for the varz, connz and routez
		var varz *server.Varz
		go func() {
			var err error
			defer wg.Done()

			result, err := Request("/varz", opts)
			if err != nil {
				log.Fatalf("Could not get /varz: %v", err)
			}

			if varzVal, ok := result.(*server.Varz); ok {
				varz = varzVal
			}
		}()

		var connz *server.Connz
		go func() {
			var err error
			defer wg.Done()

			result, err := Request("/connz", opts)
			if err != nil {
				log.Fatalf("Could not get /connz: %v", err)
			}

			if connzVal, ok := result.(*server.Connz); ok {
				connz = connzVal
			}
		}()
		wg.Wait()

		cpu := varz.CPU
		numConns := connz.NumConns
		memVal := varz.Mem

		// Periodic snapshot to get per sec metrics
		inMsgsVal := varz.InMsgs
		outMsgsVal := varz.OutMsgs
		inBytesVal := varz.InBytes
		outBytesVal := varz.OutBytes

		inMsgsDelta = inMsgsVal - inMsgsLastVal
		outMsgsDelta = outMsgsVal - outMsgsLastVal
		inBytesDelta = inBytesVal - inBytesLastVal
		outBytesDelta = outBytesVal - outBytesLastVal

		inMsgsLastVal = inMsgsVal
		outMsgsLastVal = outMsgsVal
		inBytesLastVal = inBytesVal
		outBytesLastVal = outBytesVal

		now := time.Now()
		tdelta := pollTime.Sub(now)
		pollTime = now

		// Calculate rates but the first time
		if !first {
			inMsgsRate = float64(inMsgsDelta) - tdelta.Seconds()
			outMsgsRate = float64(outMsgsDelta) - tdelta.Seconds()
			inBytesRate = float64(inBytesDelta) - tdelta.Seconds()
			outBytesRate = float64(outBytesDelta) - tdelta.Seconds()
		}

		mem := Psize(memVal)
		inMsgs := Psize(inMsgsVal)
		outMsgs := Psize(outMsgsVal)
		inBytes := Psize(inBytesVal)
		outBytes := Psize(outBytesVal)

		info := "\nServer:\n  Load: CPU: %.1f%%  Memory: %s\n"
		info += "  In:   Msgs: %s  Bytes: %s  Msgs/Sec: %.1f  Bytes/Sec: %.1f\n"
		info += "  Out:  Msgs: %s  Bytes: %s  Msgs/Sec: %.1f  Bytes/Sec: %.1f"

		text := fmt.Sprintf(info, cpu, mem,
			inMsgs, inBytes, inMsgsRate, inBytesRate,
			outMsgs, outBytes, outMsgsRate, outBytesRate)
		text += fmt.Sprintf("\n\nConnections: %d\n", numConns)

		connHeader := "  %-20s %-8s %-6s  %-10s  %-10s  %-10s  %-10s  %-10s  %-10s  %-10s\n"
		connRows := fmt.Sprintf(connHeader, "HOST", "CID", "SUBS", "PENDING",
			                            "MSGS_TO", "MSGS_FROM", "BYTES_TO", "BYTES_FROM",
			                            "LANG", "VERSION")
		text += connRows

		connValues := "  %-20s %-8d %-6d  %-10d  %-10s  %-10s  %-10s  %-10s  %-10s  %-10s\n"
		for _, conn := range connz.Conns {
			host := fmt.Sprintf("%s:%d", conn.IP, conn.Port)
			connLine := fmt.Sprintf(connValues, host, conn.Cid, conn.NumSubs, conn.Pending,
				Psize(conn.OutMsgs), Psize(conn.InMsgs), Psize(conn.OutBytes), Psize(conn.InBytes),
				conn.Lang, conn.Version)
			text += connLine
		}
		fmt.Print(text)

		if first {
			first = false
		}

		if val, ok := opts["delay"].(int); ok {
			time.Sleep(time.Duration(val) * time.Second)
			clearScreen()
		} else {
			log.Fatalf("error: could not use %s as a refreshing interval", opts["delay"])
			break
		}
	}
}
<|endoftext|>


<commit_before>package app

import (
	"github.com/hashicorp/memberlist"
	"github.com/tpjg/goriakpbc"
	"log"
	"math"
	"math/rand"
	"strconv"
	"time"
)

type Queues struct {
	// a container for all queues
	QueueMap map[string]Queue
	// Settings for Queues in general, ie queue list
	Config *riak.RDtMap
}

type Queue struct {
	// the definition of a queue
	// name of the queue
	Name string
	// the partitions of the queue
	Parts Partitions
	// Individual settings for the queue
	Config *riak.RDtMap
}

// get a message from the queue
func (queue Queue) Get(cfg Config, list *memberlist.Memberlist, batchsize uint32) ([]riak.RObject, error) {
	// get the top and bottom partitions
	partBottom, partTop, err := queue.Parts.GetPartition(cfg, queue.Name, list)

	if err != nil {
		return nil, err
	}
	// grab a riak client
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	//set the bucket
	bucket, err := client.NewBucket(queue.Name)
	if err != nil {
		log.Printf("Error%v", err)
	}
	//get a list of batchsize message ids
	messageIds, _, err := bucket.IndexQueryRangePage("id_int", strconv.Itoa(partBottom), strconv.Itoa(partTop), batchsize, "")

	if err != nil {
		log.Printf("Error%v", err)
	}
	log.Println("Message retrieved ", len(messageIds))
	return queue.RetrieveMessages(messageIds, cfg), err
}

// Put a Message onto the queue
func (queue Queue) Put(cfg Config, message string) string {
	//Grab our bucket
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		//Retrieve a UUID
		rand.Seed(time.Now().UnixNano())
		randInt := rand.Int63n(math.MaxInt64)
		uuid := strconv.FormatInt(randInt, 10)

		messageObj := bucket.NewObject(uuid)
		messageObj.Indexes["id_int"] = []string{uuid}
		messageObj.Data = []byte(message)
		messageObj.Store()
		return uuid
	} else {
		//Actually want to handle this in some other way
		return ""
	}
}

// Delete a Message from the queue
func (queue Queue) Delete(cfg Config, id string) bool {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		log.Println("Deleting: ", id)
		err = bucket.Delete(id)
		if err == nil {
			return true
		}
	}
	if err != nil {
		log.Println(err)
	}
	// if we got here we're borked
	return false
}

// helpers
func (queue Queue) RetrieveMessages(ids []string, cfg Config) []riak.RObject {
	var rObjectArrayChan = make(chan []riak.RObject, len(ids))
	var rKeys = make(chan string, len(ids))

	start := time.Now()
	//fmt.Println("In get multi")
	for i := 0; i < len(ids); i++ {
		go func() {
			var riakKey string
			client := cfg.RiakConnection()
			defer cfg.ReleaseRiakConnection(client)
			//fmt.Println("Getting bucket")
			bucket, _ := client.NewBucket(queue.Name)
			riakKey = <-rKeys
			//fmt.Println("Getting value")
			rObject, _ := bucket.Get(riakKey)

			rObjectArrayChan <- []riak.RObject{*rObject}

			//fmt.Println("Returning value")
		}()
		rKeys <- ids[i]
	}
	returnVals := make([]riak.RObject, 0)
	for i := 0; i < len(ids); i++ {
		var rObjectArray = <-rObjectArrayChan
		//If the key isn't blank, we've got a meaningful object to deal with
		if len(rObjectArray) == 1 {
			returnVals = append(returnVals, rObjectArray[0])
		}
	}
	elapsed := time.Since(start)
	log.Printf("Get Multi Took %s\n", elapsed)
	return returnVals
}

func (queues Queues) syncConfig(cfg Config) {
	for {
		log.Println("syncing Queue config with Riak")
		client := cfg.RiakConnection()
		bucket, err := client.NewBucketType("maps", CONFIGURATION_BUCKET)
		if err != nil {
			log.Println(err)
		}
		queues.Config, err = bucket.FetchMap(QUEUE_CONFIG_NAME)
		if err != nil {
			log.Println(err)
		}

		//iterate the map and add or remove topics that need to be destroyed
		queueSlice := queues.Config.FetchSet(QUEUE_SET_NAME).GetValue()
		if queueSlice == nil {
			//bail if there aren't any topics
			//but not before sleeping
			cfg.ReleaseRiakConnection(client)
			time.Sleep(cfg.Core.SyncConfigInterval * time.Second)
			continue
		}

		//Is there a better way to do this?
		//iterate over the queues in riak and add the missing ones
		queuesToKeep := make(map[string]bool)
		for _, queue := range queueSlice {
			var present bool
			_, present = queues.QueueMap[string(queue)]
			if present != true {
				// analog of topics.InitTopic
			}
			queuesToKeep[string(queue)] = true
		}

		//iterate over the topics in topics.TopicMap and delete the ones no longer used
		for queue, _ := range queues.QueueMap {
			var present bool
			_, present = queuesToKeep[queue]
			if present != true {
				delete(queues.QueueMap, queue)
			}
		}

		//sync all topics with riak

		for _, queue := range queues.QueueMap {
			queue.syncConfig(cfg)
		}
		//sleep for the configured interval
		cfg.ReleaseRiakConnection(client)
		time.Sleep(cfg.Core.SyncConfigInterval * time.Millisecond)
	}
}

func (queue Queue) syncConfig(cfg Config) {
	//refresh the queue RDtMap
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	recordName := queueConfigRecordName(queue.Name)
	queue.Config, _ = bucket.FetchMap(recordName)
	cfg.Queues.QueueMap[queue.Name] = queue
}
<commit_msg>Filling in //TODO to init queue on riak sync of queue data<commit_after>package app

import (
	"github.com/hashicorp/memberlist"
	"github.com/tpjg/goriakpbc"
	"log"
	"math"
	"math/rand"
	"strconv"
	"time"
)

type Queues struct {
	// a container for all queues
	QueueMap map[string]Queue
	// Settings for Queues in general, ie queue list
	Config *riak.RDtMap
}

type Queue struct {
	// the definition of a queue
	// name of the queue
	Name string
	// the partitions of the queue
	Parts Partitions
	// Individual settings for the queue
	Config *riak.RDtMap
}

// get a message from the queue
func (queue Queue) Get(cfg Config, list *memberlist.Memberlist, batchsize uint32) ([]riak.RObject, error) {
	// get the top and bottom partitions
	partBottom, partTop, err := queue.Parts.GetPartition(cfg, queue.Name, list)

	if err != nil {
		return nil, err
	}
	// grab a riak client
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	//set the bucket
	bucket, err := client.NewBucket(queue.Name)
	if err != nil {
		log.Printf("Error%v", err)
	}
	//get a list of batchsize message ids
	messageIds, _, err := bucket.IndexQueryRangePage("id_int", strconv.Itoa(partBottom), strconv.Itoa(partTop), batchsize, "")

	if err != nil {
		log.Printf("Error%v", err)
	}
	log.Println("Message retrieved ", len(messageIds))
	return queue.RetrieveMessages(messageIds, cfg), err
}

// Put a Message onto the queue
func (queue Queue) Put(cfg Config, message string) string {
	//Grab our bucket
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		//Retrieve a UUID
		rand.Seed(time.Now().UnixNano())
		randInt := rand.Int63n(math.MaxInt64)
		uuid := strconv.FormatInt(randInt, 10)

		messageObj := bucket.NewObject(uuid)
		messageObj.Indexes["id_int"] = []string{uuid}
		messageObj.Data = []byte(message)
		messageObj.Store()
		return uuid
	} else {
		//Actually want to handle this in some other way
		return ""
	}
}

// Delete a Message from the queue
func (queue Queue) Delete(cfg Config, id string) bool {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, err := client.NewBucket(queue.Name)
	if err == nil {
		log.Println("Deleting: ", id)
		err = bucket.Delete(id)
		if err == nil {
			return true
		}
	}
	if err != nil {
		log.Println(err)
	}
	// if we got here we're borked
	return false
}

// helpers
func (queue Queue) RetrieveMessages(ids []string, cfg Config) []riak.RObject {
	var rObjectArrayChan = make(chan []riak.RObject, len(ids))
	var rKeys = make(chan string, len(ids))

	start := time.Now()
	//fmt.Println("In get multi")
	for i := 0; i < len(ids); i++ {
		go func() {
			var riakKey string
			client := cfg.RiakConnection()
			defer cfg.ReleaseRiakConnection(client)
			//fmt.Println("Getting bucket")
			bucket, _ := client.NewBucket(queue.Name)
			riakKey = <-rKeys
			//fmt.Println("Getting value")
			rObject, _ := bucket.Get(riakKey)

			rObjectArrayChan <- []riak.RObject{*rObject}

			//fmt.Println("Returning value")
		}()
		rKeys <- ids[i]
	}
	returnVals := make([]riak.RObject, 0)
	for i := 0; i < len(ids); i++ {
		var rObjectArray = <-rObjectArrayChan
		//If the key isn't blank, we've got a meaningful object to deal with
		if len(rObjectArray) == 1 {
			returnVals = append(returnVals, rObjectArray[0])
		}
	}
	elapsed := time.Since(start)
	log.Printf("Get Multi Took %s\n", elapsed)
	return returnVals
}

func (queues Queues) syncConfig(cfg Config) {
	for {
		log.Println("syncing Queue config with Riak")
		client := cfg.RiakConnection()
		bucket, err := client.NewBucketType("maps", CONFIGURATION_BUCKET)
		if err != nil {
			log.Println(err)
		}
		queues.Config, err = bucket.FetchMap(QUEUE_CONFIG_NAME)
		if err != nil {
			log.Println(err)
		}

		//iterate the map and add or remove topics that need to be destroyed
		queueSlice := queues.Config.FetchSet(QUEUE_SET_NAME).GetValue()
		if queueSlice == nil {
			//bail if there aren't any topics
			//but not before sleeping
			cfg.ReleaseRiakConnection(client)
			time.Sleep(cfg.Core.SyncConfigInterval * time.Second)
			continue
		}

		//Is there a better way to do this?
		//iterate over the queues in riak and add the missing ones
		queuesToKeep := make(map[string]bool)
		for _, queue := range queueSlice {
			queueName := string(queue)
			var present bool
			_, present = queues.QueueMap[queueName]
			if present != true {
				initQueueFromRiak(cfg, queueName)
			}
			queuesToKeep[queueName] = true
		}

		//iterate over the topics in topics.TopicMap and delete the ones no longer used
		for queue, _ := range queues.QueueMap {
			var present bool
			_, present = queuesToKeep[queue]
			if present != true {
				delete(queues.QueueMap, queue)
			}
		}

		//sync all topics with riak

		for _, queue := range queues.QueueMap {
			queue.syncConfig(cfg)
		}
		//sleep for the configured interval
		cfg.ReleaseRiakConnection(client)
		time.Sleep(cfg.Core.SyncConfigInterval * time.Millisecond)
	}
}

func initQueueFromRiak(cfg Config, queueName string) {
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)

	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	config, _ := bucket.FetchMap(queueConfigRecordName(queueName))

	queue := Queue{
		Name:   queueName,
		Parts:  InitPartitions(cfg, queueName),
		Config: config,
	}

	cfg.Queues.QueueMap[queueName] = queue
}

func (queue Queue) syncConfig(cfg Config) {
	//refresh the queue RDtMap
	client := cfg.RiakConnection()
	defer cfg.ReleaseRiakConnection(client)
	bucket, _ := client.NewBucketType("maps", CONFIGURATION_BUCKET)
	recordName := queueConfigRecordName(queue.Name)
	queue.Config, _ = bucket.FetchMap(recordName)
	cfg.Queues.QueueMap[queue.Name] = queue
}
<|endoftext|>


<commit_before>package crypto

///////////////////////////////////////////////////////////////////////
// import external declarations

import (
	"crypto/rand"
	"fmt"
	"math/big"
	"testing"
)

///////////////////////////////////////////////////////////////////////
// Test cases

func TestCounter(t *testing.T) {

	fmt.Println("********************************************")
	fmt.Println("crypto/counter Test")
	fmt.Println("********************************************")
	fmt.Println()

	for size := 128; size <= 2048; size *= 2 {
		fmt.Printf("%d: ", size)
		priv, err := NewPaillierPrivateKey(size)
		if err != nil {
			fmt.Printf("Error generating private key of size %d bits.\n", size)
			t.Fail()
		}
		for i := 0; i < 10; i++ {
			if !test_counter(priv) {
				t.Fail()
			}
			fmt.Print(".")
		}
		fmt.Println()
	}
}

///////////////////////////////////////////////////////////////////////
/*
 * Test Counter
 */
func test_counter(priv *PaillierPrivateKey) bool {

	max := big.NewInt(2)
	pub := priv.GetPublicKey()
	cnt, err := NewCounter(pub)
	if err != nil {
		return false
	}
	var inc int64 = 0
	for i := 0; i < 100; i++ {
		v, err := rand.Int(rand.Reader, max)
		if err != nil {
			return false
		}
		cnt.Increment(v)
		if v.Bit(0) == 1 {
			inc++
		}
	}
	t := cnt.Get()
	t,err = priv.Decrypt (t)
	if err != nil {
		return false
	}
	v := t.Int64()
	if v == inc {
		return true
	}
	fmt.Printf ("Counter mismatch: %d -- %d\n", v, inc)
	return false
}
<commit_msg>Automatic code re-formatting.<commit_after>package crypto

///////////////////////////////////////////////////////////////////////
// import external declarations

import (
	"crypto/rand"
	"fmt"
	"math/big"
	"testing"
)

///////////////////////////////////////////////////////////////////////
// Test cases

func TestCounter(t *testing.T) {

	fmt.Println("********************************************")
	fmt.Println("crypto/counter Test")
	fmt.Println("********************************************")
	fmt.Println()

	for size := 128; size <= 2048; size *= 2 {
		fmt.Printf("%d: ", size)
		priv, err := NewPaillierPrivateKey(size)
		if err != nil {
			fmt.Printf("Error generating private key of size %d bits.\n", size)
			t.Fail()
		}
		for i := 0; i < 10; i++ {
			if !test_counter(priv) {
				t.Fail()
			}
			fmt.Print(".")
		}
		fmt.Println()
	}
}

///////////////////////////////////////////////////////////////////////
/*
 * Test Counter
 */
func test_counter(priv *PaillierPrivateKey) bool {

	max := big.NewInt(2)
	pub := priv.GetPublicKey()
	cnt, err := NewCounter(pub)
	if err != nil {
		return false
	}
	var inc int64 = 0
	for i := 0; i < 100; i++ {
		v, err := rand.Int(rand.Reader, max)
		if err != nil {
			return false
		}
		cnt.Increment(v)
		if v.Bit(0) == 1 {
			inc++
		}
	}
	t := cnt.Get()
	t, err = priv.Decrypt(t)
	if err != nil {
		return false
	}
	v := t.Int64()
	if v == inc {
		return true
	}
	fmt.Printf("Counter mismatch: %d -- %d\n", v, inc)
	return false
}
<|endoftext|>


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

package s3util

import (
	"github.com/jacobsa/aws/s3"
)

// List all keys currently contained by the bucket.
func ListAllKeys(b s3.Bucket) (keys []string, error)
<commit_msg>Added a stub implementation.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package s3util

import (
	"fmt"
	"github.com/jacobsa/aws/s3"
)

// List all keys currently contained by the bucket.
func ListAllKeys(b s3.Bucket) (keys []string, err error) {
	err = fmt.Errorf("TODO")
	return
}
<|endoftext|>


<commit_before>// Package racy deals with file cryptography
package racy

import (
	"crypto/sha512"
	"hash"
	"io"
	"sync"
)

// NewHash returns a new hash.Hash
var NewHash = sha512.New

type cachedMap struct {
	mu *sync.Mutex
	m  map[string][]byte
}

// Racy is used in hashing targets.
// https://www.kernel.org/pub/software/scm/git/docs/technical/racy-git.txt
type Racy struct {
	hash.Hash

	mu *sync.Mutex
	m  map[string][]byte
}

// New takes nothing and returns a new Racy
func New() *Racy {
	return &Racy{
		NewHash(),

		&sync.Mutex{},
		make(map[string][]byte),
	}
}

func hashString(s string) []byte {
	h := NewHash()
	io.WriteString(h, s)
	return h.Sum(nil)
}

// HashString takes a hash and one or more strings and writes the strings to the
// hash.
func HashString(h io.Writer, strs ...string) {
	for _, str := range strs {
		io.WriteString(h, str)
	}
}
<commit_msg>racy: change the racy API<commit_after>// Package racy deals with file cryptography
package racy

import (
	"crypto/sha512"
	"fmt"
	"hash"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// NewHash returns a new hash.Hash
var NewHash = sha512.New

var (
	cacheMutex = new(sync.Mutex)
	hashCache  = make(map[string][]byte)
)

// Racy is used in hashing targets.
// https://www.kernel.org/pub/software/scm/git/docs/technical/racy-git.txt
type Racy struct {
	hash.Hash

	exts []string

	mu *sync.Mutex
	m  map[string][]byte
}

type Option func(*Racy)

// New takes nothing and returns a new Racy
func New(options ...Option) *Racy {
	x := &Racy{
		NewHash(),

		[]string{}, // by default we'll hash everything checkout `AllowExtension` option to limit the files hashed
		&sync.Mutex{},
		make(map[string][]byte),
	}
	for _, option := range options {
		option(x)
	}
	return x
}

func AllowExtension(ext string) Option {
	ext = strings.TrimLeft(ext, ".")
	return func(r *Racy) { r.exts = append(r.exts, ext) }
}

func hashString(s string) []byte {
	h := NewHash()
	io.WriteString(h, s)
	return h.Sum(nil)
}

func hashFile(file string) []byte {
CHECK:
	if sum, ok := hashCache[file]; ok {
		return sum
	}
	// For speediness, we'll cache the hash of the file
	// and hash that instead of opening the file,
	// reading it's contents everytime
	cacheMutex.Lock()
	f, err := os.Open(file)
	if err != nil {
		log.Fatalf("racy.hashFile: error opening file %q: %v", file, err)
	}
	tmpHash := NewHash()
	io.Copy(tmpHash, f)
	hashCache[file] = tmpHash.Sum(nil)
	f.Close()
	cacheMutex.Unlock()
	goto CHECK
}

func (r *Racy) hashDir(dir string) {
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			fmt.Printf("prevent panic by handling failure accessing a path %q: %v\n", dir, err)
			return err
		}
		if info.IsDir() {
			r.hashDir(path)
		} else {
			r.hashFile(path)
		}
		return nil
	})
	log.Fatalf("racy.hashDir: error walking dir %q: %v", dir, err)
}

func (r *Racy) hashFile(file string) {
	if !r.allowedExt(file) {
		return
	}
	r.Write(hashFile(file))
}

func (r *Racy) allowedExt(file string) bool {
	ext := strings.TrimLeft(filepath.Ext(file), ".")
	for _, allowedExt := range r.exts {
		if ext == allowedExt {
			return true
		}
	}
	return false
}

func (r *Racy) HashFiles(files ...string) {
	for _, file := range files {
		stat, err := os.Stat(file)
		if err != nil {
			log.Fatalf("racy.HashFiles: error opening file: %v", err)
		}
		if !stat.IsDir() {
			r.hashFile(file)
		} else {
			r.hashDir(file)
		}
	}
	return
}

// HashStrings hashes strings written to Racy
func (r *Racy) HashStrings(strs ...string) {
	for _, str := range strs {
		io.WriteString(r, str)
	}
}<|endoftext|>


<commit_before>package raft

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/celrenheit/sandglass"

	"github.com/celrenheit/sandglass/logy"
	"github.com/celrenheit/sandglass/topic"

	"github.com/hashicorp/raft"
	raftboltdb "github.com/hashicorp/raft-boltdb"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 1 * time.Second
)

const (
	CreateTopicOp            = "CreateTopicOp"
	SetPartitionLeaderOp     = "SetPartitionLeader"
	SetPartitionLeaderBulkOp = "SetPartitionLeaderBulkOp"
)

type Config struct {
	Name          string
	Addr          string
	Dir           string
	StartAsLeader bool
}

type Store struct {
	conf   Config
	raft   *raft.Raft
	logger logy.Logger
	mu     sync.RWMutex

	state            *state
	leaderChangeChan chan bool
	newTopicChan     chan *topic.Topic

	wg sync.WaitGroup
}

type state struct {
	topics           map[string]*topic.Topic
	partitionLeaders map[string]map[string]string
}

func newState() *state {
	return &state{
		topics:           map[string]*topic.Topic{},
		partitionLeaders: map[string]map[string]string{},
	}
}

func New(conf Config, logger logy.Logger) *Store {
	return &Store{
		conf:         conf,
		logger:       logger,
		state:        newState(),
		newTopicChan: make(chan *topic.Topic),
	}
}

func (s *Store) Init(bootstrap bool) error {
	config := raft.DefaultConfig()

	address := s.conf.Addr
	serverId := s.conf.Name
	if serverId == "" {
		serverId = address
	}
	config.Logger = log.New(os.Stdout, "raft["+serverId+"] ", log.LstdFlags)
	config.LocalID = raft.ServerID(serverId)
	config.StartAsLeader = s.conf.StartAsLeader

	addr, err := net.ResolveTCPAddr("tcp", s.conf.Addr)
	if err != nil {
		return err
	}

	transport, err := raft.NewTCPTransport(s.conf.Addr, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	snapshots, err := raft.NewFileSnapshotStore(s.conf.Dir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return err
	}

	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.conf.Dir, "raft.db"))
	if err != nil {
		return err
	}

	if bootstrap {
		hasState, err := raft.HasExistingState(logStore, logStore, snapshots)
		if err != nil {
			return err
		}

		if !hasState {
			configuration := raft.Configuration{}
			configuration.Servers = append(configuration.Servers, raft.Server{
				Suffrage: raft.Voter,
				ID:       raft.ServerID(serverId),
				Address:  raft.ServerAddress(address),
			})
			err := raft.BootstrapCluster(config, logStore, logStore, snapshots, transport, configuration)
			if err != nil {
				return err
			}
		}
	}

	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra
	return nil
}

func (s *Store) AddNode(n *sandglass.Node) error {
	s.logger.Info("adding node: %v", n)
	return s.raft.AddVoter(raft.ServerID(n.Name), raft.ServerAddress(n.RAFTAddr), 0, raftTimeout).Error()
}

func (s *Store) RemoveNode(n *sandglass.Node) error {
	return s.raft.RemoveServer(raft.ServerID(n.Name), 0, raftTimeout).Error()
}

func (s *Store) Stop() error {
	return s.raft.Shutdown().Error()
}

func (s *Store) Leader() string {
	return string(s.raft.Leader())
}

func (s *Store) LeaderCh() <-chan bool {
	return s.raft.LeaderCh()
}

func (s *Store) NewTopicChan() chan *topic.Topic {
	return s.newTopicChan
}

func (s *Store) GetTopics() []*topic.Topic {
	s.mu.RLock()
	defer s.mu.RUnlock()

	out := make([]*topic.Topic, 0, len(s.state.topics))
	for _, t := range s.state.topics {
		out = append(out, t)
	}

	return out
}

func (s *fsm) HasTopic(name string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	for _, t := range s.state.topics {
		if t.Name == name {
			return true
		}
	}

	return false
}

type fsm Store

func (f *fsm) Apply(l *raft.Log) (value interface{}) {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	var val interface{}
	switch c.Op {
	case CreateTopicOp:
		val = f.applySetTopic(c.Payload)
	case SetPartitionLeaderBulkOp:
		val = f.applySetPartitionLeaderBulk(c.Payload)
	default:
		panic(fmt.Sprintf("unrecognized command op: %s", c.Op))
	}

	if val != nil {
		f.logger.Debug("Apply err: %v", val)
		return val
	}

	// return f.raft.Barrier(raftTimeout).Error()
	return nil
}

func (f *fsm) applySetTopic(b []byte) error {
	var t topic.Topic
	if err := json.Unmarshal(b, &t); err != nil {
		return err
	}
	f.logger.Info("applyCreateTopic %v", t.Name)

	if !f.HasTopic(t.Name) {
		err := t.InitStore(f.conf.Dir)
		if err != nil {
			return err
		}
		f.mu.Lock()
		f.state.topics[t.Name] = &t
		f.mu.Unlock()
	}
	// select {
	// case f.newTopicChan <- &t:
	// default:
	// }

	return nil
}

func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	f.mu.RLock()
	defer f.mu.RUnlock()

	state := newState()
	for _, topic := range f.state.topics {
		state.topics[topic.Name] = topic
	}

	for topic, partitions := range f.state.partitionLeaders {
		state.partitionLeaders[topic] = make(map[string]string)
		for part, leader := range partitions {
			state.partitionLeaders[topic][part] = leader
		}
	}

	return &fsmSnapshot{state}, nil
}

// Restore stores the key-value store to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	var restoredState state
	if err := json.NewDecoder(rc).Decode(&f); err != nil {
		return err
	}

	*f.state = restoredState
	return nil
}

func (f *fsm) applySetPartitionLeaderBulk(d []byte) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()

	var p setPartitionLeaderBulk
	if err := json.Unmarshal(d, &p); err != nil {
		return err
	}

	for topic, partitions := range p.State {
		if f.state.partitionLeaders[topic] == nil {
			f.state.partitionLeaders[topic] = map[string]string{}
		}

		for part, leader := range partitions {
			f.state.partitionLeaders[topic][part] = leader
		}
	}

	return nil
}

type command struct {
	Op      string          `json:"op,omitempty"`
	Payload json.RawMessage `json:"payload,omitempty"`
}

func (s *Store) CreateTopic(t *topic.Topic) error {
	return s.raftApplyCommand(CreateTopicOp, t)
}

func (s *Store) SetPartitionLeaderOp(topic, partition, leader string) error {
	return s.raftApplyCommand(SetPartitionLeaderOp, &setPartitionLeaderPayload{
		topic: topic, partition: partition, leader: leader,
	})
}

func (s *Store) GetTopic(name string) *topic.Topic {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.state.topics[name]
}

type setPartitionLeaderBulk struct {
	State map[string]map[string]string
}

func (s *Store) SetPartitionLeaderBulkOp(v map[string]map[string]string) error {
	return s.raftApplyCommand(SetPartitionLeaderBulkOp, setPartitionLeaderBulk{
		State: v,
	})
}

func (s *Store) GetPartitionLeader(topic, partition string) (string, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	partitions, ok := s.state.partitionLeaders[topic]
	if !ok {
		return "", false
	}

	leader, ok := partitions[partition]
	return leader, ok
}

type setPartitionLeaderPayload struct {
	topic, partition, leader string
}

func (s *Store) raftApplyCommand(op string, v interface{}) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

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

	cmd := &command{
		Op:      op,
		Payload: d,
	}
	b, err := json.Marshal(cmd)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	return f.Error()
}

type fsmSnapshot struct {
	state *state
}

func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Encode data.
		b, err := json.Marshal(f.state)
		if err != nil {
			return err
		}

		if _, err := sink.Write(b); err != nil {
			return err
		}

		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

func (f *fsmSnapshot) Release() {}
<commit_msg>raft: log when node is removed<commit_after>package raft

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/celrenheit/sandglass"

	"github.com/celrenheit/sandglass/logy"
	"github.com/celrenheit/sandglass/topic"

	"github.com/hashicorp/raft"
	raftboltdb "github.com/hashicorp/raft-boltdb"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 1 * time.Second
)

const (
	CreateTopicOp            = "CreateTopicOp"
	SetPartitionLeaderOp     = "SetPartitionLeader"
	SetPartitionLeaderBulkOp = "SetPartitionLeaderBulkOp"
)

type Config struct {
	Name          string
	Addr          string
	Dir           string
	StartAsLeader bool
}

type Store struct {
	conf   Config
	raft   *raft.Raft
	logger logy.Logger
	mu     sync.RWMutex

	state            *state
	leaderChangeChan chan bool
	newTopicChan     chan *topic.Topic

	wg sync.WaitGroup
}

type state struct {
	topics           map[string]*topic.Topic
	partitionLeaders map[string]map[string]string
}

func newState() *state {
	return &state{
		topics:           map[string]*topic.Topic{},
		partitionLeaders: map[string]map[string]string{},
	}
}

func New(conf Config, logger logy.Logger) *Store {
	return &Store{
		conf:         conf,
		logger:       logger,
		state:        newState(),
		newTopicChan: make(chan *topic.Topic),
	}
}

func (s *Store) Init(bootstrap bool) error {
	config := raft.DefaultConfig()

	address := s.conf.Addr
	serverId := s.conf.Name
	if serverId == "" {
		serverId = address
	}
	config.Logger = log.New(os.Stdout, "raft["+serverId+"] ", log.LstdFlags)
	config.LocalID = raft.ServerID(serverId)
	config.StartAsLeader = s.conf.StartAsLeader

	addr, err := net.ResolveTCPAddr("tcp", s.conf.Addr)
	if err != nil {
		return err
	}

	transport, err := raft.NewTCPTransport(s.conf.Addr, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	snapshots, err := raft.NewFileSnapshotStore(s.conf.Dir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return err
	}

	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.conf.Dir, "raft.db"))
	if err != nil {
		return err
	}

	if bootstrap {
		hasState, err := raft.HasExistingState(logStore, logStore, snapshots)
		if err != nil {
			return err
		}

		if !hasState {
			configuration := raft.Configuration{}
			configuration.Servers = append(configuration.Servers, raft.Server{
				Suffrage: raft.Voter,
				ID:       raft.ServerID(serverId),
				Address:  raft.ServerAddress(address),
			})
			err := raft.BootstrapCluster(config, logStore, logStore, snapshots, transport, configuration)
			if err != nil {
				return err
			}
		}
	}

	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra
	return nil
}

func (s *Store) AddNode(n *sandglass.Node) error {
	s.logger.Info("adding node: %+v", n)
	return s.raft.AddVoter(raft.ServerID(n.Name), raft.ServerAddress(n.RAFTAddr), 0, raftTimeout).Error()
}

func (s *Store) RemoveNode(n *sandglass.Node) error {
	s.logger.Info("removing node: %+v", n)
	return s.raft.RemoveServer(raft.ServerID(n.Name), 0, raftTimeout).Error()
}

func (s *Store) Stop() error {
	return s.raft.Shutdown().Error()
}

func (s *Store) Leader() string {
	return string(s.raft.Leader())
}

func (s *Store) LeaderCh() <-chan bool {
	return s.raft.LeaderCh()
}

func (s *Store) NewTopicChan() chan *topic.Topic {
	return s.newTopicChan
}

func (s *Store) GetTopics() []*topic.Topic {
	s.mu.RLock()
	defer s.mu.RUnlock()

	out := make([]*topic.Topic, 0, len(s.state.topics))
	for _, t := range s.state.topics {
		out = append(out, t)
	}

	return out
}

func (s *fsm) HasTopic(name string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	for _, t := range s.state.topics {
		if t.Name == name {
			return true
		}
	}

	return false
}

type fsm Store

func (f *fsm) Apply(l *raft.Log) (value interface{}) {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	var val interface{}
	switch c.Op {
	case CreateTopicOp:
		val = f.applySetTopic(c.Payload)
	case SetPartitionLeaderBulkOp:
		val = f.applySetPartitionLeaderBulk(c.Payload)
	default:
		panic(fmt.Sprintf("unrecognized command op: %s", c.Op))
	}

	if val != nil {
		f.logger.Debug("Apply err: %v", val)
		return val
	}

	return nil
}

func (f *fsm) applySetTopic(b []byte) error {
	var t topic.Topic
	if err := json.Unmarshal(b, &t); err != nil {
		return err
	}
	f.logger.Info("applyCreateTopic %v", t.Name)

	if !f.HasTopic(t.Name) {
		err := t.InitStore(f.conf.Dir)
		if err != nil {
			return err
		}
		f.mu.Lock()
		f.state.topics[t.Name] = &t
		f.mu.Unlock()
	}

	return nil
}

func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	f.mu.RLock()
	defer f.mu.RUnlock()

	state := newState()
	for _, topic := range f.state.topics {
		state.topics[topic.Name] = topic
	}

	for topic, partitions := range f.state.partitionLeaders {
		state.partitionLeaders[topic] = make(map[string]string)
		for part, leader := range partitions {
			state.partitionLeaders[topic][part] = leader
		}
	}

	return &fsmSnapshot{state}, nil
}

// Restore stores the key-value store to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	var restoredState state
	if err := json.NewDecoder(rc).Decode(&f); err != nil {
		return err
	}

	*f.state = restoredState
	return nil
}

func (f *fsm) applySetPartitionLeaderBulk(d []byte) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()

	var p setPartitionLeaderBulk
	if err := json.Unmarshal(d, &p); err != nil {
		return err
	}

	for topic, partitions := range p.State {
		if f.state.partitionLeaders[topic] == nil {
			f.state.partitionLeaders[topic] = map[string]string{}
		}

		for part, leader := range partitions {
			f.state.partitionLeaders[topic][part] = leader
		}
	}

	return nil
}

type command struct {
	Op      string          `json:"op,omitempty"`
	Payload json.RawMessage `json:"payload,omitempty"`
}

func (s *Store) CreateTopic(t *topic.Topic) error {
	return s.raftApplyCommand(CreateTopicOp, t)
}

func (s *Store) SetPartitionLeaderOp(topic, partition, leader string) error {
	return s.raftApplyCommand(SetPartitionLeaderOp, &setPartitionLeaderPayload{
		topic: topic, partition: partition, leader: leader,
	})
}

func (s *Store) GetTopic(name string) *topic.Topic {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.state.topics[name]
}

type setPartitionLeaderBulk struct {
	State map[string]map[string]string
}

func (s *Store) SetPartitionLeaderBulkOp(v map[string]map[string]string) error {
	return s.raftApplyCommand(SetPartitionLeaderBulkOp, setPartitionLeaderBulk{
		State: v,
	})
}

func (s *Store) GetPartitionLeader(topic, partition string) (string, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	partitions, ok := s.state.partitionLeaders[topic]
	if !ok {
		return "", false
	}

	leader, ok := partitions[partition]
	return leader, ok
}

type setPartitionLeaderPayload struct {
	topic, partition, leader string
}

func (s *Store) raftApplyCommand(op string, v interface{}) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

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

	cmd := &command{
		Op:      op,
		Payload: d,
	}
	b, err := json.Marshal(cmd)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	return f.Error()
}

type fsmSnapshot struct {
	state *state
}

func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Encode data.
		b, err := json.Marshal(f.state)
		if err != nil {
			return err
		}

		if _, err := sink.Write(b); err != nil {
			return err
		}

		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

func (f *fsmSnapshot) Release() {}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
)

// Barista makes coffee
type Barista interface {
	MakeOrder(Order) Order
	String() string
}

type barista struct {
	speed int
	id    int
}

func (b *barista) MakeOrder(order Order) Order {
	time.Sleep(time.Duration(b.id) * time.Second)
	return order
}

func (b *barista) String() string {
	return fmt.Sprintf("Barista %d", b.id)
}

// NewBarista creates a new Barista
var NewBarista = func(id int) Barista {
	return &barista{id: id, speed: 2}
}

// RandomGroupOfBaristas creates a random group of baristas
var RandomGroupOfBaristas = func(n int) []Barista {
	baristas := make([]Barista, n)
	for i := 0; i < n; i++ {
		baristas[i] = NewBarista(i + 1)
	}
	return baristas
}
<commit_msg>make the baristas twice as slow<commit_after>package main

import (
	"fmt"
	"time"
)

// Barista makes coffee
type Barista interface {
	MakeOrder(Order) Order
	String() string
}

type barista struct {
	speed int
	id    int
}

func (b *barista) MakeOrder(order Order) Order {
	time.Sleep(time.Duration(b.id) * time.Second)
	return order
}

func (b *barista) String() string {
	return fmt.Sprintf("Barista %d", b.id)
}

// NewBarista creates a new Barista
var NewBarista = func(id int) Barista {
	return &barista{id: id, speed: 5}
}

// RandomGroupOfBaristas creates a random group of baristas
var RandomGroupOfBaristas = func(n int) []Barista {
	baristas := make([]Barista, n)
	for i := 0; i < n; i++ {
		baristas[i] = NewBarista(i + 1)
	}
	return baristas
}
<|endoftext|>


<commit_before>package rc

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

	fmt "github.com/jhunt/go-ansi"
	"gopkg.in/yaml.v2"
)

type Config struct {
	Version int               `yaml:"version"`
	Current string            `yaml:"current"`
	Vaults  map[string]*Vault `yaml:"vaults"`
}

type Vault struct {
	URL        string `yaml:"url"`
	Token      string `yaml:"token"`
	SkipVerify bool   `yaml:"skip_verify"`
}

type oldConfig struct {
	Current    string                 `yaml:"Current"`
	Targets    map[string]interface{} `yaml:"Targets"`
	Aliases    map[string]string      `yaml:"Aliases"`
	SkipVerify map[string]bool        `yaml:"SkipVerify"`
}

func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("USERPROFILE")
		if home == "" {
			home = os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		}
		return home
	}
	return os.Getenv("HOME")
}

func saferc() string {
	return fmt.Sprintf("%s/.saferc", userHomeDir())
}

func svtoken() string {
	return fmt.Sprintf("%s/.svtoken", userHomeDir())
}

func (legacy *oldConfig) convert() Config {
	c := Config{
		Version: 1,
		Current: legacy.Current,
		Vaults:  make(map[string]*Vault),
	}

	for alias, url := range legacy.Aliases {
		v := &Vault{
			URL: url,
		}
		if skip, ok := legacy.SkipVerify[url]; ok {
			v.SkipVerify = skip
		}
		if token, ok := legacy.Targets[url]; ok && token != nil {
			v.Token = token.(string)
		}
		c.Vaults[alias] = v
	}

	return c
}

func (c *Config) credentials() (string, string, bool, error) {
	if c.Current == "" {
		return "", "", false, nil
	}

	v, ok, _ := c.Find(c.Current)
	if !ok {
		return "", "", false, fmt.Errorf("Current target vault '%s' not found in ~/.saferc", c.Current)
	}

	return v.URL, v.Token, v.SkipVerify, nil
}

func Read() Config {
	var c Config

	b, err := ioutil.ReadFile(saferc())
	if err != nil {
		return Config{Version: 1}
	}

	if err = yaml.Unmarshal(b, &c); err != nil {
		return Config{Version: 1}
	}
	if c.Version == 0 {
		var legacy oldConfig
		if err = yaml.Unmarshal(b, &legacy); err != nil {
			fmt.Fprintf(os.Stderr, "@R{!!! %s}\n", err)
			os.Exit(1)
		}
		c = legacy.convert()
	}

	return c
}

func Apply(use string) Config {
	c := Read()

	if err := c.Apply(use); err != nil {
		fmt.Fprintf(os.Stderr, "@R{!!! %s}\n", err)
		os.Exit(1)
	}
	return c
}

func (c *Config) Write() error {
	b, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

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

	v, err := c.Vault("")
	if err != nil {
		return err
	}
	if v == nil {
		os.Remove(svtoken())
		return nil
	}

	sv := struct {
		Vault      string `yaml:"vault"` /* this is different than Vault.URL */
		Token      string `yaml:"token"`
		SkipVerify bool   `yaml:"skip_verify"`
	}{v.URL, v.Token, v.SkipVerify}
	b, err = yaml.Marshal(sv)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(svtoken(), b, 0600)
}

func (c *Config) Apply(use string) error {
	v, err := c.Vault(use)
	if err != nil {
		fmt.Fprintf(os.Stderr, "@R{!!! %s}\n", err)
		os.Exit(1)
	}

	if v != nil {
		os.Setenv("VAULT_ADDR", v.URL)
		os.Setenv("VAULT_TOKEN", v.Token)
		if v.SkipVerify {
			os.Setenv("VAULT_SKIP_VERIFY", "1")
		}
	} else {
		if os.Getenv("VAULT_TOKEN") == "" {
			tokenFile := fmt.Sprintf("%s/.vault-token", os.Getenv("HOME"))
			b, err := ioutil.ReadFile(tokenFile)
			if err == nil {
				os.Setenv("VAULT_TOKEN", strings.TrimSpace(string(b)))
			}
		}
	}
	return nil
}

func (c *Config) SetCurrent(alias string, reskip bool) error {
	v, ok, err := c.Find(alias)
	if err != nil {
		return err
	}
	if !ok {
		return fmt.Errorf("Unknown target '%s'", alias)
	}
	c.Current = alias
	if reskip {
		v.SkipVerify = true
	}
	return nil
}

func (c *Config) SetTarget(alias, url string, skipverify bool) error {
	if c.Vaults == nil {
		c.Vaults = make(map[string]*Vault)
	}

	c.Current = alias
	c.Vaults[alias] = &Vault{
		URL:        url,
		SkipVerify: skipverify,
	}

	return nil
}

func (c *Config) SetToken(token string) error {
	if c.Current == "" {
		return fmt.Errorf("No target selected")
	}
	v, ok, _ := c.Find(c.Current)
	if !ok {
		return fmt.Errorf("Unknown target '%s'", c.Current)
	}
	v.Token = token
	return nil
}

func (c *Config) URL() string {
	if v, ok, _ := c.Find(c.Current); ok {
		return v.URL
	}
	return ""
}

func (c *Config) Verified() bool {
	if v, ok, _ := c.Find(c.Current); ok {
		return !v.SkipVerify
	}
	return false
}

func (c *Config) Find(alias string) (*Vault, bool, error) {
	if v, ok := c.Vaults[alias]; ok {
		return v, true, nil
	}

	var v *Vault
	n := 0
	want := strings.TrimSuffix(alias, "/")

	for _, maybe := range c.Vaults {
		if strings.TrimSuffix(maybe.URL, "/") == want {
			n++
			v = maybe
		}
	}
	if n == 1 {
		return v, true, nil
	}
	if n > 1 {
		return nil, true, fmt.Errorf("More than one target for Vault at '%s' (maybe try an alias?)", alias)
	}

	return nil, false, nil
}

func (c *Config) Vault(which string) (*Vault, error) {
	if which == "" {
		which = c.Current
	}

	if which == "" {
		return nil, nil /* not an error */
	}

	v, ok, err := c.Find(which)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("Current target '%s' not found in ~/.saferc", which)
	}
	return v, nil
}
<commit_msg>Don't delete token on target if address is the same<commit_after>package rc

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

	fmt "github.com/jhunt/go-ansi"
	"gopkg.in/yaml.v2"
)

type Config struct {
	Version int               `yaml:"version"`
	Current string            `yaml:"current"`
	Vaults  map[string]*Vault `yaml:"vaults"`
}

type Vault struct {
	URL        string `yaml:"url"`
	Token      string `yaml:"token"`
	SkipVerify bool   `yaml:"skip_verify"`
}

type oldConfig struct {
	Current    string                 `yaml:"Current"`
	Targets    map[string]interface{} `yaml:"Targets"`
	Aliases    map[string]string      `yaml:"Aliases"`
	SkipVerify map[string]bool        `yaml:"SkipVerify"`
}

func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("USERPROFILE")
		if home == "" {
			home = os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		}
		return home
	}
	return os.Getenv("HOME")
}

func saferc() string {
	return fmt.Sprintf("%s/.saferc", userHomeDir())
}

func svtoken() string {
	return fmt.Sprintf("%s/.svtoken", userHomeDir())
}

func (legacy *oldConfig) convert() Config {
	c := Config{
		Version: 1,
		Current: legacy.Current,
		Vaults:  make(map[string]*Vault),
	}

	for alias, url := range legacy.Aliases {
		v := &Vault{
			URL: url,
		}
		if skip, ok := legacy.SkipVerify[url]; ok {
			v.SkipVerify = skip
		}
		if token, ok := legacy.Targets[url]; ok && token != nil {
			v.Token = token.(string)
		}
		c.Vaults[alias] = v
	}

	return c
}

func (c *Config) credentials() (string, string, bool, error) {
	if c.Current == "" {
		return "", "", false, nil
	}

	v, ok, _ := c.Find(c.Current)
	if !ok {
		return "", "", false, fmt.Errorf("Current target vault '%s' not found in ~/.saferc", c.Current)
	}

	return v.URL, v.Token, v.SkipVerify, nil
}

func Read() Config {
	var c Config

	b, err := ioutil.ReadFile(saferc())
	if err != nil {
		return Config{Version: 1}
	}

	if err = yaml.Unmarshal(b, &c); err != nil {
		return Config{Version: 1}
	}
	if c.Version == 0 {
		var legacy oldConfig
		if err = yaml.Unmarshal(b, &legacy); err != nil {
			fmt.Fprintf(os.Stderr, "@R{!!! %s}\n", err)
			os.Exit(1)
		}
		c = legacy.convert()
	}

	return c
}

func Apply(use string) Config {
	c := Read()

	if err := c.Apply(use); err != nil {
		fmt.Fprintf(os.Stderr, "@R{!!! %s}\n", err)
		os.Exit(1)
	}
	return c
}

func (c *Config) Write() error {
	b, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

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

	v, err := c.Vault("")
	if err != nil {
		return err
	}
	if v == nil {
		os.Remove(svtoken())
		return nil
	}

	sv := struct {
		Vault      string `yaml:"vault"` /* this is different than Vault.URL */
		Token      string `yaml:"token"`
		SkipVerify bool   `yaml:"skip_verify"`
	}{v.URL, v.Token, v.SkipVerify}
	b, err = yaml.Marshal(sv)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(svtoken(), b, 0600)
}

func (c *Config) Apply(use string) error {
	v, err := c.Vault(use)
	if err != nil {
		fmt.Fprintf(os.Stderr, "@R{!!! %s}\n", err)
		os.Exit(1)
	}

	if v != nil {
		os.Setenv("VAULT_ADDR", v.URL)
		os.Setenv("VAULT_TOKEN", v.Token)
		if v.SkipVerify {
			os.Setenv("VAULT_SKIP_VERIFY", "1")
		}
	} else {
		if os.Getenv("VAULT_TOKEN") == "" {
			tokenFile := fmt.Sprintf("%s/.vault-token", os.Getenv("HOME"))
			b, err := ioutil.ReadFile(tokenFile)
			if err == nil {
				os.Setenv("VAULT_TOKEN", strings.TrimSpace(string(b)))
			}
		}
	}
	return nil
}

func (c *Config) SetCurrent(alias string, reskip bool) error {
	v, ok, err := c.Find(alias)
	if err != nil {
		return err
	}
	if !ok {
		return fmt.Errorf("Unknown target '%s'", alias)
	}
	c.Current = alias
	if reskip {
		v.SkipVerify = true
	}
	return nil
}

func (c *Config) SetTarget(alias, url string, skipverify bool) error {
	if c.Vaults == nil {
		c.Vaults = make(map[string]*Vault)
	}

	var token string

	c.Current = alias
	if existingAlias, found := c.Vaults[alias]; found {
		if url == existingAlias.URL {
			token = existingAlias.Token
		}
	}

	c.Vaults[alias] = &Vault{
		URL:        url,
		SkipVerify: skipverify,
		Token:      token,
	}

	return nil
}

func (c *Config) SetToken(token string) error {
	if c.Current == "" {
		return fmt.Errorf("No target selected")
	}
	v, ok, _ := c.Find(c.Current)
	if !ok {
		return fmt.Errorf("Unknown target '%s'", c.Current)
	}
	v.Token = token
	return nil
}

func (c *Config) URL() string {
	if v, ok, _ := c.Find(c.Current); ok {
		return v.URL
	}
	return ""
}

func (c *Config) Verified() bool {
	if v, ok, _ := c.Find(c.Current); ok {
		return !v.SkipVerify
	}
	return false
}

func (c *Config) Find(alias string) (*Vault, bool, error) {
	if v, ok := c.Vaults[alias]; ok {
		return v, true, nil
	}

	var v *Vault
	n := 0
	want := strings.TrimSuffix(alias, "/")

	for _, maybe := range c.Vaults {
		if strings.TrimSuffix(maybe.URL, "/") == want {
			n++
			v = maybe
		}
	}
	if n == 1 {
		return v, true, nil
	}
	if n > 1 {
		return nil, true, fmt.Errorf("More than one target for Vault at '%s' (maybe try an alias?)", alias)
	}

	return nil, false, nil
}

func (c *Config) Vault(which string) (*Vault, error) {
	if which == "" {
		which = c.Current
	}

	if which == "" {
		return nil, nil /* not an error */
	}

	v, ok, err := c.Find(which)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("Current target '%s' not found in ~/.saferc", which)
	}
	return v, nil
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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 contains functions for querying the aggregation results with hierarchical DPF key expansion.
package query

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"

	"golang.org/x/sync/errgroup"
	"github.com/pborman/uuid"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/cryptoio"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/dpfaggregator"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/ioutils"

	dpfpb "github.com/google/distributed_point_functions/dpf/distributed_point_function_go_proto"
	cryptopb "github.com/google/privacy-sandbox-aggregation-service/pipeline/crypto_go_proto"
	grpcpb "github.com/google/privacy-sandbox-aggregation-service/service/service_go_grpc_proto"
	servicepb "github.com/google/privacy-sandbox-aggregation-service/service/service_go_grpc_proto"
)

const elementBitSize = 64

// ExpansionConfig contains the parameters that define the hierarchy for how the DPF key will be expanded.
type ExpansionConfig struct {
	PrefixLengths               []int32
	ExpansionThresholdPerPrefix []uint64
}

// WriteExpansionConfigFile writes the ExpansionConfig into a file.
func WriteExpansionConfigFile(ctx context.Context, config *ExpansionConfig, filename string) error {
	bc, err := json.Marshal(config)
	if err != nil {
		return err
	}
	return ioutils.WriteBytes(ctx, bc, filename)
}

func validateExpansionConfig(config *ExpansionConfig) error {
	if len(config.PrefixLengths) == 0 {
		return errors.New("expect nonempty PrefixLengths")
	}

	if got, want := len(config.ExpansionThresholdPerPrefix), len(config.PrefixLengths); got != want {
		return fmt.Errorf("expect len(ExpansionThreshold)=%d, got %d", want, got)
	}
	var cur int32
	for _, l := range config.PrefixLengths {
		if l <= 0 {
			return fmt.Errorf("prefix length should be positive, got %d", l)
		}
		if l <= cur {
			return errors.New("prefix lengths should be in ascending order")
		}
		cur = l
	}
	return nil
}

// ReadExpansionConfigFile reads the ExpansionConfig from a file and validate it.
func ReadExpansionConfigFile(ctx context.Context, filename string) (*ExpansionConfig, error) {
	bc, err := ioutils.ReadBytes(ctx, filename)
	if err != nil {
		return nil, err
	}
	config := &ExpansionConfig{}
	if err := json.Unmarshal(bc, config); err != nil {
		return nil, err
	}
	return config, validateExpansionConfig(config)
}

type aggregateParams struct {
	PrefixesFile                                 string
	SumParamsFile                                string
	PartialHistogramFile1, PartialHistogramFile2 string
	PartialReportFile1, PartialReportFile2       string
}

func aggregateReports(ctx context.Context, params aggregateParams, client1, client2 grpcpb.AggregatorClient) error {
	g, ctx := errgroup.WithContext(ctx)
	g.Go(func() error {
		_, err := client1.AggregateDpfPartialReport(ctx, &servicepb.AggregateDpfPartialReportRequest{
			PartialReportFile:    params.PartialReportFile1,
			PartialHistogramFile: params.PartialHistogramFile1,
			PrefixesFile:         params.PrefixesFile,
			SumDpfParametersFile: params.SumParamsFile,
		})
		return err
	})

	g.Go(func() error {
		_, err := client2.AggregateDpfPartialReport(ctx, &servicepb.AggregateDpfPartialReportRequest{
			PartialReportFile:    params.PartialReportFile2,
			PartialHistogramFile: params.PartialHistogramFile2,
			PrefixesFile:         params.PrefixesFile,
			SumDpfParametersFile: params.SumParamsFile,
		})
		return err
	})

	return g.Wait()
}

// PrefixHistogramParams contains the parameters for querying the histogram of given prefixes.
type PrefixHistogramParams struct {
	Prefixes                               *cryptopb.HierarchicalPrefixes
	SumParams                              *cryptopb.IncrementalDpfParameters
	PartialReportFile1, PartialReportFile2 string
	PartialAggregationDir                  string
	ParamsDir                              string
	Helper1, Helper2                       grpcpb.AggregatorClient
}

// getPrefixHistogram calls the RPC methods on both helpers and merges the generated partial aggregation results.
func getPrefixHistogram(ctx context.Context, params *PrefixHistogramParams) ([]dpfaggregator.CompleteHistogram, error) {
	tempID := uuid.New()
	prefixFile := fmt.Sprintf("%s/prefixes%s.txt", params.ParamsDir, tempID)
	if err := cryptoio.SavePrefixes(ctx, prefixFile, params.Prefixes); err != nil {
		return nil, err
	}
	sumParamsFile := fmt.Sprintf("%s/sum_params%s.txt", params.ParamsDir, tempID)
	if err := cryptoio.SaveDPFParameters(ctx, sumParamsFile, params.SumParams); err != nil {
		return nil, err
	}

	tempPartialResultFile1 := fmt.Sprintf("%s/%s_1.txt", params.PartialAggregationDir, tempID)
	tempPartialResultFile2 := fmt.Sprintf("%s/%s_2.txt", params.PartialAggregationDir, tempID)

	if err := aggregateReports(ctx, aggregateParams{
		PrefixesFile:          prefixFile,
		SumParamsFile:         sumParamsFile,
		PartialHistogramFile1: tempPartialResultFile1,
		PartialHistogramFile2: tempPartialResultFile2,
	}, params.Helper1, params.Helper2); err != nil {
		return nil, err
	}
	partial1, err := dpfaggregator.ReadPartialHistogram(ctx, tempPartialResultFile1)
	if err != nil {
		return nil, err
	}
	partial2, err := dpfaggregator.ReadPartialHistogram(ctx, tempPartialResultFile2)
	if err != nil {
		return nil, err
	}

	for idx := range partial2 {
		if _, ok := partial1[idx]; !ok {
			return nil, fmt.Errorf("index %d appears in partial2, missing in partial1", idx)
		}
	}

	var result []dpfaggregator.CompleteHistogram
	for idx := range partial1 {
		if _, ok := partial2[idx]; !ok {
			return nil, fmt.Errorf("index %d appears in partial1, missing in partial2", idx)
		}
		result = append(result, dpfaggregator.CompleteHistogram{
			Index: idx,
			Sum:   partial1[idx].PartialSum + partial2[idx].PartialSum,
		})
	}
	return result, nil
}

func getNextNonemptyPrefixes(result []dpfaggregator.CompleteHistogram, threshold uint64) []uint64 {
	var prefixes []uint64
	for _, r := range result {
		if r.Sum >= threshold {
			prefixes = append(prefixes, r.Index)
		}
	}
	return prefixes
}

func extendPrefixDomains(sumParams *cryptopb.IncrementalDpfParameters, prefixLength int32) {
	sumParams.Params = append(sumParams.Params, &dpfpb.DpfParameters{
		LogDomainSize:  prefixLength,
		ElementBitsize: elementBitSize,
	})
}

// HierarchicalResult records the aggregation result at certain prefix length.
//
// TODO: Add PrivacyBudgetConsumed field
type HierarchicalResult struct {
	PrefixLength int32
	Histogram    []dpfaggregator.CompleteHistogram
	// The threshold applied to the above result, which generates prefixes for the next-level expansion.
	ExpansionThreshold uint64
}

// HierarchicalAggregation queries the hierarchical aggregation results.
func HierarchicalAggregation(ctx context.Context, params *PrefixHistogramParams, config *ExpansionConfig) ([]HierarchicalResult, error) {
	var results []HierarchicalResult
	for i, threshold := range config.ExpansionThresholdPerPrefix {
		extendPrefixDomains(params.SumParams, config.PrefixLengths[i])
		result, err := getPrefixHistogram(ctx, params)
		if err != nil {
			return nil, err
		}
		params.Prefixes.Prefixes = append(params.Prefixes.Prefixes, &cryptopb.DomainPrefixes{Prefix: getNextNonemptyPrefixes(result, threshold)})
		results = append(results, HierarchicalResult{PrefixLength: config.PrefixLengths[i], Histogram: result, ExpansionThreshold: threshold})
	}
	return results, nil
}

// WriteHierarchicalResultsFile writes the hierarchical query results into a file.
func WriteHierarchicalResultsFile(results []HierarchicalResult, filename string) error {
	br, err := json.Marshal(results)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, br, os.ModePerm)
}

// ReadHierarchicalResultsFile reads the hierarchical query results from a file.
func ReadHierarchicalResultsFile(filename string) ([]HierarchicalResult, error) {
	br, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	var results []HierarchicalResult
	if err := json.Unmarshal(br, &results); err != nil {
		return nil, err
	}
	return results, nil
}
<commit_msg>Add the missing parameters for aggregating partial reports.<commit_after>// Copyright 2021 Google LLC
//
// 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 contains functions for querying the aggregation results with hierarchical DPF key expansion.
package query

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"

	"golang.org/x/sync/errgroup"
	"github.com/pborman/uuid"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/cryptoio"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/dpfaggregator"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/ioutils"

	dpfpb "github.com/google/distributed_point_functions/dpf/distributed_point_function_go_proto"
	cryptopb "github.com/google/privacy-sandbox-aggregation-service/pipeline/crypto_go_proto"
	grpcpb "github.com/google/privacy-sandbox-aggregation-service/service/service_go_grpc_proto"
	servicepb "github.com/google/privacy-sandbox-aggregation-service/service/service_go_grpc_proto"
)

const elementBitSize = 64

// ExpansionConfig contains the parameters that define the hierarchy for how the DPF key will be expanded.
type ExpansionConfig struct {
	PrefixLengths               []int32
	ExpansionThresholdPerPrefix []uint64
}

// WriteExpansionConfigFile writes the ExpansionConfig into a file.
func WriteExpansionConfigFile(ctx context.Context, config *ExpansionConfig, filename string) error {
	bc, err := json.Marshal(config)
	if err != nil {
		return err
	}
	return ioutils.WriteBytes(ctx, bc, filename)
}

func validateExpansionConfig(config *ExpansionConfig) error {
	if len(config.PrefixLengths) == 0 {
		return errors.New("expect nonempty PrefixLengths")
	}

	if got, want := len(config.ExpansionThresholdPerPrefix), len(config.PrefixLengths); got != want {
		return fmt.Errorf("expect len(ExpansionThreshold)=%d, got %d", want, got)
	}
	var cur int32
	for _, l := range config.PrefixLengths {
		if l <= 0 {
			return fmt.Errorf("prefix length should be positive, got %d", l)
		}
		if l <= cur {
			return errors.New("prefix lengths should be in ascending order")
		}
		cur = l
	}
	return nil
}

// ReadExpansionConfigFile reads the ExpansionConfig from a file and validate it.
func ReadExpansionConfigFile(ctx context.Context, filename string) (*ExpansionConfig, error) {
	bc, err := ioutils.ReadBytes(ctx, filename)
	if err != nil {
		return nil, err
	}
	config := &ExpansionConfig{}
	if err := json.Unmarshal(bc, config); err != nil {
		return nil, err
	}
	return config, validateExpansionConfig(config)
}

type aggregateParams struct {
	PrefixesFile                                 string
	SumParamsFile                                string
	PartialHistogramFile1, PartialHistogramFile2 string
	PartialReportFile1, PartialReportFile2       string
}

func aggregateReports(ctx context.Context, params aggregateParams, client1, client2 grpcpb.AggregatorClient) error {
	g, ctx := errgroup.WithContext(ctx)
	g.Go(func() error {
		_, err := client1.AggregateDpfPartialReport(ctx, &servicepb.AggregateDpfPartialReportRequest{
			PartialReportFile:    params.PartialReportFile1,
			PartialHistogramFile: params.PartialHistogramFile1,
			PrefixesFile:         params.PrefixesFile,
			SumDpfParametersFile: params.SumParamsFile,
		})
		return err
	})

	g.Go(func() error {
		_, err := client2.AggregateDpfPartialReport(ctx, &servicepb.AggregateDpfPartialReportRequest{
			PartialReportFile:    params.PartialReportFile2,
			PartialHistogramFile: params.PartialHistogramFile2,
			PrefixesFile:         params.PrefixesFile,
			SumDpfParametersFile: params.SumParamsFile,
		})
		return err
	})

	return g.Wait()
}

// PrefixHistogramParams contains the parameters for querying the histogram of given prefixes.
type PrefixHistogramParams struct {
	Prefixes                               *cryptopb.HierarchicalPrefixes
	SumParams                              *cryptopb.IncrementalDpfParameters
	PartialReportFile1, PartialReportFile2 string
	PartialAggregationDir                  string
	ParamsDir                              string
	Helper1, Helper2                       grpcpb.AggregatorClient
}

// getPrefixHistogram calls the RPC methods on both helpers and merges the generated partial aggregation results.
func getPrefixHistogram(ctx context.Context, params *PrefixHistogramParams) ([]dpfaggregator.CompleteHistogram, error) {
	tempID := uuid.New()
	prefixFile := fmt.Sprintf("%s/prefixes%s.txt", params.ParamsDir, tempID)
	if err := cryptoio.SavePrefixes(ctx, prefixFile, params.Prefixes); err != nil {
		return nil, err
	}
	sumParamsFile := fmt.Sprintf("%s/sum_params%s.txt", params.ParamsDir, tempID)
	if err := cryptoio.SaveDPFParameters(ctx, sumParamsFile, params.SumParams); err != nil {
		return nil, err
	}

	tempPartialResultFile1 := fmt.Sprintf("%s/%s_1.txt", params.PartialAggregationDir, tempID)
	tempPartialResultFile2 := fmt.Sprintf("%s/%s_2.txt", params.PartialAggregationDir, tempID)

	if err := aggregateReports(ctx, aggregateParams{
		PrefixesFile:          prefixFile,
		SumParamsFile:         sumParamsFile,
		PartialReportFile1:    params.PartialReportFile1,
		PartialReportFile2:    params.PartialReportFile2,
		PartialHistogramFile1: tempPartialResultFile1,
		PartialHistogramFile2: tempPartialResultFile2,
	}, params.Helper1, params.Helper2); err != nil {
		return nil, err
	}
	partial1, err := dpfaggregator.ReadPartialHistogram(ctx, tempPartialResultFile1)
	if err != nil {
		return nil, err
	}
	partial2, err := dpfaggregator.ReadPartialHistogram(ctx, tempPartialResultFile2)
	if err != nil {
		return nil, err
	}

	for idx := range partial2 {
		if _, ok := partial1[idx]; !ok {
			return nil, fmt.Errorf("index %d appears in partial2, missing in partial1", idx)
		}
	}

	var result []dpfaggregator.CompleteHistogram
	for idx := range partial1 {
		if _, ok := partial2[idx]; !ok {
			return nil, fmt.Errorf("index %d appears in partial1, missing in partial2", idx)
		}
		result = append(result, dpfaggregator.CompleteHistogram{
			Index: idx,
			Sum:   partial1[idx].PartialSum + partial2[idx].PartialSum,
		})
	}
	return result, nil
}

func getNextNonemptyPrefixes(result []dpfaggregator.CompleteHistogram, threshold uint64) []uint64 {
	var prefixes []uint64
	for _, r := range result {
		if r.Sum >= threshold {
			prefixes = append(prefixes, r.Index)
		}
	}
	return prefixes
}

func extendPrefixDomains(sumParams *cryptopb.IncrementalDpfParameters, prefixLength int32) {
	sumParams.Params = append(sumParams.Params, &dpfpb.DpfParameters{
		LogDomainSize:  prefixLength,
		ElementBitsize: elementBitSize,
	})
}

// HierarchicalResult records the aggregation result at certain prefix length.
//
// TODO: Add PrivacyBudgetConsumed field
type HierarchicalResult struct {
	PrefixLength int32
	Histogram    []dpfaggregator.CompleteHistogram
	// The threshold applied to the above result, which generates prefixes for the next-level expansion.
	ExpansionThreshold uint64
}

// HierarchicalAggregation queries the hierarchical aggregation results.
func HierarchicalAggregation(ctx context.Context, params *PrefixHistogramParams, config *ExpansionConfig) ([]HierarchicalResult, error) {
	var results []HierarchicalResult
	for i, threshold := range config.ExpansionThresholdPerPrefix {
		extendPrefixDomains(params.SumParams, config.PrefixLengths[i])
		result, err := getPrefixHistogram(ctx, params)
		if err != nil {
			return nil, err
		}
		params.Prefixes.Prefixes = append(params.Prefixes.Prefixes, &cryptopb.DomainPrefixes{Prefix: getNextNonemptyPrefixes(result, threshold)})
		results = append(results, HierarchicalResult{PrefixLength: config.PrefixLengths[i], Histogram: result, ExpansionThreshold: threshold})
	}
	return results, nil
}

// WriteHierarchicalResultsFile writes the hierarchical query results into a file.
func WriteHierarchicalResultsFile(results []HierarchicalResult, filename string) error {
	br, err := json.Marshal(results)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, br, os.ModePerm)
}

// ReadHierarchicalResultsFile reads the hierarchical query results from a file.
func ReadHierarchicalResultsFile(filename string) ([]HierarchicalResult, error) {
	br, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	var results []HierarchicalResult
	if err := json.Unmarshal(br, &results); err != nil {
		return nil, err
	}
	return results, nil
}
<|endoftext|>