<commit_before>// repmbox handles repbin mailboxes.
//
// Release 20150603.
// You can reach me with
// repclient --recipientPubKey 7VW3oPLzQc7VS2anLyDtrdARDdSwa7QTF7h3N2t6J2VN_DTshmJFEDa7XM2w9nHBq4CgtvK4kYdBp8G3wFPkYcGd1
// Don't forget to put your own key into your message! (create with repmbox -add)
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strconv"
	"strings"
)

type sender struct {
	Sender     string
	PublicKey  string
	PrivateKey string
}

type config struct {
	PrivateKey string
	Server     string
	Start      int
	Count      int
	Sender     []sender
}

type senderKey struct {
	user    string
	privKey string
}

func genConfig() (*config, error) {
	// generate private key with repclient
	cmd := exec.Command("repclient", "--genkey", "--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	var cfg config
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			cfg.PrivateKey = parts[1]
			break
		}
	}
	if cfg.PrivateKey == "" {
		return nil, fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	cfg.Start = -1
	cfg.Count = 100
	return &cfg, nil
}

func loadConfig(configFile, configDir string) (*config, error) {
	if _, err := os.Stat(configFile); err != nil {
		return nil, fmt.Errorf("config file '%s' doesn't exist, generate with:\n"+
			"mkdir -p %s && repmbox --genconfig > %s", configFile, configDir, configFile)
	}
	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, err
	}
	var cfg config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	// check server entry
	if cfg.Server == "" {
		return nil, fmt.Errorf("specify Server entry in config file '%s'", configFile)
	}
	return &cfg, nil
}

func (cfg *config) save(configFile string) error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(configFile, jsn, 0600); err != nil {
		return err
	}
	return nil
}

func (cfg *config) show() error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	fmt.Println(string(jsn))
	return nil
}

func (cfg *config) addSender(addSender, configFile string) error {
	var snd sender
	snd.Sender = addSender
	// generate key pair for sender
	cmd := exec.Command("repclient", "--gentemp",
		"--privkey", cfg.PrivateKey,
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return err
	}
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			snd.PrivateKey = parts[1]
		}
		if len(parts) >= 2 && parts[0] == "STATUS (PublicKey):" {
			snd.PublicKey = parts[1]
			break
		}
	}
	if snd.PrivateKey == "" {
		return fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	if snd.PublicKey == "" {
		return fmt.Errorf("could not parse public key from output:\n%s", out.String())
	}
	cfg.Sender = append(cfg.Sender, snd)
	if err := cfg.save(configFile); err != nil {
		return err
	}
	fmt.Printf("mailbox for %s added:\n", addSender)
	fmt.Println(snd.PublicKey)
	return nil
}

func (cfg *config) getList() ([]string, bool, error) {
	cmd := exec.Command("repclient",
		"--index",
		"--server", cfg.Server,
		"--privkey", cfg.PrivateKey,
		"--start", strconv.Itoa(cfg.Start+1),
		"--count", strconv.Itoa(cfg.Count),
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		fmt.Print(out.String())
		return nil, false, err
	}
	var list []string
	var more bool
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 {
			if parts[0] == "STATUS (MessageList):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 2 {
					list = append(list, particles[1])
				} else {
					return nil, false, fmt.Errorf("could not parse line: %s", line)
				}
			} else if parts[0] == "STATUS (ListResult):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 3 {
					var err error
					more, err = strconv.ParseBool(particles[2])
					if err != nil {
						return nil, false, err
					}
					break
				}
			}
		}
	}
	return list, more, nil
}

func (cfg *config) getSenderKeys() map[string]senderKey {
	keys := make(map[string]senderKey)
	for _, sdr := range cfg.Sender {
		pubKeys := strings.Split(sdr.PublicKey, "_")
		privKeys := strings.Split(sdr.PrivateKey, "_")
		keys[pubKeys[0]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[0],
		}
		keys[pubKeys[1]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[1],
		}
	}
	return keys
}

func (cfg *config) getMessages(list []string, outdir, stmdir string, verbose bool) error {
	// make sure STM directory exists
	if err := os.MkdirAll(stmdir, 0700); err != nil {
		return err
	}
	for _, msg := range list {
		fmt.Printf("retrieving message %s\n", msg)
		cmd := exec.Command("repclient",
			"--decrypt",
			"--keymgt", "3",
			"--server", cfg.Server,
			"--stmdir", stmdir,
			"--appdata",
			msg)
		var out bytes.Buffer
		cmd.Stdout = &out
		stderr, err := cmd.StderrPipe()
		if err != nil {
			return err
		}
		r, w, err := os.Pipe()
		if err != nil {
			return err
		}
		cmd.ExtraFiles = append(cmd.ExtraFiles, r)
		if err := cmd.Start(); err != nil {
			return err
		}
		var user string
		var stmfile string
		keys := cfg.getSenderKeys()
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			line := scanner.Text()
			if verbose {
				fmt.Println(line)
			}
			parts := strings.Split(line, "\t")
			if len(parts) >= 2 {
				if parts[0] == "STATUS (KeyMGTRequest):" {
					key, ok := keys[parts[1]]
					if !ok {
						return fmt.Errorf("key for %s not found", parts[1])
					}
					if _, err := w.WriteString(key.privKey + "\n"); err != nil {
						return err
					}
					user = key.user
				} else if parts[0] == "STATUS (STMFile):" {
					stmfile = parts[1]
				}
			} else {
				return fmt.Errorf("could not parse repclient output: %s", line)
			}
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		if err := cmd.Wait(); err != nil {
			return fmt.Errorf("repclient call failed with: %s", err)
		}
		if stmfile != "" {
			fmt.Printf("new repost message from %s written to %s\n", user, stmfile)
		} else {
			// make sure output directory exists
			resdir := path.Join(outdir, user)
			if err := os.MkdirAll(resdir, 0700); err != nil {
				return err
			}
			// write message
			filename := path.Join(resdir, msg)
			if err := ioutil.WriteFile(path.Join(resdir, msg), out.Bytes(), 0600); err != nil {
				return err
			}
			fmt.Printf("new message from %s written to %s\n", user, filename)
		}
	}
	return nil
}

func appDataDir(appName string) (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return path.Join(user.HomeDir, ".config", appName), nil
}

func fatal(err error) {
	fmt.Fprintf(os.Stderr, "%s: error: %s\n", os.Args[0], err)
	os.Exit(1)
}

func main() {
	configDir, err := appDataDir("repmbox")
	if err != nil {
		fatal(err)
	}
	// define options
	addSender := flag.String("add", "", "Add new mailbox for given sender name")
	configFile := flag.String("configfile", path.Join(configDir, "repmbox.config"), "Path to configuration file")
	genconfig := flag.Bool("genconfig", false, "Generate config file")
	listSender := flag.Bool("list", false, "List mailboxes")
	outdir := flag.String("outdir", path.Join(configDir, "messages"), "Directory for downloaded messages")
	stmdir := flag.String("stmdir", path.Join(configDir, "stmdir"), "Directory for STM messages")
	verbose := flag.Bool("v", false, "Be verbose")
	// parse options
	flag.Parse()
	if *genconfig {
		cfg, err := genConfig()
		if err != nil {
			fatal(err)
		}
		if err := cfg.show(); err != nil {
			fatal(err)
		}
		return
	}
	// load config file
	cfg, err := loadConfig(*configFile, configDir)
	if err != nil {
		fatal(err)
	}
	if *addSender != "" {
		if err := cfg.addSender(*addSender, *configFile); err != nil {
			fatal(err)
		}
	} else if *listSender {
		for _, snd := range cfg.Sender {
			fmt.Println(snd.Sender, snd.PublicKey)
		}
	} else {
		more := true
		for more {
			var list []string
			// download new messages
			list, more, err = cfg.getList()
			if err != nil {
				fatal(err)
			}
			// get new messages
			if err := cfg.getMessages(list, *outdir, *stmdir, *verbose); err != nil {
				fatal(err)
			}
			// increase start
			cfg.Start += len(list)
			if err := cfg.save(*configFile); err != nil {
				fatal(err)
			}
			if *verbose {
				fmt.Printf("more: %s\n", strconv.FormatBool(more))
			}
		}
	}
}
<commit_msg>handle unavailable keys<commit_after>// repmbox handles repbin mailboxes.
//
// Release 20150605.
// You can reach me with
// repclient --recipientPubKey 7VW3oPLzQc7VS2anLyDtrdARDdSwa7QTF7h3N2t6J2VN_DTshmJFEDa7XM2w9nHBq4CgtvK4kYdBp8G3wFPkYcGd1
// Don't forget to put your own key into your message! (create with repmbox -add)
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strconv"
	"strings"
)

type sender struct {
	Sender     string
	PublicKey  string
	PrivateKey string
}

type config struct {
	PrivateKey string
	Server     string
	Start      int
	Count      int
	Sender     []sender
}

type senderKey struct {
	user    string
	privKey string
}

func genConfig() (*config, error) {
	// generate private key with repclient
	cmd := exec.Command("repclient", "--genkey", "--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	var cfg config
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			cfg.PrivateKey = parts[1]
			break
		}
	}
	if cfg.PrivateKey == "" {
		return nil, fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	cfg.Start = -1
	cfg.Count = 100
	return &cfg, nil
}

func loadConfig(configFile, configDir string) (*config, error) {
	if _, err := os.Stat(configFile); err != nil {
		return nil, fmt.Errorf("config file '%s' doesn't exist, generate with:\n"+
			"mkdir -p %s && repmbox --genconfig > %s", configFile, configDir, configFile)
	}
	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, err
	}
	var cfg config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	// check server entry
	if cfg.Server == "" {
		return nil, fmt.Errorf("specify Server entry in config file '%s'", configFile)
	}
	return &cfg, nil
}

func (cfg *config) save(configFile string) error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(configFile, jsn, 0600); err != nil {
		return err
	}
	return nil
}

func (cfg *config) show() error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	fmt.Println(string(jsn))
	return nil
}

func (cfg *config) addSender(addSender, configFile string) error {
	var snd sender
	snd.Sender = addSender
	// generate key pair for sender
	cmd := exec.Command("repclient", "--gentemp",
		"--privkey", cfg.PrivateKey,
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return err
	}
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			snd.PrivateKey = parts[1]
		}
		if len(parts) >= 2 && parts[0] == "STATUS (PublicKey):" {
			snd.PublicKey = parts[1]
			break
		}
	}
	if snd.PrivateKey == "" {
		return fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	if snd.PublicKey == "" {
		return fmt.Errorf("could not parse public key from output:\n%s", out.String())
	}
	cfg.Sender = append(cfg.Sender, snd)
	if err := cfg.save(configFile); err != nil {
		return err
	}
	fmt.Printf("mailbox for %s added:\n", addSender)
	fmt.Println(snd.PublicKey)
	return nil
}

func (cfg *config) getList() ([]string, bool, error) {
	cmd := exec.Command("repclient",
		"--index",
		"--server", cfg.Server,
		"--privkey", cfg.PrivateKey,
		"--start", strconv.Itoa(cfg.Start+1),
		"--count", strconv.Itoa(cfg.Count),
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		fmt.Print(out.String())
		return nil, false, err
	}
	var list []string
	var more bool
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 {
			if parts[0] == "STATUS (MessageList):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 2 {
					list = append(list, particles[1])
				} else {
					return nil, false, fmt.Errorf("could not parse line: %s", line)
				}
			} else if parts[0] == "STATUS (ListResult):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 3 {
					var err error
					more, err = strconv.ParseBool(particles[2])
					if err != nil {
						return nil, false, err
					}
					break
				}
			}
		}
	}
	return list, more, nil
}

func (cfg *config) getSenderKeys() map[string]senderKey {
	keys := make(map[string]senderKey)
	for _, sdr := range cfg.Sender {
		pubKeys := strings.Split(sdr.PublicKey, "_")
		privKeys := strings.Split(sdr.PrivateKey, "_")
		keys[pubKeys[0]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[0],
		}
		keys[pubKeys[1]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[1],
		}
	}
	return keys
}

func (cfg *config) getMessages(list []string, outdir, stmdir string, verbose bool) error {
	// make sure STM directory exists
	if err := os.MkdirAll(stmdir, 0700); err != nil {
		return err
	}
	for _, msg := range list {
		fmt.Printf("retrieving message %s\n", msg)
		cmd := exec.Command("repclient",
			"--decrypt",
			"--keymgt", "3",
			"--server", cfg.Server,
			"--stmdir", stmdir,
			"--appdata",
			msg)
		var out bytes.Buffer
		cmd.Stdout = &out
		stderr, err := cmd.StderrPipe()
		if err != nil {
			return err
		}
		r, w, err := os.Pipe()
		if err != nil {
			return err
		}
		cmd.ExtraFiles = append(cmd.ExtraFiles, r)
		if err := cmd.Start(); err != nil {
			return err
		}
		var user string
		var stmfile string
		var keyNotAvailable bool
		keys := cfg.getSenderKeys()
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			line := scanner.Text()
			if verbose {
				fmt.Println(line)
			}
			parts := strings.Split(line, "\t")
			if len(parts) >= 2 {
				if parts[0] == "STATUS (KeyMGTRequest):" {
					key, ok := keys[parts[1]]
					if !ok {
						if _, err := w.WriteString("not available\n"); err != nil {
							return err
						}
						keyNotAvailable = true
						break
					} else {
						if _, err := w.WriteString(key.privKey + "\n"); err != nil {
							return err
						}
						user = key.user
					}
				} else if parts[0] == "STATUS (STMFile):" {
					stmfile = parts[1]
				}
			} else {
				return fmt.Errorf("could not parse repclient output: %s", line)
			}
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		err = cmd.Wait()
		if !keyNotAvailable && err != nil {
			return fmt.Errorf("repclient call failed with: %s", err)
		}
		if stmfile != "" {
			fmt.Printf("new repost message from %s written to %s\n", user, stmfile)
		} else if keyNotAvailable {
			fmt.Println("key not available, skipping message forever!")
		} else {
			// make sure output directory exists
			resdir := path.Join(outdir, user)
			if err := os.MkdirAll(resdir, 0700); err != nil {
				return err
			}
			// write message
			filename := path.Join(resdir, msg)
			if err := ioutil.WriteFile(path.Join(resdir, msg), out.Bytes(), 0600); err != nil {
				return err
			}
			fmt.Printf("new message from %s written to %s\n", user, filename)
		}
	}
	return nil
}

func appDataDir(appName string) (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return path.Join(user.HomeDir, ".config", appName), nil
}

func fatal(err error) {
	fmt.Fprintf(os.Stderr, "%s: error: %s\n", os.Args[0], err)
	os.Exit(1)
}

func main() {
	configDir, err := appDataDir("repmbox")
	if err != nil {
		fatal(err)
	}
	// define options
	addSender := flag.String("add", "", "Add new mailbox for given sender name")
	configFile := flag.String("configfile", path.Join(configDir, "repmbox.config"), "Path to configuration file")
	genconfig := flag.Bool("genconfig", false, "Generate config file")
	listSender := flag.Bool("list", false, "List mailboxes")
	outdir := flag.String("outdir", path.Join(configDir, "messages"), "Directory for downloaded messages")
	stmdir := flag.String("stmdir", path.Join(configDir, "stmdir"), "Directory for STM messages")
	verbose := flag.Bool("v", false, "Be verbose")
	// parse options
	flag.Parse()
	if *genconfig {
		cfg, err := genConfig()
		if err != nil {
			fatal(err)
		}
		if err := cfg.show(); err != nil {
			fatal(err)
		}
		return
	}
	// load config file
	cfg, err := loadConfig(*configFile, configDir)
	if err != nil {
		fatal(err)
	}
	if *addSender != "" {
		if err := cfg.addSender(*addSender, *configFile); err != nil {
			fatal(err)
		}
	} else if *listSender {
		for _, snd := range cfg.Sender {
			fmt.Println(snd.Sender, snd.PublicKey)
		}
	} else {
		more := true
		for more {
			var list []string
			// download new messages
			list, more, err = cfg.getList()
			if err != nil {
				fatal(err)
			}
			// get new messages
			if err := cfg.getMessages(list, *outdir, *stmdir, *verbose); err != nil {
				fatal(err)
			}
			// increase start
			cfg.Start += len(list)
			if err := cfg.save(*configFile); err != nil {
				fatal(err)
			}
			if *verbose {
				fmt.Printf("more: %s\n", strconv.FormatBool(more))
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
  A partitioner assigns partitions within a consumer group like a
  consistent hash would. When group membership changes the least
  number of partitions are reassigned.

  I don't actually use a consistent hash because such algos take
  either a lot of CPU (the max(N-hash-functions) algo) or a lot
  of ram (the points-on-the-circle algo), or aren't applicable
  when members come and go (jump-hash).

  Instead I take advantage of the fact that I can know the current
  assignments of the remaining group members when constructing
  the new assignments. That makes keep the partitioning stable
  relatively easy to do:
    1) calculate an ideal load for each consumer (I assume equal
	   weight since I don't need anything more complicated...yet)
	2) remove excess partitions from overloaded consumers
	3) add the abandonned and excess partitions to the least loaded
	   consumers.

  Note that the state of the partition at time T depends on
  the history of the consumer group memberships from the time when
  there was 1 consumer until time T.

  Copyright 2016 MistSys
*/

package stable

import (
	"fmt"
	"log"
	"sort"

	"github.com/Shopify/sarama"
)

// a partitioner that assigns partitions to consumers such that as consumers come and go the least number of partitions are reassigned
type stablePartitioner string

// global instance of the stable partitioner
const Stable stablePartitioner = "stable"

// print a debug message
func dbgf(format string, args ...interface{}) {
	log.Printf(format, args...)
}

func (sp stablePartitioner) PrepareJoin(jreq *sarama.JoinGroupRequest, topics []string, current_assignments map[string][]int32) {
	// encode the current assignments in a manner proprietary to this partitioner
	var data = data{
		version:     1,
		assignments: current_assignments,
	}

	jreq.AddGroupProtocolMetadata(string(sp),
		&sarama.ConsumerGroupMemberMetadata{
			Version:  1,
			Topics:   topics,
			UserData: data.marshal(),
		})
}

// for each topic in jresp, assign the topic's partitions to the members requesting the topic
func (sp stablePartitioner) Partition(sreq *sarama.SyncGroupRequest, jresp *sarama.JoinGroupResponse, client sarama.Client) error {
	if jresp.GroupProtocol != string(sp) {
		return fmt.Errorf("sarama.JoinGroupResponse.GroupProtocol %q unexpected; expected %q", jresp.GroupProtocol != string(sp))
	}
	by_member, err := jresp.GetMembers() //  map of member to ConsumerGroupMemberMetadata
	dbgf("by_member = %v", by_member)
	if err != nil {
		return err
	}

	// invert the data, so we have the requests grouped by topic (they arrived grouped by member, since the kafka broker treats the data from each consumer as an opaque blob, so it couldn't do this step for us)
	by_topic := make(map[string]map[string][]int32) // map of topic to members and members to current partition assignment
	for member, request := range by_member {
		if request.Version != 1 {
			// skip unsupported versions. we'll only assign to clients we can understand. Since we are such a client
			// we won't block all consumers (at least for those topics we consume). If this ends up a bad idea, we
			// can always change this code to return an error.
			continue
		}
		var data data
		err := data.unmarshal(request.UserData)
		if err != nil {
			return err
		}

		for _, topic := range request.Topics {
			members, ok := by_topic[topic]
			if !ok {
				members = make(map[string][]int32)
				by_topic[topic] = members
			}
			members[member] = data.assignments[topic] // NOTE: might be nil, which is OK. It just means the member wants to consume the partition but isn't doing so currently
		}
	}
	dbgf("by_topic = %v", by_topic)

	// lookup the partitions in each topic. since we are asking for all partitions, not just the online ones, the numbering
	// appears to always be 0...N-1. But in case I don't understand kafka and there is some corner case where the numbering
	// of partitions is different I keep careful track of the exact numbers I've received.
	var partitions_by_topic = make(map[string]partitionslist)
	for topic := range by_topic {
		// note: calls to client.Partitions() hit the metadata cache in the sarama client, so we don't gain much by asking concurrently
		partitions, err := client.Partitions(topic)
		if err != nil {
			// what to do? we could maybe skip the topic, assigning it to no-one. But I/O errors are likely to happen again.
			// so let's stop partitioning and return the error.
			return err
		}
		pl := make(partitionslist, len(partitions)) // we must make a copy before we sort the list
		for i, p := range partitions {
			pl[i] = p
		}
		sort.Sort(pl)
		partitions_by_topic[topic] = pl
	}
	dbgf("partitions_by_topic = %v", partitions_by_topic)

	// and compute the sorted set of members of each topic
	var members_by_topic = make(map[string]memberslist)
	for topic, members := range by_topic {
		ml := make(memberslist, 0, len(members))
		for m := range members {
			ml = append(ml, m)
		}
		sort.Sort(ml)
		members_by_topic[topic] = ml
	}
	dbgf("members_by_topic = %v", members_by_topic)

	// I want topics with the same # of partitions and the same consumer group membership to result in the same partition assignments.
	// That way messages published under identical partition keys in those topics will all end up consumed by the same member.
	// So organize topics into groups which will be partitioned identically
	var matched_topics = make(map[string]string) // map from each topic to the 'master' topic with the same # of partitions and group membership. Topics which are unique are their own master
	for topic, members := range members_by_topic {
		// see if a match exists
		num_partitions := len(partitions_by_topic[topic])
		for t := range matched_topics { // TODO if # of topics gets large enough this shows up in the profiler, change this to some sort of a map lookup, rather than this O(N^2) search
			if num_partitions == len(partitions_by_topic[t]) && members.Equal(members_by_topic[t]) {
				// match; have topic 'topic' be partitioned the same as topic 't'
				matched_topics[topic] = matched_topics[t]
				continue
			}
			// no existing topic matches this one, so it is its own match
			matched_topics[topic] = topic
		}
	}
	dbgf("matched_topics = %v", matched_topics)

	// adjust the partitioning of each master topic in by_topic
	for topic, match := range matched_topics {
		if topic == match {
			adjust_partitioning(by_topic[topic], partitions_by_topic[topic])
		} // else it is not a master topic. once the master has been partitioned we'll simply copy the result
	}

	// set matched topics to the same assignment as their master
	for topic, match := range matched_topics {
		if topic != match {
			by_topic[topic] = by_topic[match]
		}
	}

	// invert by_topic into the equivalent organized by member, and then by topic
	assignments := make(map[string]map[string][]int32) // map of member to topics, and topic to partitions
	for topic, members := range by_topic {
		for member, partitions := range members {
			topics, ok := assignments[member]
			if !ok {
				topics = make(map[string][]int32)
				assignments[member] = topics
			}
			topics[topic] = partitions
		}
	}
	dbgf("assignments = %v", assignments)

	// and encode the assignments in the sync request
	for member, topics := range assignments {
		sreq.AddGroupAssignmentMember(member,
			&sarama.ConsumerGroupMemberAssignment{
				Version: 1,
				Topics:  topics,
			})
	}

	return nil
}

func (stablePartitioner) ParseSync(sresp *sarama.SyncGroupResponse) (map[string][]int32, error) {
	if len(sresp.MemberAssignment) == 0 {
		// in the corner case that we ask for no topics, we get nothing back. However sarama fd498173ae2bf (head of master branch Nov 6th 2016) will return a useless error if we call sresp.GetMemberAssignment() in this case
		return nil, nil
	}
	ma, err := sresp.GetMemberAssignment()
	dbgf("MemberAssignment %v", ma)
	if err != nil {
		return nil, err
	}
	if ma.Version != 1 {
		return nil, fmt.Errorf("unsupported MemberAssignment version %d", ma.Version)
	}
	return ma.Topics, nil
}

// ----------------------------------

// adjust_partitioning does the main work. it adjusts the partition assignment map it is passed in-place
func adjust_partitioning(assignment map[string][]int32, partitions partitionslist) {
	dbgf("adjust_partitioning(assignment = %v, partitions = %v)", assignment, partitions)
	num_members := len(assignment)
	dbgf("num_members = %v", num_members)
	if num_members == 0 {
		// no one wants this topic; stop now before we /0
		return
	}
	num_partitions := len(partitions)
	dbgf("num_partitions = %v", num_partitions)
	level := (num_partitions + num_members - 1) / num_members // the maximum # of partitions each member should receive
	dbgf("level = %v", level)

	unassigned := make(map[int32]struct{}) // set of unassigned partition ids

	// initially all partitions are unassigned
	for _, p := range partitions {
		unassigned[p] = struct{}{}
	}
	dbgf("unassigned = %v", unassigned)

	// let each member keep up to 'level' of its current assignment
	for m, a := range assignment {
		if len(a) > level {
			a = a[:level]
			assignment[m] = a
		}
		for _, p := range a {
			delete(unassigned, p)
		}
	}
	dbgf("assignment = %v", assignment)
	dbgf("unassigned = %v", unassigned)

	// assign the unassigned partitions to any member with < level partitions
	for p := range unassigned {
	assignment_loop:
		for m, a := range assignment {
			if len(a) < level {
				a = append(a, p)
				assignment[m] = a
				break assignment_loop
			}
		}
	}

	// and we're done
	dbgf("assignment = %v", assignment)
}

// ----------------------------------

// a sortable, comparible list of members
type memberslist []string // a list of the member ids

// implement sort.Interface
func (ml memberslist) Len() int           { return len(ml) }
func (ml memberslist) Swap(i, j int)      { ml[i], ml[j] = ml[j], ml[i] }
func (ml memberslist) Less(i, j int) bool { return ml[i] < ml[j] }

// and compare two sorted memberslist for equality
func (ml memberslist) Equal(ml2 memberslist) bool {
	if len(ml) != len(ml2) {
		return false
	}
	for i, m := range ml {
		if m != ml2[i] {
			return false
		}
	}
	return true
}

// ----------------------------------

// a sortable, comparible list of partition ids
type partitionslist []int32 // a list of the partition ids

// implement sort.Interface
func (pl partitionslist) Len() int           { return len(pl) }
func (pl partitionslist) Swap(i, j int)      { pl[i], pl[j] = pl[j], pl[i] }
func (pl partitionslist) Less(i, j int) bool { return pl[i] < pl[j] }

// and compare two sorted partitionslist for equality
func (pl partitionslist) Equal(pl2 partitionslist) bool {
	if len(pl) != len(pl2) {
		return false
	}
	for i, m := range pl {
		if m != pl2[i] {
			return false
		}
	}
	return true
}
<commit_msg>fix computation of matched_topics so it is right<commit_after>/*
  A partitioner assigns partitions within a consumer group like a
  consistent hash would. When group membership changes the least
  number of partitions are reassigned.

  I don't actually use a consistent hash because such algos take
  either a lot of CPU (the max(N-hash-functions) algo) or a lot
  of ram (the points-on-the-circle algo), or aren't applicable
  when members come and go (jump-hash).

  Instead I take advantage of the fact that I can know the current
  assignments of the remaining group members when constructing
  the new assignments. That makes keep the partitioning stable
  relatively easy to do:
    1) calculate an ideal load for each consumer (I assume equal
	   weight since I don't need anything more complicated...yet)
	2) remove excess partitions from overloaded consumers
	3) add the abandonned and excess partitions to the least loaded
	   consumers.

  Note that the state of the partition at time T depends on
  the history of the consumer group memberships from the time when
  there was 1 consumer until time T.

  Copyright 2016 MistSys
*/

package stable

import (
	"fmt"
	"log"
	"sort"

	"github.com/Shopify/sarama"
)

// a partitioner that assigns partitions to consumers such that as consumers come and go the least number of partitions are reassigned
type stablePartitioner string

// global instance of the stable partitioner
const Stable stablePartitioner = "stable"

// print a debug message
func dbgf(format string, args ...interface{}) {
	log.Printf(format, args...)
}

func (sp stablePartitioner) PrepareJoin(jreq *sarama.JoinGroupRequest, topics []string, current_assignments map[string][]int32) {
	// encode the current assignments in a manner proprietary to this partitioner
	var data = data{
		version:     1,
		assignments: current_assignments,
	}

	jreq.AddGroupProtocolMetadata(string(sp),
		&sarama.ConsumerGroupMemberMetadata{
			Version:  1,
			Topics:   topics,
			UserData: data.marshal(),
		})
}

// for each topic in jresp, assign the topic's partitions to the members requesting the topic
func (sp stablePartitioner) Partition(sreq *sarama.SyncGroupRequest, jresp *sarama.JoinGroupResponse, client sarama.Client) error {
	if jresp.GroupProtocol != string(sp) {
		return fmt.Errorf("sarama.JoinGroupResponse.GroupProtocol %q unexpected; expected %q", jresp.GroupProtocol != string(sp))
	}
	by_member, err := jresp.GetMembers() //  map of member to ConsumerGroupMemberMetadata
	dbgf("by_member = %v", by_member)
	if err != nil {
		return err
	}

	// invert the data, so we have the requests grouped by topic (they arrived grouped by member, since the kafka broker treats the data from each consumer as an opaque blob, so it couldn't do this step for us)
	by_topic := make(map[string]map[string][]int32) // map of topic to members and members to current partition assignment
	for member, request := range by_member {
		if request.Version != 1 {
			// skip unsupported versions. we'll only assign to clients we can understand. Since we are such a client
			// we won't block all consumers (at least for those topics we consume). If this ends up a bad idea, we
			// can always change this code to return an error.
			continue
		}
		var data data
		err := data.unmarshal(request.UserData)
		if err != nil {
			return err
		}

		for _, topic := range request.Topics {
			members, ok := by_topic[topic]
			if !ok {
				members = make(map[string][]int32)
				by_topic[topic] = members
			}
			members[member] = data.assignments[topic] // NOTE: might be nil, which is OK. It just means the member wants to consume the partition but isn't doing so currently
		}
	}
	dbgf("by_topic = %v", by_topic)

	// lookup the partitions in each topic. since we are asking for all partitions, not just the online ones, the numbering
	// appears to always be 0...N-1. But in case I don't understand kafka and there is some corner case where the numbering
	// of partitions is different I keep careful track of the exact numbers I've received.
	var partitions_by_topic = make(map[string]partitionslist)
	for topic := range by_topic {
		// note: calls to client.Partitions() hit the metadata cache in the sarama client, so we don't gain much by asking concurrently
		partitions, err := client.Partitions(topic)
		if err != nil {
			// what to do? we could maybe skip the topic, assigning it to no-one. But I/O errors are likely to happen again.
			// so let's stop partitioning and return the error.
			return err
		}
		pl := make(partitionslist, len(partitions)) // we must make a copy before we sort the list
		for i, p := range partitions {
			pl[i] = p
		}
		sort.Sort(pl)
		partitions_by_topic[topic] = pl
	}
	dbgf("partitions_by_topic = %v", partitions_by_topic)

	// and compute the sorted set of members of each topic
	var members_by_topic = make(map[string]memberslist)
	for topic, members := range by_topic {
		ml := make(memberslist, 0, len(members))
		for m := range members {
			ml = append(ml, m)
		}
		sort.Sort(ml)
		members_by_topic[topic] = ml
	}
	dbgf("members_by_topic = %v", members_by_topic)

	// I want topics with the same # of partitions and the same consumer group membership to result in the same partition assignments.
	// That way messages published under identical partition keys in those topics will all end up consumed by the same member.
	// So organize topics into groups which will be partitioned identically
	var matched_topics = make(map[string]string) // map from each topic to the 'master' topic with the same # of partitions and group membership. Topics which are unique are their own master
topic_match_loop:
	for topic, members := range members_by_topic {
		// see if a match exists
		num_partitions := len(partitions_by_topic[topic])
		for t := range matched_topics { // TODO if # of topics gets large enough this shows up in the profiler, change this to some sort of a map lookup, rather than this O(N^2) search
			if num_partitions == len(partitions_by_topic[t]) && members.Equal(members_by_topic[t]) {
				// match; have topic 'topic' be partitioned the same way as topic 't'
				matched_topics[topic] = matched_topics[t]
				continue topic_match_loop
			}
		}
		// no existing topic matches this one, so it is its own match
		matched_topics[topic] = topic
	}
	dbgf("matched_topics = %v", matched_topics)

	// adjust the partitioning of each master topic in by_topic
	for topic, match := range matched_topics {
		if topic == match {
			adjust_partitioning(by_topic[topic], partitions_by_topic[topic])
		} // else it is not a master topic. once the master has been partitioned we'll simply copy the result
	}

	// set matched topics to the same assignment as their master
	for topic, match := range matched_topics {
		if topic != match {
			by_topic[topic] = by_topic[match]
		}
	}

	// invert by_topic into the equivalent organized by member, and then by topic
	assignments := make(map[string]map[string][]int32) // map of member to topics, and topic to partitions
	for topic, members := range by_topic {
		for member, partitions := range members {
			topics, ok := assignments[member]
			if !ok {
				topics = make(map[string][]int32)
				assignments[member] = topics
			}
			topics[topic] = partitions
		}
	}
	dbgf("assignments = %v", assignments)

	// and encode the assignments in the sync request
	for member, topics := range assignments {
		sreq.AddGroupAssignmentMember(member,
			&sarama.ConsumerGroupMemberAssignment{
				Version: 1,
				Topics:  topics,
			})
	}

	return nil
}

func (stablePartitioner) ParseSync(sresp *sarama.SyncGroupResponse) (map[string][]int32, error) {
	if len(sresp.MemberAssignment) == 0 {
		// in the corner case that we ask for no topics, we get nothing back. However sarama fd498173ae2bf (head of master branch Nov 6th 2016) will return a useless error if we call sresp.GetMemberAssignment() in this case
		return nil, nil
	}
	ma, err := sresp.GetMemberAssignment()
	dbgf("MemberAssignment %v", ma)
	if err != nil {
		return nil, err
	}
	if ma.Version != 1 {
		return nil, fmt.Errorf("unsupported MemberAssignment version %d", ma.Version)
	}
	return ma.Topics, nil
}

// ----------------------------------

// adjust_partitioning does the main work. it adjusts the partition assignment map it is passed in-place
func adjust_partitioning(assignment map[string][]int32, partitions partitionslist) {
	dbgf("adjust_partitioning(assignment = %v, partitions = %v)", assignment, partitions)
	num_members := len(assignment)
	dbgf("num_members = %v", num_members)
	if num_members == 0 {
		// no one wants this topic; stop now before we /0
		return
	}
	num_partitions := len(partitions)
	dbgf("num_partitions = %v", num_partitions)
	level := (num_partitions + num_members - 1) / num_members // the maximum # of partitions each member should receive
	dbgf("level = %v", level)

	unassigned := make(map[int32]struct{}) // set of unassigned partition ids

	// initially all partitions are unassigned
	for _, p := range partitions {
		unassigned[p] = struct{}{}
	}
	dbgf("unassigned = %v", unassigned)

	// let each member keep up to 'level' of its current assignment
	for m, a := range assignment {
		if len(a) > level {
			a = a[:level]
			assignment[m] = a
		}
		for _, p := range a {
			delete(unassigned, p)
		}
	}
	dbgf("assignment = %v", assignment)
	dbgf("unassigned = %v", unassigned)

	// assign the unassigned partitions to any member with < level partitions
	for p := range unassigned {
	assignment_loop:
		for m, a := range assignment {
			if len(a) < level {
				a = append(a, p)
				assignment[m] = a
				break assignment_loop
			}
		}
	}

	// and we're done
	dbgf("assignment = %v", assignment)
}

// ----------------------------------

// a sortable, comparible list of members
type memberslist []string // a list of the member ids

// implement sort.Interface
func (ml memberslist) Len() int           { return len(ml) }
func (ml memberslist) Swap(i, j int)      { ml[i], ml[j] = ml[j], ml[i] }
func (ml memberslist) Less(i, j int) bool { return ml[i] < ml[j] }

// and compare two sorted memberslist for equality
func (ml memberslist) Equal(ml2 memberslist) bool {
	if len(ml) != len(ml2) {
		return false
	}
	for i, m := range ml {
		if m != ml2[i] {
			return false
		}
	}
	return true
}

// ----------------------------------

// a sortable, comparible list of partition ids
type partitionslist []int32 // a list of the partition ids

// implement sort.Interface
func (pl partitionslist) Len() int           { return len(pl) }
func (pl partitionslist) Swap(i, j int)      { pl[i], pl[j] = pl[j], pl[i] }
func (pl partitionslist) Less(i, j int) bool { return pl[i] < pl[j] }

// and compare two sorted partitionslist for equality
func (pl partitionslist) Equal(pl2 partitionslist) bool {
	if len(pl) != len(pl2) {
		return false
	}
	for i, m := range pl {
		if m != pl2[i] {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"encoding/json"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils"
	"gopkg.in/juju/charm.v4"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"
)

var metricsLogger = loggo.GetLogger("juju.state.metrics")

const (
	CleanupAge = time.Hour * 24
)

// MetricBatch represents a batch of metrics reported from a unit.
// These will be received from the unit in batches.
// The main contents of the metric (key, value) is defined
// by the charm author and sent from the unit via a call to
// add-metric
type MetricBatch struct {
	st  *State
	doc metricBatchDoc
}

type metricBatchDoc struct {
	UUID     string    `bson:"_id"`
	EnvUUID  string    `bson:"env-uuid"`
	Unit     string    `bson:"unit"`
	CharmUrl string    `bson:"charmurl"`
	Sent     bool      `bson:"sent"`
	Created  time.Time `bson:"created"`
	Metrics  []Metric  `bson:"metrics"`
}

// Metric represents a single Metric.
type Metric struct {
	Key         string    `bson:"key"`
	Value       string    `bson:"value"`
	Time        time.Time `bson:"time"`
	Credentials []byte    `bson:"credentials"`
}

// validate checks that the MetricBatch contains valid metrics.
func (m *MetricBatch) validate() error {
	charmUrl, err := charm.ParseURL(m.doc.CharmUrl)
	if err != nil {
		return err
	}
	chrm, err := m.st.Charm(charmUrl)
	if err != nil {
		return err
	}
	chrmMetrics := chrm.Metrics()
	if chrmMetrics == nil {
		return errors.Errorf("charm doesn't implement metrics")
	}
	for _, m := range m.doc.Metrics {
		if err := chrmMetrics.ValidateMetric(m.Key, m.Value); err != nil {
			return err
		}
	}
	return nil
}

// AddMetric adds a new batch of metrics to the database.
// A UUID for the metric will be generated and the new MetricBatch will be returned
func (st *State) addMetrics(unitTag names.UnitTag, charmUrl *charm.URL, created time.Time, metrics []Metric) (*MetricBatch, error) {
	if len(metrics) == 0 {
		return nil, errors.New("cannot add a batch of 0 metrics")
	}
	uuid, err := utils.NewUUID()
	if err != nil {
		return nil, err
	}

	metric := &MetricBatch{
		st: st,
		doc: metricBatchDoc{
			UUID:     uuid.String(),
			EnvUUID:  st.EnvironUUID(),
			Unit:     unitTag.Id(),
			CharmUrl: charmUrl.String(),
			Sent:     false,
			Created:  created,
			Metrics:  metrics,
		}}
	if err := metric.validate(); err != nil {
		return nil, err
	}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			notDead, err := isNotDead(st.db, unitsC, st.docID(unitTag.Id()))
			if err != nil || !notDead {
				return nil, errors.NotFoundf(unitTag.Id())
			}
		}
		ops := []txn.Op{{
			C:      unitsC,
			Id:     st.docID(unitTag.Id()),
			Assert: notDeadDoc,
		}, {
			C:      metricsC,
			Id:     metric.UUID(),
			Assert: txn.DocMissing,
			Insert: &metric.doc,
		}}
		return ops, nil
	}
	err = st.run(buildTxn)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return metric, nil
}

// MetricBatches returns all metric batches currently stored in state.
// TODO (tasdomas): this method is currently only used in the uniter worker test -
//                  it needs to be modified to restrict the scope of the values it
//                  returns if it is to be used outside of tests.
func (st *State) MetricBatches() ([]MetricBatch, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	docs := []metricBatchDoc{}
	err := c.Find(nil).All(&docs)
	if err != nil {
		return nil, errors.Trace(err)
	}
	results := make([]MetricBatch, len(docs))
	for i, doc := range docs {
		results[i] = MetricBatch{st: st, doc: doc}
	}
	return results, nil
}

// MetricBatch returns the metric batch with the given id.
func (st *State) MetricBatch(id string) (*MetricBatch, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	doc := metricBatchDoc{}
	err := c.Find(bson.M{"_id": id}).One(&doc)
	if err == mgo.ErrNotFound {
		return nil, errors.NotFoundf("metric %v", id)
	}
	if err != nil {
		return nil, err
	}
	return &MetricBatch{st: st, doc: doc}, nil
}

// CleanupOldMetrics looks for metrics that are 24 hours old (or older)
// and have been sent. Any metrics it finds are deleted.
func (st *State) CleanupOldMetrics() error {
	age := time.Now().Add(-(CleanupAge))
	c, closer := st.getCollection(metricsC)
	defer closer()
	// Nothing else in the system will interact with sent metrics, and nothing needs
	// to watch them either; so in this instance it's safe to do an end run around the
	// mgo/txn package. See State.cleanupRelationSettings for a similar situation.
	err := c.Remove(bson.M{
		"sent":    true,
		"created": bson.M{"$lte": age},
	})
	if err == mgo.ErrNotFound {
		metricsLogger.Infof("no metrics found to cleanup")
		return nil
	}
	return err
}

// MetricsToSend returns batchSize metrics that need to be sent
// to the collector
func (st *State) MetricsToSend(batchSize int) ([]*MetricBatch, error) {
	var docs []metricBatchDoc
	c, closer := st.getCollection(metricsC)
	defer closer()
	err := c.Find(bson.M{
		"sent": false,
	}).Limit(batchSize).All(&docs)
	if err != nil {
		return nil, errors.Trace(err)
	}

	batch := make([]*MetricBatch, len(docs))
	for i, doc := range docs {
		batch[i] = &MetricBatch{st: st, doc: doc}

	}

	return batch, nil
}

// CountofUnsentMetrics returns the number of metrics that
// haven't been sent to the collection service.
func (st *State) CountofUnsentMetrics() (int, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	return c.Find(bson.M{
		"sent": false,
	}).Count()
}

// CountofSentMetrics returns the number of metrics that
// have been sent to the collection service and have not
// been removed by the cleanup worker.
func (st *State) CountofSentMetrics() (int, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	return c.Find(bson.M{
		"sent": true,
	}).Count()
}

// MarshalJSON defines how the MetricBatch type should be
// converted to json.
func (m *MetricBatch) MarshalJSON() ([]byte, error) {
	return json.Marshal(m.doc)
}

// UUID returns to uuid of the metric.
func (m *MetricBatch) UUID() string {
	return m.doc.UUID
}

// EnvUUID returns the environment UUID this metric applies to.
func (m *MetricBatch) EnvUUID() string {
	return m.doc.EnvUUID
}

// Unit returns the name of the unit this metric was generated in.
func (m *MetricBatch) Unit() string {
	return m.doc.Unit
}

// CharmURL returns the charm url for the charm this metric was generated in.
func (m *MetricBatch) CharmURL() string {
	return m.doc.CharmUrl
}

// Created returns the time this metric batch was created.
func (m *MetricBatch) Created() time.Time {
	return m.doc.Created
}

// Sent returns a flag to tell us if this metric has been sent to the metric
// collection service
func (m *MetricBatch) Sent() bool {
	return m.doc.Sent
}

// Metrics returns the metrics in this batch.
func (m *MetricBatch) Metrics() []Metric {
	result := make([]Metric, len(m.doc.Metrics))
	copy(result, m.doc.Metrics)
	return result
}

// SetSent sets the sent flag to true
func (m *MetricBatch) SetSent() error {
	ops := setSentOps([]string{m.UUID()})
	if err := m.st.runTransaction(ops); err != nil {
		return errors.Annotatef(err, "cannot set metric sent for metric %q", m.UUID())
	}

	m.doc.Sent = true
	return nil
}

func setSentOps(batchUUIDs []string) []txn.Op {
	ops := make([]txn.Op, len(batchUUIDs))
	for i, u := range batchUUIDs {
		ops[i] = txn.Op{
			C:      metricsC,
			Id:     u,
			Assert: txn.DocExists,
			Update: bson.M{"$set": bson.M{"sent": true}},
		}
	}
	return ops
}

// SetMetricBatchesSent sets sent on each MetricBatch corresponding to the uuids provided.
func (st *State) SetMetricBatchesSent(batchUUIDs []string) error {
	ops := setSentOps(batchUUIDs)
	if err := st.runTransaction(ops); err != nil {
		return errors.Annotatef(err, "cannot set metric sent in bulk call")
	}
	return nil
}
<commit_msg>return errors.trace instead of bare err<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"encoding/json"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils"
	"gopkg.in/juju/charm.v4"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"
)

var metricsLogger = loggo.GetLogger("juju.state.metrics")

const (
	CleanupAge = time.Hour * 24
)

// MetricBatch represents a batch of metrics reported from a unit.
// These will be received from the unit in batches.
// The main contents of the metric (key, value) is defined
// by the charm author and sent from the unit via a call to
// add-metric
type MetricBatch struct {
	st  *State
	doc metricBatchDoc
}

type metricBatchDoc struct {
	UUID     string    `bson:"_id"`
	EnvUUID  string    `bson:"env-uuid"`
	Unit     string    `bson:"unit"`
	CharmUrl string    `bson:"charmurl"`
	Sent     bool      `bson:"sent"`
	Created  time.Time `bson:"created"`
	Metrics  []Metric  `bson:"metrics"`
}

// Metric represents a single Metric.
type Metric struct {
	Key         string    `bson:"key"`
	Value       string    `bson:"value"`
	Time        time.Time `bson:"time"`
	Credentials []byte    `bson:"credentials"`
}

// validate checks that the MetricBatch contains valid metrics.
func (m *MetricBatch) validate() error {
	charmUrl, err := charm.ParseURL(m.doc.CharmUrl)
	if err != nil {
		return errors.Trace(err)
	}
	chrm, err := m.st.Charm(charmUrl)
	if err != nil {
		return errors.Trace(err)
	}
	chrmMetrics := chrm.Metrics()
	if chrmMetrics == nil {
		return errors.Errorf("charm doesn't implement metrics")
	}
	for _, m := range m.doc.Metrics {
		if err := chrmMetrics.ValidateMetric(m.Key, m.Value); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

// AddMetric adds a new batch of metrics to the database.
// A UUID for the metric will be generated and the new MetricBatch will be returned
func (st *State) addMetrics(unitTag names.UnitTag, charmUrl *charm.URL, created time.Time, metrics []Metric) (*MetricBatch, error) {
	if len(metrics) == 0 {
		return nil, errors.New("cannot add a batch of 0 metrics")
	}
	uuid, err := utils.NewUUID()
	if err != nil {
		return nil, err
	}

	metric := &MetricBatch{
		st: st,
		doc: metricBatchDoc{
			UUID:     uuid.String(),
			EnvUUID:  st.EnvironUUID(),
			Unit:     unitTag.Id(),
			CharmUrl: charmUrl.String(),
			Sent:     false,
			Created:  created,
			Metrics:  metrics,
		}}
	if err := metric.validate(); err != nil {
		return nil, err
	}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			notDead, err := isNotDead(st.db, unitsC, st.docID(unitTag.Id()))
			if err != nil || !notDead {
				return nil, errors.NotFoundf(unitTag.Id())
			}
		}
		ops := []txn.Op{{
			C:      unitsC,
			Id:     st.docID(unitTag.Id()),
			Assert: notDeadDoc,
		}, {
			C:      metricsC,
			Id:     metric.UUID(),
			Assert: txn.DocMissing,
			Insert: &metric.doc,
		}}
		return ops, nil
	}
	err = st.run(buildTxn)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return metric, nil
}

// MetricBatches returns all metric batches currently stored in state.
// TODO (tasdomas): this method is currently only used in the uniter worker test -
//                  it needs to be modified to restrict the scope of the values it
//                  returns if it is to be used outside of tests.
func (st *State) MetricBatches() ([]MetricBatch, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	docs := []metricBatchDoc{}
	err := c.Find(nil).All(&docs)
	if err != nil {
		return nil, errors.Trace(err)
	}
	results := make([]MetricBatch, len(docs))
	for i, doc := range docs {
		results[i] = MetricBatch{st: st, doc: doc}
	}
	return results, nil
}

// MetricBatch returns the metric batch with the given id.
func (st *State) MetricBatch(id string) (*MetricBatch, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	doc := metricBatchDoc{}
	err := c.Find(bson.M{"_id": id}).One(&doc)
	if err == mgo.ErrNotFound {
		return nil, errors.NotFoundf("metric %v", id)
	}
	if err != nil {
		return nil, err
	}
	return &MetricBatch{st: st, doc: doc}, nil
}

// CleanupOldMetrics looks for metrics that are 24 hours old (or older)
// and have been sent. Any metrics it finds are deleted.
func (st *State) CleanupOldMetrics() error {
	age := time.Now().Add(-(CleanupAge))
	c, closer := st.getCollection(metricsC)
	defer closer()
	// Nothing else in the system will interact with sent metrics, and nothing needs
	// to watch them either; so in this instance it's safe to do an end run around the
	// mgo/txn package. See State.cleanupRelationSettings for a similar situation.
	err := c.Remove(bson.M{
		"sent":    true,
		"created": bson.M{"$lte": age},
	})
	if err == mgo.ErrNotFound {
		metricsLogger.Infof("no metrics found to cleanup")
		return nil
	}
	return errors.Trace(err)
}

// MetricsToSend returns batchSize metrics that need to be sent
// to the collector
func (st *State) MetricsToSend(batchSize int) ([]*MetricBatch, error) {
	var docs []metricBatchDoc
	c, closer := st.getCollection(metricsC)
	defer closer()
	err := c.Find(bson.M{
		"sent": false,
	}).Limit(batchSize).All(&docs)
	if err != nil {
		return nil, errors.Trace(err)
	}

	batch := make([]*MetricBatch, len(docs))
	for i, doc := range docs {
		batch[i] = &MetricBatch{st: st, doc: doc}

	}

	return batch, nil
}

// CountofUnsentMetrics returns the number of metrics that
// haven't been sent to the collection service.
func (st *State) CountofUnsentMetrics() (int, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	return c.Find(bson.M{
		"sent": false,
	}).Count()
}

// CountofSentMetrics returns the number of metrics that
// have been sent to the collection service and have not
// been removed by the cleanup worker.
func (st *State) CountofSentMetrics() (int, error) {
	c, closer := st.getCollection(metricsC)
	defer closer()
	return c.Find(bson.M{
		"sent": true,
	}).Count()
}

// MarshalJSON defines how the MetricBatch type should be
// converted to json.
func (m *MetricBatch) MarshalJSON() ([]byte, error) {
	return json.Marshal(m.doc)
}

// UUID returns to uuid of the metric.
func (m *MetricBatch) UUID() string {
	return m.doc.UUID
}

// EnvUUID returns the environment UUID this metric applies to.
func (m *MetricBatch) EnvUUID() string {
	return m.doc.EnvUUID
}

// Unit returns the name of the unit this metric was generated in.
func (m *MetricBatch) Unit() string {
	return m.doc.Unit
}

// CharmURL returns the charm url for the charm this metric was generated in.
func (m *MetricBatch) CharmURL() string {
	return m.doc.CharmUrl
}

// Created returns the time this metric batch was created.
func (m *MetricBatch) Created() time.Time {
	return m.doc.Created
}

// Sent returns a flag to tell us if this metric has been sent to the metric
// collection service
func (m *MetricBatch) Sent() bool {
	return m.doc.Sent
}

// Metrics returns the metrics in this batch.
func (m *MetricBatch) Metrics() []Metric {
	result := make([]Metric, len(m.doc.Metrics))
	copy(result, m.doc.Metrics)
	return result
}

// SetSent sets the sent flag to true
func (m *MetricBatch) SetSent() error {
	ops := setSentOps([]string{m.UUID()})
	if err := m.st.runTransaction(ops); err != nil {
		return errors.Annotatef(err, "cannot set metric sent for metric %q", m.UUID())
	}

	m.doc.Sent = true
	return nil
}

func setSentOps(batchUUIDs []string) []txn.Op {
	ops := make([]txn.Op, len(batchUUIDs))
	for i, u := range batchUUIDs {
		ops[i] = txn.Op{
			C:      metricsC,
			Id:     u,
			Assert: txn.DocExists,
			Update: bson.M{"$set": bson.M{"sent": true}},
		}
	}
	return ops
}

// SetMetricBatchesSent sets sent on each MetricBatch corresponding to the uuids provided.
func (st *State) SetMetricBatchesSent(batchUUIDs []string) error {
	ops := setSentOps(batchUUIDs)
	if err := st.runTransaction(ops); err != nil {
		return errors.Annotatef(err, "cannot set metric sent in bulk call")
	}
	return nil
}
<|endoftext|>


<commit_before>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"strconv"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]pstore.PeerInfo
	addrs  []ma.Multiaddr
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]pstore.PeerInfo),
		disconnect: make(chan struct{}, 1),
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	if ar.addrs != nil && ar.autonat.Status() == autonat.NATStatusPrivate {
		return ar.addrs
	} else {
		return ar.addrsF(addrs)
	}
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			// invalidate addrs
			ar.mx.Lock()
			if ar.addrs != nil {
				ar.addrs = nil
				push = true
			}
			ar.mx.Unlock()

			// if we had previously announced relay addrs, push our public addrs
			if push {
				push = false
				ar.host.PushIdentify()
			}

		case autonat.NATStatusPrivate:
			push = false // clear, findRelays pushes as needed
			ar.findRelays(ctx)
		}

		select {
		case <-ar.disconnect:
			// invalidate addrs
			ar.mx.Lock()
			if ar.addrs != nil {
				ar.addrs = nil
				push = true
			}
			ar.mx.Unlock()
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) {
again:
	ar.mx.Lock()
	haveRelays := len(ar.relays)
	if haveRelays >= DesiredRelays {
		ar.mx.Unlock()
		// this dance is necessary to cover the Private->Public->Private transition
		// where we were already connected to enough relays while private and dropped
		// the addrs while public
		if ar.addrs == nil {
			ar.updateAddrs()
		}
		return
	}
	need := DesiredRelays - len(ar.relays)
	ar.mx.Unlock()

	limit := 1000

	dctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	pis, err := discovery.FindPeers(dctx, ar.discover, RelayRendezvous, limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())
		return
	}

	pis = ar.selectRelays(ctx, pis, 20, 50)
	update := 0

	for _, pi := range pis {
		ar.mx.Lock()
		if _, ok := ar.relays[pi.ID]; ok {
			ar.mx.Unlock()
			continue
		}
		ar.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 15*time.Second)
		err = ar.host.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		ar.mx.Lock()
		ar.relays[pi.ID] = pi
		haveRelays++
		ar.mx.Unlock()

		// tag the connection as very important
		ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)

		update++
		need--
		if need == 0 {
			break
		}
	}

	if haveRelays == 0 {
		// we failed to find any relays and we are not connected to any!
		// wait a little and try again, the discovery query might have returned only dead peers
		select {
		case <-time.After(30 * time.Second):
			goto again
		case <-ctx.Done():
			return
		}
	}

	if update > 0 || ar.addrs == nil {
		ar.updateAddrs()
	}
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo, count, maxq int) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	if len(pis) == 0 {
		log.Debugf("no relays discovered")
		return pis
	}

	if len(pis) < count {
		count = len(pis)
	}

	if len(pis) < maxq {
		maxq = len(pis)
	}

	// only select relays that can be found by routing
	type queryResult struct {
		pi  pstore.PeerInfo
		err error
	}
	result := make([]pstore.PeerInfo, 0, count)
	resultCh := make(chan queryResult, maxq)

	qctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	// shuffle to randomize the order of queries
	shuffleRelays(pis)
	for _, pi := range pis[:maxq] {
		go func(p peer.ID) {
			pi, err := ar.router.FindPeer(qctx, p)
			if err != nil {
				log.Debugf("error finding relay peer %s: %s", p, err.Error())
			}
			resultCh <- queryResult{pi: pi, err: err}
		}(pi.ID)
	}

	rcount := 0
	for len(result) < count && rcount < maxq {
		select {
		case qr := <-resultCh:
			rcount++
			if qr.err == nil {
				pi := cleanupAddressSet(qr.pi)
				if len(pi.Addrs) > 0 {
					result = append(result, pi)
				} else {
					log.Debugf("ignoring relay peer %s: cleaned up address set is empty", pi.ID)
				}
			}

		case <-qctx.Done():
			break
		}
	}

	shuffleRelays(result)
	return result
}

func (ar *AutoRelay) updateAddrs() {
	ar.doUpdateAddrs()
	ar.host.PushIdentify()
}

// This function updates our NATed advertised addrs (ar.addrs)
// The public addrs are rewritten so that they only retain the public IP part; they
// become undialable but are useful as a hint to the dialer to determine whether or not
// to dial private addrs.
// The non-public addrs are included verbatim so that peers behind the same NAT/firewall
// can still dial us directly.
// On top of those, we add the relay-specific addrs for the relays to which we are
// connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
// through which we can be dialed.
func (ar *AutoRelay) doUpdateAddrs() {
	ar.mx.Lock()
	defer ar.mx.Unlock()

	addrs := ar.baseAddrs()
	raddrs := make([]ma.Multiaddr, 0, len(addrs)+len(ar.relays))

	// remove our public addresses from the list
	for _, addr := range addrs {
		if manet.IsPublicAddr(addr) {
			continue
		}
		raddrs = append(raddrs, addr)
	}

	// add relay specific addrs to the list
	for _, pi := range ar.relays {
		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", pi.ID.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range pi.Addrs {
			if !manet.IsPrivateAddr(addr) {
				pub := addr.Encapsulate(circuit)
				raddrs = append(raddrs, pub)
			}
		}
	}

	ar.addrs = raddrs
}

// This function cleans up a relay's address set to remove private addresses and curtail
// addrsplosion. For the latter, we use the following heuristic:
// - if the address set includes a (tcp) address with the default port 4001,
//   we remove all tcp addrs with a different port
func cleanupAddressSet(pi pstore.PeerInfo) pstore.PeerInfo {
	// pass-1: find default port
	has4001 := false
	for _, addr := range pi.Addrs {
		port, err := tcpPort(addr)
		if err != nil {
			continue
		}
		if port == 4001 {
			has4001 = true
			break
		}
	}

	// pass-2: cleanup
	var newAddrs []ma.Multiaddr

	for _, addr := range pi.Addrs {
		if manet.IsPrivateAddr(addr) {
			continue
		}

		if has4001 {
			port, err := tcpPort(addr)
			if err == nil && port != 4001 {
				continue
			}
		}

		newAddrs = append(newAddrs, addr)
	}

	return pstore.PeerInfo{ID: pi.ID, Addrs: newAddrs}
}

func tcpPort(addr ma.Multiaddr) (int, error) {
	val, err := addr.ValueForProtocol(ma.P_TCP)
	if err != nil {
		// not tcp
		return 0, err
	}
	return strconv.Atoi(val)
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<commit_msg>some better logging<commit_after>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"strconv"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]pstore.PeerInfo
	addrs  []ma.Multiaddr
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]pstore.PeerInfo),
		disconnect: make(chan struct{}, 1),
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	defer ar.mx.Unlock()
	if ar.addrs != nil && ar.autonat.Status() == autonat.NATStatusPrivate {
		return ar.addrs
	} else {
		return ar.addrsF(addrs)
	}
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			// invalidate addrs
			ar.mx.Lock()
			if ar.addrs != nil {
				ar.addrs = nil
				push = true
			}
			ar.mx.Unlock()

			// if we had previously announced relay addrs, push our public addrs
			if push {
				push = false
				ar.host.PushIdentify()
			}

		case autonat.NATStatusPrivate:
			push = false // clear, findRelays pushes as needed
			ar.findRelays(ctx)
		}

		select {
		case <-ar.disconnect:
			// invalidate addrs
			ar.mx.Lock()
			if ar.addrs != nil {
				ar.addrs = nil
				push = true
			}
			ar.mx.Unlock()
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) {
again:
	ar.mx.Lock()
	haveRelays := len(ar.relays)
	if haveRelays >= DesiredRelays {
		ar.mx.Unlock()
		// this dance is necessary to cover the Private->Public->Private transition
		// where we were already connected to enough relays while private and dropped
		// the addrs while public
		if ar.addrs == nil {
			ar.updateAddrs()
		}
		return
	}
	need := DesiredRelays - len(ar.relays)
	ar.mx.Unlock()

	limit := 1000

	dctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	pis, err := discovery.FindPeers(dctx, ar.discover, RelayRendezvous, limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())
		return
	}

	log.Debugf("discovered %d relays", len(pis))

	pis = ar.selectRelays(ctx, pis, 20, 50)
	update := 0

	for _, pi := range pis {
		ar.mx.Lock()
		if _, ok := ar.relays[pi.ID]; ok {
			ar.mx.Unlock()
			continue
		}
		ar.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 15*time.Second)
		err = ar.host.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		ar.mx.Lock()
		ar.relays[pi.ID] = pi
		haveRelays++
		ar.mx.Unlock()

		// tag the connection as very important
		ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)

		update++
		need--
		if need == 0 {
			break
		}
	}

	if haveRelays == 0 {
		// we failed to find any relays and we are not connected to any!
		// wait a little and try again, the discovery query might have returned only dead peers
		log.Debug("no relays connected; retrying in 30s")
		select {
		case <-time.After(30 * time.Second):
			goto again
		case <-ctx.Done():
			return
		}
	}

	if update > 0 || ar.addrs == nil {
		ar.updateAddrs()
	}
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo, count, maxq int) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	if len(pis) == 0 {
		return pis
	}

	if len(pis) < count {
		count = len(pis)
	}

	if len(pis) < maxq {
		maxq = len(pis)
	}

	// only select relays that can be found by routing
	type queryResult struct {
		pi  pstore.PeerInfo
		err error
	}
	result := make([]pstore.PeerInfo, 0, count)
	resultCh := make(chan queryResult, maxq)

	qctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	// shuffle to randomize the order of queries
	shuffleRelays(pis)
	for _, pi := range pis[:maxq] {
		go func(p peer.ID) {
			pi, err := ar.router.FindPeer(qctx, p)
			if err != nil {
				log.Debugf("error finding relay peer %s: %s", p, err.Error())
			}
			resultCh <- queryResult{pi: pi, err: err}
		}(pi.ID)
	}

	rcount := 0
	for len(result) < count && rcount < maxq {
		select {
		case qr := <-resultCh:
			rcount++
			if qr.err == nil {
				pi := cleanupAddressSet(qr.pi)
				if len(pi.Addrs) > 0 {
					result = append(result, pi)
				} else {
					log.Debugf("ignoring relay peer %s: cleaned up address set is empty", pi.ID)
				}
			}

		case <-qctx.Done():
			break
		}
	}

	shuffleRelays(result)
	return result
}

func (ar *AutoRelay) updateAddrs() {
	ar.doUpdateAddrs()
	ar.host.PushIdentify()
}

// This function updates our NATed advertised addrs (ar.addrs)
// The public addrs are rewritten so that they only retain the public IP part; they
// become undialable but are useful as a hint to the dialer to determine whether or not
// to dial private addrs.
// The non-public addrs are included verbatim so that peers behind the same NAT/firewall
// can still dial us directly.
// On top of those, we add the relay-specific addrs for the relays to which we are
// connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
// through which we can be dialed.
func (ar *AutoRelay) doUpdateAddrs() {
	ar.mx.Lock()
	defer ar.mx.Unlock()

	addrs := ar.baseAddrs()
	raddrs := make([]ma.Multiaddr, 0, len(addrs)+len(ar.relays))

	// remove our public addresses from the list
	for _, addr := range addrs {
		if manet.IsPublicAddr(addr) {
			continue
		}
		raddrs = append(raddrs, addr)
	}

	// add relay specific addrs to the list
	for _, pi := range ar.relays {
		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", pi.ID.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range pi.Addrs {
			if !manet.IsPrivateAddr(addr) {
				pub := addr.Encapsulate(circuit)
				raddrs = append(raddrs, pub)
			}
		}
	}

	ar.addrs = raddrs
}

// This function cleans up a relay's address set to remove private addresses and curtail
// addrsplosion. For the latter, we use the following heuristic:
// - if the address set includes a (tcp) address with the default port 4001,
//   we remove all tcp addrs with a different port
func cleanupAddressSet(pi pstore.PeerInfo) pstore.PeerInfo {
	// pass-1: find default port
	has4001 := false
	for _, addr := range pi.Addrs {
		port, err := tcpPort(addr)
		if err != nil {
			continue
		}
		if port == 4001 {
			has4001 = true
			break
		}
	}

	// pass-2: cleanup
	var newAddrs []ma.Multiaddr

	for _, addr := range pi.Addrs {
		if manet.IsPrivateAddr(addr) {
			continue
		}

		if has4001 {
			port, err := tcpPort(addr)
			if err == nil && port != 4001 {
				continue
			}
		}

		newAddrs = append(newAddrs, addr)
	}

	return pstore.PeerInfo{ID: pi.ID, Addrs: newAddrs}
}

func tcpPort(addr ma.Multiaddr) (int, error) {
	val, err := addr.ValueForProtocol(ma.P_TCP)
	if err != nil {
		// not tcp
		return 0, err
	}
	return strconv.Atoi(val)
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<|endoftext|>


<commit_before>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	discovery "github.com/libp2p/go-libp2p-discovery"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

var (
	DesiredRelays = 3

	BootDelay = 90 * time.Second
)

// AutoRelayHost is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelayHost struct {
	*basic.BasicHost
	discover discovery.Discoverer
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]pstore.PeerInfo
	addrs  []ma.Multiaddr
}

func NewAutoRelayHost(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer) *AutoRelayHost {
	autonat := autonat.NewAutoNAT(ctx, bhost)
	h := &AutoRelayHost{
		BasicHost:  bhost,
		discover:   discover,
		autonat:    autonat,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]pstore.PeerInfo),
		disconnect: make(chan struct{}, 1),
	}
	bhost.AddrsFactory = h.hostAddrs
	bhost.Network().Notify(h)
	go h.background(ctx)
	return h
}

func (h *AutoRelayHost) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	h.mx.Lock()
	defer h.mx.Unlock()
	if h.addrs != nil && h.autonat.Status() == autonat.NATStatusPrivate {
		return h.addrs
	} else {
		return h.addrsF(addrs)
	}
}

func (h *AutoRelayHost) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	for {
		wait := autonat.AutoNATRefreshInterval
		switch h.autonat.Status() {
		case autonat.NATStatusUnknown:
			wait = autonat.AutoNATRetryInterval
		case autonat.NATStatusPublic:
		case autonat.NATStatusPrivate:
			h.findRelays(ctx)
		}

		select {
		case <-h.disconnect:
			// invalidate addrs
			h.mx.Lock()
			h.addrs = nil
			h.mx.Unlock()
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (h *AutoRelayHost) findRelays(ctx context.Context) {
	h.mx.Lock()
	if len(h.relays) >= DesiredRelays {
		h.mx.Unlock()
		return
	}
	need := DesiredRelays - len(h.relays)
	h.mx.Unlock()

	limit := 20
	for ; need > limit; limit *= 2 {
	}

	dctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	pis, err := discovery.FindPeers(dctx, h.discover, "/libp2p/relay", limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())
		return
	}

	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric
	shuffleRelays(pis)

	update := 0

	for _, pi := range pis {
		h.mx.Lock()
		if _, ok := h.relays[pi.ID]; ok {
			h.mx.Unlock()
			continue
		}
		h.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 60*time.Second)
		err = h.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		h.mx.Lock()
		h.relays[pi.ID] = pi
		h.mx.Unlock()

		update++
		need--
		if need == 0 {
			break
		}
	}

	if update > 0 || h.addrs == nil {
		h.updateAddrs()
	}
}

func (h *AutoRelayHost) updateAddrs() {
	h.doUpdateAddrs()
	h.PushIdentify()
}

func (h *AutoRelayHost) doUpdateAddrs() {
	h.mx.Lock()
	defer h.mx.Unlock()

	addrs := h.addrsF(h.AllAddrs())
	raddrs := make([]ma.Multiaddr, 0, len(addrs)+len(h.relays))

	// remove our public addresses from the list and replace them by just the public IP
	for _, addr := range addrs {
		if manet.IsPublicAddr(addr) {
			ip, err := addr.ValueForProtocol(ma.P_IP4)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip4/%s", ip))
				if err != nil {
					panic(err)
				}

				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}

			ip, err = addr.ValueForProtocol(ma.P_IP6)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip6/%s", ip))
				if err != nil {
					panic(err)
				}
				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}
		} else {
			raddrs = append(raddrs, addr)
		}
	}

	circuit, err := ma.NewMultiaddr("/p2p-circuit")
	if err != nil {
		panic(err)
	}

	for _, pi := range h.relays {
		for _, addr := range pi.Addrs {
			if !manet.IsPrivateAddr(addr) {
				pub := addr.Encapsulate(circuit)
				raddrs = append(raddrs, pub)
			}
		}
	}

	h.addrs = raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

func containsAddr(lst []ma.Multiaddr, addr ma.Multiaddr) bool {
	for _, xaddr := range lst {
		if xaddr.Equal(addr) {
			return true
		}
	}
	return false
}

// notify
func (h *AutoRelayHost) Listen(inet.Network, ma.Multiaddr)      {}
func (h *AutoRelayHost) ListenClose(inet.Network, ma.Multiaddr) {}
func (h *AutoRelayHost) Connected(inet.Network, inet.Conn)      {}

func (h *AutoRelayHost) Disconnected(_ inet.Network, c inet.Conn) {
	p := c.RemotePeer()
	h.mx.Lock()
	defer h.mx.Unlock()
	if _, ok := h.relays[p]; ok {
		delete(h.relays, p)
		select {
		case h.disconnect <- struct{}{}:
		default:
		}
	}
}

func (h *AutoRelayHost) OpenedStream(inet.Network, inet.Stream) {}
func (h *AutoRelayHost) ClosedStream(inet.Network, inet.Stream) {}

var _ host.Host = (*AutoRelayHost)(nil)
<commit_msg>use AllAddrs as the address factory in autonat<commit_after>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	discovery "github.com/libp2p/go-libp2p-discovery"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

var (
	DesiredRelays = 3

	BootDelay = 90 * time.Second
)

// AutoRelayHost is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelayHost struct {
	*basic.BasicHost
	discover discovery.Discoverer
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]pstore.PeerInfo
	addrs  []ma.Multiaddr
}

func NewAutoRelayHost(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer) *AutoRelayHost {
	autonat := autonat.NewAutoNAT(ctx, bhost, bhost.AllAddrs)
	h := &AutoRelayHost{
		BasicHost:  bhost,
		discover:   discover,
		autonat:    autonat,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]pstore.PeerInfo),
		disconnect: make(chan struct{}, 1),
	}
	bhost.AddrsFactory = h.hostAddrs
	bhost.Network().Notify(h)
	go h.background(ctx)
	return h
}

func (h *AutoRelayHost) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	h.mx.Lock()
	defer h.mx.Unlock()
	if h.addrs != nil && h.autonat.Status() == autonat.NATStatusPrivate {
		return h.addrs
	} else {
		return h.addrsF(addrs)
	}
}

func (h *AutoRelayHost) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	for {
		wait := autonat.AutoNATRefreshInterval
		switch h.autonat.Status() {
		case autonat.NATStatusUnknown:
			wait = autonat.AutoNATRetryInterval
		case autonat.NATStatusPublic:
		case autonat.NATStatusPrivate:
			h.findRelays(ctx)
		}

		select {
		case <-h.disconnect:
			// invalidate addrs
			h.mx.Lock()
			h.addrs = nil
			h.mx.Unlock()
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (h *AutoRelayHost) findRelays(ctx context.Context) {
	h.mx.Lock()
	if len(h.relays) >= DesiredRelays {
		h.mx.Unlock()
		return
	}
	need := DesiredRelays - len(h.relays)
	h.mx.Unlock()

	limit := 20
	for ; need > limit; limit *= 2 {
	}

	dctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	pis, err := discovery.FindPeers(dctx, h.discover, "/libp2p/relay", limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())
		return
	}

	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric
	shuffleRelays(pis)

	update := 0

	for _, pi := range pis {
		h.mx.Lock()
		if _, ok := h.relays[pi.ID]; ok {
			h.mx.Unlock()
			continue
		}
		h.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 60*time.Second)
		err = h.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		h.mx.Lock()
		h.relays[pi.ID] = pi
		h.mx.Unlock()

		update++
		need--
		if need == 0 {
			break
		}
	}

	if update > 0 || h.addrs == nil {
		h.updateAddrs()
	}
}

func (h *AutoRelayHost) updateAddrs() {
	h.doUpdateAddrs()
	h.PushIdentify()
}

func (h *AutoRelayHost) doUpdateAddrs() {
	h.mx.Lock()
	defer h.mx.Unlock()

	addrs := h.addrsF(h.AllAddrs())
	raddrs := make([]ma.Multiaddr, 0, len(addrs)+len(h.relays))

	// remove our public addresses from the list and replace them by just the public IP
	for _, addr := range addrs {
		if manet.IsPublicAddr(addr) {
			ip, err := addr.ValueForProtocol(ma.P_IP4)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip4/%s", ip))
				if err != nil {
					panic(err)
				}

				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}

			ip, err = addr.ValueForProtocol(ma.P_IP6)
			if err == nil {
				pub, err := ma.NewMultiaddr(fmt.Sprintf("/ip6/%s", ip))
				if err != nil {
					panic(err)
				}
				if !containsAddr(raddrs, pub) {
					raddrs = append(raddrs, pub)
				}
				continue
			}
		} else {
			raddrs = append(raddrs, addr)
		}
	}

	circuit, err := ma.NewMultiaddr("/p2p-circuit")
	if err != nil {
		panic(err)
	}

	for _, pi := range h.relays {
		for _, addr := range pi.Addrs {
			if !manet.IsPrivateAddr(addr) {
				pub := addr.Encapsulate(circuit)
				raddrs = append(raddrs, pub)
			}
		}
	}

	h.addrs = raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

func containsAddr(lst []ma.Multiaddr, addr ma.Multiaddr) bool {
	for _, xaddr := range lst {
		if xaddr.Equal(addr) {
			return true
		}
	}
	return false
}

// notify
func (h *AutoRelayHost) Listen(inet.Network, ma.Multiaddr)      {}
func (h *AutoRelayHost) ListenClose(inet.Network, ma.Multiaddr) {}
func (h *AutoRelayHost) Connected(inet.Network, inet.Conn)      {}

func (h *AutoRelayHost) Disconnected(_ inet.Network, c inet.Conn) {
	p := c.RemotePeer()
	h.mx.Lock()
	defer h.mx.Unlock()
	if _, ok := h.relays[p]; ok {
		delete(h.relays, p)
		select {
		case h.disconnect <- struct{}{}:
		default:
		}
	}
}

func (h *AutoRelayHost) OpenedStream(inet.Network, inet.Stream) {}
func (h *AutoRelayHost) ClosedStream(inet.Network, inet.Stream) {}

var _ host.Host = (*AutoRelayHost)(nil)
<|endoftext|>