<commit_before>package server

import (
	"bytes"
	"context"
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	workerpkg "github.com/pachyderm/pachyderm/src/server/pkg/worker"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/clientv3/mirror"
	"go.pedge.io/lion/proto"
	"google.golang.org/grpc"
)

const (
	workerEtcdPrefix = "workers"
)

// An input/output pair for a single datum. When a worker has finished
// processing 'data', it writes the resulting hashtree to 'resp' (each job has
// its own response channel)
type datumAndResp struct {
	jobID   string // This is passed to workers, so they can annotate their logs
	datum   []*pfs.FileInfo
	respCh  chan hashtree.HashTree
	errCh   chan struct{}
	retCh   chan *datumAndResp
	retries int
}

// WorkerPool represents a pool of workers that can be used to process datums.
type WorkerPool interface {
	DataCh() chan *datumAndResp
}

type worker struct {
	ctx          context.Context
	cancel       func()
	addr         string
	workerClient workerpkg.WorkerClient
	pachClient   *client.APIClient
	retries      int
}

func (w *worker) run(dataCh chan *datumAndResp) {
	defer func() {
		protolion.Infof("goro for worker %s is exiting", w.addr)
	}()
	returnDatum := func(dr *datumAndResp) {
		dr.retries++
		select {
		case dr.retCh <- dr:
		case <-w.ctx.Done():
		}
	}
	for {
		var dr *datumAndResp
		select {
		case dr = <-dataCh:
		case <-w.ctx.Done():
			return
		}
		if dr.retries > w.retries {
			close(dr.errCh)
			continue
		}
		resp, err := w.workerClient.Process(w.ctx, &workerpkg.ProcessRequest{
			JobID: dr.jobID,
			Data:  dr.datum,
		})
		if err != nil || resp.Failed {
			protolion.Errorf("worker %s failed to process datum %v with error %s", w.addr, dr.datum, err)
			returnDatum(dr)
			continue
		}
		if resp.Tag != nil {
			var buffer bytes.Buffer
			if err := w.pachClient.GetTag(resp.Tag.Name, &buffer); err != nil {
				protolion.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", w.addr, dr.datum, err)
				returnDatum(dr)
				continue
			}
			tree, err := hashtree.Deserialize(buffer.Bytes())
			if err != nil {
				protolion.Errorf("failed to serialize hashtree after worker %s has ostensibly processed the datum %v; this is likely a bug: %v", w.addr, dr.datum, err)
				returnDatum(dr)
				continue
			}
			dr.respCh <- tree
		} else {
			protolion.Errorf("unrecognized response from worker %s when processing datum %v; this is likely a bug", w.addr, dr.datum)
			returnDatum(dr)
			continue
		}
	}
}

type workerPool struct {
	// Worker pool recieves work via this channel
	dataCh chan *datumAndResp
	// Parent of all worker contexts (see workersMap)
	ctx context.Context
	// The prefix in etcd where new workers can be discovered
	workerDir string
	// Map of worker address to workers
	workersMap map[string]*worker
	// RWMutex to protect workersMap
	workersMapMu sync.RWMutex
	// Used to check for workers added/deleted in etcd
	etcdClient *etcd.Client
	// The number of times to retry failures
	retries int
}

func (w *workerPool) discoverWorkers(ctx context.Context) {
	b := backoff.NewInfiniteBackOff()
	if err := backoff.RetryNotify(func() error {
		syncer := mirror.NewSyncer(w.etcdClient, w.workerDir, 0)
		respCh, errCh := syncer.SyncBase(ctx)
	getBaseWorkers:
		for {
			select {
			case resp, ok := <-respCh:
				if !ok {
					if err := <-errCh; err != nil {
						return err
					}
					break getBaseWorkers
				}
				for _, kv := range resp.Kvs {
					addr := path.Base(string(kv.Key))
					if err := w.addWorker(addr); err != nil {
						return err
					}
				}
			case err := <-errCh:
				if err != nil {
					return err
				}
			}
		}
		watchCh := syncer.SyncUpdates(ctx)
		protolion.Infof("watching `%s` for workers", w.workerDir)
		for {
			resp, ok := <-watchCh
			if !ok {
				return fmt.Errorf("watcher for prefix %s closed for unknown reasons", w.workerDir)
			}
			if err := resp.Err(); err != nil {
				return err
			}
			for _, event := range resp.Events {
				addr := path.Base(string(event.Kv.Key))
				switch event.Type {
				case etcd.EventTypePut:
					if err := w.addWorker(addr); err != nil {
						return err
					}
				case etcd.EventTypeDelete:
					if err := w.delWorker(addr); err != nil {
						return err
					}
				}
			}
		}
		panic("unreachable")
	}, b, func(err error, d time.Duration) error {
		if err == context.Canceled {
			return err
		}
		protolion.Errorf("error discovering workers: %v; retrying in %v", err, d)
		return nil
	}); err != context.Canceled {
		panic(fmt.Sprintf("the retry loop should not exit with a non-context-cancelled error: %v", err))
	}
}

func (w *workerPool) addWorker(addr string) error {
	w.workersMapMu.RLock()
	if worker, ok := w.workersMap[addr]; ok {
		worker.cancel()
	}
	w.workersMapMu.RUnlock()

	conn, err := grpc.Dial(fmt.Sprintf("%s:%d", addr, client.PPSWorkerPort), grpc.WithInsecure(), grpc.WithTimeout(5*time.Second))
	if err != nil {
		return err
	}
	childCtx, cancelFn := context.WithCancel(w.ctx)

	pachClient, err := client.NewInCluster()
	if err != nil {
		return err
	}

	wr := &worker{
		ctx:          childCtx,
		cancel:       cancelFn,
		addr:         addr,
		workerClient: workerpkg.NewWorkerClient(conn),
		pachClient:   pachClient,
		retries:      w.retries,
	}
	w.workersMapMu.Lock()
	w.workersMap[addr] = wr
	w.workersMapMu.Unlock()
	protolion.Infof("launching new worker at %v", addr)
	go wr.run(w.dataCh)
	return nil
}

func (w *workerPool) delWorker(addr string) error {
	w.workersMapMu.RLock()
	defer w.workersMapMu.RUnlock()
	worker, ok := w.workersMap[addr]
	if !ok {
		return fmt.Errorf("deleting worker %s which is not in worker pool", addr)
	}
	worker.cancel()
	return nil
}

func (w *workerPool) DataCh() chan *datumAndResp {
	return w.dataCh
}

func status(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]*pps.WorkerStatus, error) {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return nil, err
	}
	var result []*pps.WorkerStatus
	for _, workerClient := range workerClients {
		status, err := workerClient.Status(ctx, &types.Empty{})
		if err != nil {
			return nil, err
		}
		result = append(result, status)
	}
	return result, nil
}

func cancel(ctx context.Context, id string, etcdClient *etcd.Client,
	etcdPrefix string, jobID string, dataFilter []string) error {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return err
	}
	for _, workerClient := range workerClients {
		status, err := workerClient.Status(ctx, &types.Empty{})
		if err != nil {
			return err
		}
		if jobID == status.JobID && workerpkg.MatchDatum(dataFilter, status.Data) {
			_, err := workerClient.Cancel(ctx, &workerpkg.CancelRequest{
				DataFilters: dataFilter,
			})
			return err
		}
	}
	return nil
}

// workerPool fetches the worker pool associated with 'id', or creates one if
// none exists.
func (a *apiServer) workerPool(ctx context.Context, id string, retries int) WorkerPool {
	a.workerPoolsLock.Lock()
	defer a.workerPoolsLock.Unlock()
	workerPool, ok := a.workerPools[id]
	if !ok {
		workerPool = a.newWorkerPool(ctx, id, retries)
		a.workerPools[id] = workerPool
	}
	return workerPool
}

// newWorkerPool generates a new worker pool for the job or pipeline identified
// with 'id'.  Each 'id' used to create a new worker pool must correspond to
// a unique binary (in other words, all workers in the worker pool for 'id'
// will be running the same user binary)
func (a *apiServer) newWorkerPool(ctx context.Context, id string, retries int) WorkerPool {
	wp := &workerPool{
		ctx:        ctx,
		dataCh:     make(chan *datumAndResp),
		workerDir:  path.Join(a.etcdPrefix, workerEtcdPrefix, id),
		workersMap: make(map[string]*worker),
		etcdClient: a.etcdClient,
		retries:    retries,
	}
	// We need to make sure that the prefix ends with the trailing slash,
	// because
	if wp.workerDir[len(wp.workerDir)-1] != '/' {
		wp.workerDir += "/"
	}

	go wp.discoverWorkers(ctx)
	return wp
}

func (a *apiServer) delWorkerPool(id string) {
	a.workerPoolsLock.Lock()
	defer a.workerPoolsLock.Unlock()
	delete(a.workerPools, id)
}

func workerClients(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]workerpkg.WorkerClient, error) {
	resp, err := etcdClient.Get(ctx, path.Join(etcdPrefix, workerEtcdPrefix, id), etcd.WithPrefix())
	if err != nil {
		return nil, err
	}

	var result []workerpkg.WorkerClient
	for _, kv := range resp.Kvs {
		conn, err := grpc.Dial(fmt.Sprintf("%s:%d", string(kv.Key), client.PPSWorkerPort), grpc.WithInsecure(), grpc.WithTimeout(5*time.Second))
		if err != nil {
			return nil, err
		}
		result = append(result, workerpkg.NewWorkerClient(conn))
	}
	return result, nil
}
<commit_msg>Remove unneeded mutex.<commit_after>package server

import (
	"bytes"
	"context"
	"fmt"
	"path"
	"time"

	"github.com/gogo/protobuf/types"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	workerpkg "github.com/pachyderm/pachyderm/src/server/pkg/worker"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/clientv3/mirror"
	"go.pedge.io/lion/proto"
	"google.golang.org/grpc"
)

const (
	workerEtcdPrefix = "workers"
)

// An input/output pair for a single datum. When a worker has finished
// processing 'data', it writes the resulting hashtree to 'resp' (each job has
// its own response channel)
type datumAndResp struct {
	jobID   string // This is passed to workers, so they can annotate their logs
	datum   []*pfs.FileInfo
	respCh  chan hashtree.HashTree
	errCh   chan struct{}
	retCh   chan *datumAndResp
	retries int
}

// WorkerPool represents a pool of workers that can be used to process datums.
type WorkerPool interface {
	DataCh() chan *datumAndResp
}

type worker struct {
	ctx          context.Context
	cancel       func()
	addr         string
	workerClient workerpkg.WorkerClient
	pachClient   *client.APIClient
	retries      int
}

func (w *worker) run(dataCh chan *datumAndResp) {
	defer func() {
		protolion.Infof("goro for worker %s is exiting", w.addr)
	}()
	returnDatum := func(dr *datumAndResp) {
		dr.retries++
		select {
		case dr.retCh <- dr:
		case <-w.ctx.Done():
		}
	}
	for {
		var dr *datumAndResp
		select {
		case dr = <-dataCh:
		case <-w.ctx.Done():
			return
		}
		if dr.retries > w.retries {
			close(dr.errCh)
			continue
		}
		resp, err := w.workerClient.Process(w.ctx, &workerpkg.ProcessRequest{
			JobID: dr.jobID,
			Data:  dr.datum,
		})
		if err != nil || resp.Failed {
			protolion.Errorf("worker %s failed to process datum %v with error %s", w.addr, dr.datum, err)
			returnDatum(dr)
			continue
		}
		if resp.Tag != nil {
			var buffer bytes.Buffer
			if err := w.pachClient.GetTag(resp.Tag.Name, &buffer); err != nil {
				protolion.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", w.addr, dr.datum, err)
				returnDatum(dr)
				continue
			}
			tree, err := hashtree.Deserialize(buffer.Bytes())
			if err != nil {
				protolion.Errorf("failed to serialize hashtree after worker %s has ostensibly processed the datum %v; this is likely a bug: %v", w.addr, dr.datum, err)
				returnDatum(dr)
				continue
			}
			dr.respCh <- tree
		} else {
			protolion.Errorf("unrecognized response from worker %s when processing datum %v; this is likely a bug", w.addr, dr.datum)
			returnDatum(dr)
			continue
		}
	}
}

type workerPool struct {
	// Worker pool recieves work via this channel
	dataCh chan *datumAndResp
	// Parent of all worker contexts (see workersMap)
	ctx context.Context
	// The prefix in etcd where new workers can be discovered
	workerDir string
	// Map of worker address to workers
	workersMap map[string]*worker
	// Used to check for workers added/deleted in etcd
	etcdClient *etcd.Client
	// The number of times to retry failures
	retries int
}

func (w *workerPool) discoverWorkers(ctx context.Context) {
	b := backoff.NewInfiniteBackOff()
	if err := backoff.RetryNotify(func() error {
		syncer := mirror.NewSyncer(w.etcdClient, w.workerDir, 0)
		respCh, errCh := syncer.SyncBase(ctx)
	getBaseWorkers:
		for {
			select {
			case resp, ok := <-respCh:
				if !ok {
					if err := <-errCh; err != nil {
						return err
					}
					break getBaseWorkers
				}
				for _, kv := range resp.Kvs {
					addr := path.Base(string(kv.Key))
					if err := w.addWorker(addr); err != nil {
						return err
					}
				}
			case err := <-errCh:
				if err != nil {
					return err
				}
			}
		}
		watchCh := syncer.SyncUpdates(ctx)
		protolion.Infof("watching `%s` for workers", w.workerDir)
		for {
			resp, ok := <-watchCh
			if !ok {
				return fmt.Errorf("watcher for prefix %s closed for unknown reasons", w.workerDir)
			}
			if err := resp.Err(); err != nil {
				return err
			}
			for _, event := range resp.Events {
				addr := path.Base(string(event.Kv.Key))
				switch event.Type {
				case etcd.EventTypePut:
					if err := w.addWorker(addr); err != nil {
						return err
					}
				case etcd.EventTypeDelete:
					if err := w.delWorker(addr); err != nil {
						return err
					}
				}
			}
		}
		panic("unreachable")
	}, b, func(err error, d time.Duration) error {
		if err == context.Canceled {
			return err
		}
		protolion.Errorf("error discovering workers: %v; retrying in %v", err, d)
		return nil
	}); err != context.Canceled {
		panic(fmt.Sprintf("the retry loop should not exit with a non-context-cancelled error: %v", err))
	}
}

func (w *workerPool) addWorker(addr string) error {
	if worker, ok := w.workersMap[addr]; ok {
		worker.cancel()
	}

	conn, err := grpc.Dial(fmt.Sprintf("%s:%d", addr, client.PPSWorkerPort), grpc.WithInsecure(), grpc.WithTimeout(5*time.Second))
	if err != nil {
		return err
	}
	childCtx, cancelFn := context.WithCancel(w.ctx)

	pachClient, err := client.NewInCluster()
	if err != nil {
		return err
	}

	wr := &worker{
		ctx:          childCtx,
		cancel:       cancelFn,
		addr:         addr,
		workerClient: workerpkg.NewWorkerClient(conn),
		pachClient:   pachClient,
		retries:      w.retries,
	}
	w.workersMap[addr] = wr
	protolion.Infof("launching new worker at %v", addr)
	go wr.run(w.dataCh)
	return nil
}

func (w *workerPool) delWorker(addr string) error {
	worker, ok := w.workersMap[addr]
	if !ok {
		return fmt.Errorf("deleting worker %s which is not in worker pool", addr)
	}
	worker.cancel()
	return nil
}

func (w *workerPool) DataCh() chan *datumAndResp {
	return w.dataCh
}

func status(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]*pps.WorkerStatus, error) {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return nil, err
	}
	var result []*pps.WorkerStatus
	for _, workerClient := range workerClients {
		status, err := workerClient.Status(ctx, &types.Empty{})
		if err != nil {
			return nil, err
		}
		result = append(result, status)
	}
	return result, nil
}

func cancel(ctx context.Context, id string, etcdClient *etcd.Client,
	etcdPrefix string, jobID string, dataFilter []string) error {
	workerClients, err := workerClients(ctx, id, etcdClient, etcdPrefix)
	if err != nil {
		return err
	}
	for _, workerClient := range workerClients {
		status, err := workerClient.Status(ctx, &types.Empty{})
		if err != nil {
			return err
		}
		if jobID == status.JobID && workerpkg.MatchDatum(dataFilter, status.Data) {
			_, err := workerClient.Cancel(ctx, &workerpkg.CancelRequest{
				DataFilters: dataFilter,
			})
			return err
		}
	}
	return nil
}

// workerPool fetches the worker pool associated with 'id', or creates one if
// none exists.
func (a *apiServer) workerPool(ctx context.Context, id string, retries int) WorkerPool {
	a.workerPoolsLock.Lock()
	defer a.workerPoolsLock.Unlock()
	workerPool, ok := a.workerPools[id]
	if !ok {
		workerPool = a.newWorkerPool(ctx, id, retries)
		a.workerPools[id] = workerPool
	}
	return workerPool
}

// newWorkerPool generates a new worker pool for the job or pipeline identified
// with 'id'.  Each 'id' used to create a new worker pool must correspond to
// a unique binary (in other words, all workers in the worker pool for 'id'
// will be running the same user binary)
func (a *apiServer) newWorkerPool(ctx context.Context, id string, retries int) WorkerPool {
	wp := &workerPool{
		ctx:        ctx,
		dataCh:     make(chan *datumAndResp),
		workerDir:  path.Join(a.etcdPrefix, workerEtcdPrefix, id),
		workersMap: make(map[string]*worker),
		etcdClient: a.etcdClient,
		retries:    retries,
	}
	// We need to make sure that the prefix ends with the trailing slash,
	// because
	if wp.workerDir[len(wp.workerDir)-1] != '/' {
		wp.workerDir += "/"
	}

	go wp.discoverWorkers(ctx)
	return wp
}

func (a *apiServer) delWorkerPool(id string) {
	a.workerPoolsLock.Lock()
	defer a.workerPoolsLock.Unlock()
	delete(a.workerPools, id)
}

func workerClients(ctx context.Context, id string, etcdClient *etcd.Client, etcdPrefix string) ([]workerpkg.WorkerClient, error) {
	resp, err := etcdClient.Get(ctx, path.Join(etcdPrefix, workerEtcdPrefix, id), etcd.WithPrefix())
	if err != nil {
		return nil, err
	}

	var result []workerpkg.WorkerClient
	for _, kv := range resp.Kvs {
		conn, err := grpc.Dial(fmt.Sprintf("%s:%d", string(kv.Key), client.PPSWorkerPort), grpc.WithInsecure(), grpc.WithTimeout(5*time.Second))
		if err != nil {
			return nil, err
		}
		result = append(result, workerpkg.NewWorkerClient(conn))
	}
	return result, nil
}
<|endoftext|>


<commit_before>package gannoy

import (
	"fmt"
	"math"
	"sort"
	"sync"

	"github.com/gansidui/priority_queue"
)

type GannoyIndex struct {
	meta      meta
	maps      Maps
	tree      int
	dim       int
	distance  Distance
	random    Random
	nodes     Nodes
	K         int
	buildChan chan buildArgs
}

func NewGannoyIndex(metaFile string, distance Distance, random Random) (GannoyIndex, error) {

	meta, err := loadMeta(metaFile)
	if err != nil {
		return GannoyIndex{}, err
	}
	tree := meta.tree
	dim := meta.dim

	ann := meta.treePath()
	maps := meta.mapPath()

	// K := 3
	K := 50
	gannoy := GannoyIndex{
		meta:      meta,
		maps:      newMaps(maps),
		tree:      tree,
		dim:       dim,
		distance:  distance,
		random:    random,
		K:         K,
		nodes:     newNodes(ann, tree, dim, K),
		buildChan: make(chan buildArgs, 1),
	}
	go gannoy.builder()
	return gannoy, nil
}

func (g GannoyIndex) Tree() {
	for i, root := range g.meta.roots() {
		g.walk(i, g.nodes.getNode(root), root, 0)
	}
}

func (g *GannoyIndex) AddItem(id int, w []float64) error {
	args := buildArgs{action: ADD, id: id, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) RemoveItem(id int) error {
	args := buildArgs{action: DELETE, id: id, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g GannoyIndex) GetNnsByItem(id, n, searchK int) []int {
	m := g.nodes.getNode(g.maps.getIndex(id))
	if !m.isLeaf() {
		return []int{}
	}
	indices := g.getAllNns(m.v, n, searchK)
	ids := make([]int, len(indices))
	for i, index := range indices {
		ids[i] = g.maps.getId(index)
	}
	return ids
}

func (g GannoyIndex) getAllNns(v []float64, n, searchK int) []int {
	if searchK == -1 {
		searchK = n * g.tree
	}

	q := priority_queue.New()
	for _, root := range g.meta.roots() {
		q.Push(&Queue{priority: math.Inf(1), value: root})
	}

	nns := []int{}
	for len(nns) < searchK && q.Len() > 0 {
		top := q.Top().(*Queue)
		d := top.priority
		i := top.value

		nd := g.nodes.getNode(i)
		q.Pop()
		if nd.isLeaf() {
			nns = append(nns, i)
		} else if nd.nDescendants <= g.K {
			dst := nd.children
			nns = append(nns, dst...)
		} else {
			margin := g.distance.margin(nd, v, g.dim)
			q.Push(&Queue{priority: math.Min(d, +margin), value: nd.children[1]})
			q.Push(&Queue{priority: math.Min(d, -margin), value: nd.children[0]})
		}
	}

	sort.Ints(nns)
	nnsDist := []Dist{}
	last := -1
	for _, j := range nns {
		if j == last {
			continue
		}
		last = j
		nnsDist = append(nnsDist, Dist{distance: g.distance.distance(v, g.nodes.getNode(j).v, g.dim), item: j})
	}

	m := len(nnsDist)
	p := m
	if n < m {
		p = n
	}

	result := []int{}
	sort.Slice(nnsDist, func(i, j int) bool {
		return nnsDist[i].distance < nnsDist[j].distance
	})
	for i := 0; i < p; i++ {
		result = append(result, nnsDist[i].item)
	}

	return result
}

func (g *GannoyIndex) addItem(id int, w []float64) error {
	n := g.nodes.newNode()
	n.v = w
	n.parents = make([]int, g.tree)
	err := n.save()
	if err != nil {
		return err
	}
	// fmt.Printf("id %d\n", n.id)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for index := range buildChan {
			// fmt.Printf("root: %d\n", g.meta.roots()[index])
			g.build(index, g.meta.roots()[index], n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}

	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)
	g.maps.add(n.id, id)

	return nil
}

func (g *GannoyIndex) build(index, root int, n Node) {
	if root == -1 {
		// 最初のノード
		n.parents[index] = -1
		n.save()
		g.meta.updateRoot(index, n.id)
		return
	}
	item := g.findBranchByVector(root, n.v)
	found := g.nodes.getNode(item)
	// fmt.Printf("Found %d\n", item)

	org_parent := found.parents[index]
	if found.isBucket() && len(found.children) < g.K {
		// ノードに余裕があれば追加
		// fmt.Printf("pattern bucket\n")
		n.updateParents(index, item)
		found.nDescendants++
		found.children = append(found.children, n.id)
		found.save()
	} else {
		// ノードが上限またはリーフノードであれば新しいノードを追加
		willDelete := false
		var indices []int
		if found.isLeaf() {
			// fmt.Printf("pattern leaf node\n")
			indices = []int{item, n.id}
		} else {
			// fmt.Printf("pattern full backet\n")
			indices = append(found.children, n.id)
			willDelete = true
		}

		m := g.makeTree(index, org_parent, indices)
		// fmt.Printf("m: %d, org_parent: %d\n", m, org_parent)
		if org_parent == -1 {
			// rootノードの入れ替え
			g.meta.updateRoot(index, m)
		} else {
			parent := g.nodes.getNode(org_parent)
			parent.nDescendants++
			children := make([]int, len(parent.children))
			for i, child := range parent.children {
				if child == item {
					// 新しいノードに変更
					children[i] = m
				} else {
					// 既存のノードのまま
					children[i] = child
				}
			}
			parent.children = children
			parent.save()

		}
		if willDelete {
			found.destroy()
		}
	}
}

func (g *GannoyIndex) removeItem(id int) error {
	index := g.maps.getIndex(id)
	n := g.nodes.getNode(index)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for root := range buildChan {
			g.remove(root, n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}
	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)

	g.maps.remove(n.id, id)
	n.ref = false
	n.save()

	return nil
}

func (g *GannoyIndex) remove(root int, node Node) {
	parent := g.nodes.getNode(node.parents[root])
	if parent.isBucket() && len(parent.children) > 2 {
		// fmt.Printf("pattern bucket\n")
		target := -1
		for i, child := range parent.children {
			if child == node.id {
				target = i
			}
		}
		if target == -1 {
			return
		}
		children := append(parent.children[:target], parent.children[(target+1):]...)
		parent.nDescendants--
		parent.children = children
		parent.save()
	} else {
		// fmt.Printf("pattern leaf node\n")
		var other int
		for _, child := range parent.children {
			if child != node.id {
				other = child
			}
		}
		grandParent := g.nodes.getNode(parent.parents[root])
		children := []int{}
		for _, child := range grandParent.children {
			if child == node.parents[root] {
				children = append(children, other)
			} else {
				children = append(children, child)
			}
		}
		grandParent.nDescendants--
		grandParent.children = children
		grandParent.save()

		otherNode := g.nodes.getNode(other)
		otherNode.parents[root] = parent.parents[root]
		otherNode.save()

		parent.ref = false
		parent.save()
	}
}

func (g GannoyIndex) findBranchByVector(index int, v []float64) int {
	node := g.nodes.getNode(index)
	if node.isLeaf() || node.isBucket() {
		return index
	}
	side := g.distance.side(node, v, g.dim, g.random)
	return g.findBranchByVector(node.children[side], v)
}

func (g *GannoyIndex) makeTree(root, parent int, indices []int) int {
	if len(indices) == 1 {
		n := g.nodes.getNode(indices[0])
		if len(n.parents) == 0 {
			n.parents = make([]int, g.tree)
		}
		n.updateParents(root, parent)
		return indices[0]
	}

	if len(indices) <= g.K {
		m := g.nodes.newNode()
		m.parents = make([]int, g.tree)
		m.nDescendants = len(indices)
		m.parents[root] = parent
		m.children = indices
		m.save()
		for _, child := range indices {
			c := g.nodes.getNode(child)
			if len(c.parents) == 0 {
				c.parents = make([]int, g.tree)
			}
			c.updateParents(root, m.id)
		}
		return m.id
	}

	children := make([]Node, len(indices))
	for i, idx := range indices {
		children[i] = g.nodes.getNode(idx)
	}

	childrenIndices := [2][]int{[]int{}, []int{}}

	m := g.nodes.newNode()
	m.parents = make([]int, g.tree)
	m.nDescendants = len(indices)
	m.parents[root] = parent

	m = g.distance.createSplit(children, g.dim, g.random, m)
	for _, idx := range indices {
		n := g.nodes.getNode(idx)
		side := g.distance.side(m, n.v, g.dim, g.random)
		childrenIndices[side] = append(childrenIndices[side], idx)
	}

	for len(childrenIndices[0]) == 0 || len(childrenIndices[1]) == 0 {
		childrenIndices[0] = []int{}
		childrenIndices[1] = []int{}
		for z := 0; z < g.dim; z++ {
			m.v[z] = 0.0
		}
		for _, idx := range indices {
			side := g.random.flip()
			childrenIndices[side] = append(childrenIndices[side], idx)
		}
	}

	var flip int
	if len(childrenIndices[0]) > len(childrenIndices[1]) {
		flip = 1
	}

	m.save()
	for side := 0; side < 2; side++ {
		m.children[side^flip] = g.makeTree(root, m.id, childrenIndices[side^flip])
	}
	m.save()

	return m.id
}

type buildArgs struct {
	action int
	id     int
	w      []float64
	result chan error
}

func (g *GannoyIndex) builder() {
	for args := range g.buildChan {
		switch args.action {
		case ADD:
			args.result <- g.addItem(args.id, args.w)
		case DELETE:
			args.result <- g.removeItem(args.id)
		}
	}
}

func (g GannoyIndex) walk(root int, node Node, id, tab int) {
	for i := 0; i < tab*2; i++ {
		fmt.Print(" ")
	}
	fmt.Printf("%d [%d] (%d) [nDescendants: %d, v: %v]\n", id, g.maps.getId(id), node.parents[root], node.nDescendants, node.v)
	if !node.isLeaf() {
		for _, child := range node.children {
			g.walk(root, g.nodes.getNode(child), child, tab+1)
		}
	}
}
<commit_msg>Implemented a feature for updating node from trees.<commit_after>package gannoy

import (
	"fmt"
	"math"
	"sort"
	"sync"

	"github.com/gansidui/priority_queue"
)

type GannoyIndex struct {
	meta      meta
	maps      Maps
	tree      int
	dim       int
	distance  Distance
	random    Random
	nodes     Nodes
	K         int
	buildChan chan buildArgs
}

func NewGannoyIndex(metaFile string, distance Distance, random Random) (GannoyIndex, error) {

	meta, err := loadMeta(metaFile)
	if err != nil {
		return GannoyIndex{}, err
	}
	tree := meta.tree
	dim := meta.dim

	ann := meta.treePath()
	maps := meta.mapPath()

	// K := 3
	K := 50
	gannoy := GannoyIndex{
		meta:      meta,
		maps:      newMaps(maps),
		tree:      tree,
		dim:       dim,
		distance:  distance,
		random:    random,
		K:         K,
		nodes:     newNodes(ann, tree, dim, K),
		buildChan: make(chan buildArgs, 1),
	}
	go gannoy.builder()
	return gannoy, nil
}

func (g GannoyIndex) Tree() {
	for i, root := range g.meta.roots() {
		g.walk(i, g.nodes.getNode(root), root, 0)
	}
}

func (g *GannoyIndex) AddItem(id int, w []float64) error {
	args := buildArgs{action: ADD, id: id, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) RemoveItem(id int) error {
	args := buildArgs{action: DELETE, id: id, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) UpdateItem(id int, w []float64) error {
	args := buildArgs{action: UPDATE, id: id, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g GannoyIndex) GetNnsByItem(id, n, searchK int) []int {
	m := g.nodes.getNode(g.maps.getIndex(id))
	if !m.isLeaf() {
		return []int{}
	}
	indices := g.getAllNns(m.v, n, searchK)
	ids := make([]int, len(indices))
	for i, index := range indices {
		ids[i] = g.maps.getId(index)
	}
	return ids
}

func (g GannoyIndex) getAllNns(v []float64, n, searchK int) []int {
	if searchK == -1 {
		searchK = n * g.tree
	}

	q := priority_queue.New()
	for _, root := range g.meta.roots() {
		q.Push(&Queue{priority: math.Inf(1), value: root})
	}

	nns := []int{}
	for len(nns) < searchK && q.Len() > 0 {
		top := q.Top().(*Queue)
		d := top.priority
		i := top.value

		nd := g.nodes.getNode(i)
		q.Pop()
		if nd.isLeaf() {
			nns = append(nns, i)
		} else if nd.nDescendants <= g.K {
			dst := nd.children
			nns = append(nns, dst...)
		} else {
			margin := g.distance.margin(nd, v, g.dim)
			q.Push(&Queue{priority: math.Min(d, +margin), value: nd.children[1]})
			q.Push(&Queue{priority: math.Min(d, -margin), value: nd.children[0]})
		}
	}

	sort.Ints(nns)
	nnsDist := []Dist{}
	last := -1
	for _, j := range nns {
		if j == last {
			continue
		}
		last = j
		nnsDist = append(nnsDist, Dist{distance: g.distance.distance(v, g.nodes.getNode(j).v, g.dim), item: j})
	}

	m := len(nnsDist)
	p := m
	if n < m {
		p = n
	}

	result := []int{}
	sort.Slice(nnsDist, func(i, j int) bool {
		return nnsDist[i].distance < nnsDist[j].distance
	})
	for i := 0; i < p; i++ {
		result = append(result, nnsDist[i].item)
	}

	return result
}

func (g *GannoyIndex) addItem(id int, w []float64) error {
	n := g.nodes.newNode()
	n.v = w
	n.parents = make([]int, g.tree)
	err := n.save()
	if err != nil {
		return err
	}
	// fmt.Printf("id %d\n", n.id)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for index := range buildChan {
			// fmt.Printf("root: %d\n", g.meta.roots()[index])
			g.build(index, g.meta.roots()[index], n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}

	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)
	g.maps.add(n.id, id)

	return nil
}

func (g *GannoyIndex) build(index, root int, n Node) {
	if root == -1 {
		// 最初のノード
		n.parents[index] = -1
		n.save()
		g.meta.updateRoot(index, n.id)
		return
	}
	item := g.findBranchByVector(root, n.v)
	found := g.nodes.getNode(item)
	// fmt.Printf("Found %d\n", item)

	org_parent := found.parents[index]
	if found.isBucket() && len(found.children) < g.K {
		// ノードに余裕があれば追加
		// fmt.Printf("pattern bucket\n")
		n.updateParents(index, item)
		found.nDescendants++
		found.children = append(found.children, n.id)
		found.save()
	} else {
		// ノードが上限またはリーフノードであれば新しいノードを追加
		willDelete := false
		var indices []int
		if found.isLeaf() {
			// fmt.Printf("pattern leaf node\n")
			indices = []int{item, n.id}
		} else {
			// fmt.Printf("pattern full backet\n")
			indices = append(found.children, n.id)
			willDelete = true
		}

		m := g.makeTree(index, org_parent, indices)
		// fmt.Printf("m: %d, org_parent: %d\n", m, org_parent)
		if org_parent == -1 {
			// rootノードの入れ替え
			g.meta.updateRoot(index, m)
		} else {
			parent := g.nodes.getNode(org_parent)
			parent.nDescendants++
			children := make([]int, len(parent.children))
			for i, child := range parent.children {
				if child == item {
					// 新しいノードに変更
					children[i] = m
				} else {
					// 既存のノードのまま
					children[i] = child
				}
			}
			parent.children = children
			parent.save()

		}
		if willDelete {
			found.destroy()
		}
	}
}

func (g *GannoyIndex) removeItem(id int) error {
	index := g.maps.getIndex(id)
	n := g.nodes.getNode(index)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for root := range buildChan {
			g.remove(root, n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}
	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)

	g.maps.remove(n.id, id)
	n.ref = false
	n.save()

	return nil
}

func (g *GannoyIndex) remove(root int, node Node) {
	parent := g.nodes.getNode(node.parents[root])
	if parent.isBucket() && len(parent.children) > 2 {
		// fmt.Printf("pattern bucket\n")
		target := -1
		for i, child := range parent.children {
			if child == node.id {
				target = i
			}
		}
		if target == -1 {
			return
		}
		children := append(parent.children[:target], parent.children[(target+1):]...)
		parent.nDescendants--
		parent.children = children
		parent.save()
	} else {
		// fmt.Printf("pattern leaf node\n")
		var other int
		for _, child := range parent.children {
			if child != node.id {
				other = child
			}
		}
		grandParent := g.nodes.getNode(parent.parents[root])
		children := []int{}
		for _, child := range grandParent.children {
			if child == node.parents[root] {
				children = append(children, other)
			} else {
				children = append(children, child)
			}
		}
		grandParent.nDescendants--
		grandParent.children = children
		grandParent.save()

		otherNode := g.nodes.getNode(other)
		otherNode.parents[root] = parent.parents[root]
		otherNode.save()

		parent.ref = false
		parent.save()
	}
}

func (g GannoyIndex) findBranchByVector(index int, v []float64) int {
	node := g.nodes.getNode(index)
	if node.isLeaf() || node.isBucket() {
		return index
	}
	side := g.distance.side(node, v, g.dim, g.random)
	return g.findBranchByVector(node.children[side], v)
}

func (g *GannoyIndex) makeTree(root, parent int, indices []int) int {
	if len(indices) == 1 {
		n := g.nodes.getNode(indices[0])
		if len(n.parents) == 0 {
			n.parents = make([]int, g.tree)
		}
		n.updateParents(root, parent)
		return indices[0]
	}

	if len(indices) <= g.K {
		m := g.nodes.newNode()
		m.parents = make([]int, g.tree)
		m.nDescendants = len(indices)
		m.parents[root] = parent
		m.children = indices
		m.save()
		for _, child := range indices {
			c := g.nodes.getNode(child)
			if len(c.parents) == 0 {
				c.parents = make([]int, g.tree)
			}
			c.updateParents(root, m.id)
		}
		return m.id
	}

	children := make([]Node, len(indices))
	for i, idx := range indices {
		children[i] = g.nodes.getNode(idx)
	}

	childrenIndices := [2][]int{[]int{}, []int{}}

	m := g.nodes.newNode()
	m.parents = make([]int, g.tree)
	m.nDescendants = len(indices)
	m.parents[root] = parent

	m = g.distance.createSplit(children, g.dim, g.random, m)
	for _, idx := range indices {
		n := g.nodes.getNode(idx)
		side := g.distance.side(m, n.v, g.dim, g.random)
		childrenIndices[side] = append(childrenIndices[side], idx)
	}

	for len(childrenIndices[0]) == 0 || len(childrenIndices[1]) == 0 {
		childrenIndices[0] = []int{}
		childrenIndices[1] = []int{}
		for z := 0; z < g.dim; z++ {
			m.v[z] = 0.0
		}
		for _, idx := range indices {
			side := g.random.flip()
			childrenIndices[side] = append(childrenIndices[side], idx)
		}
	}

	var flip int
	if len(childrenIndices[0]) > len(childrenIndices[1]) {
		flip = 1
	}

	m.save()
	for side := 0; side < 2; side++ {
		m.children[side^flip] = g.makeTree(root, m.id, childrenIndices[side^flip])
	}
	m.save()

	return m.id
}

type buildArgs struct {
	action int
	id     int
	w      []float64
	result chan error
}

func (g *GannoyIndex) builder() {
	for args := range g.buildChan {
		switch args.action {
		case ADD:
			args.result <- g.addItem(args.id, args.w)
		case DELETE:
			args.result <- g.removeItem(args.id)
		case UPDATE:
			err := g.removeItem(args.id)
			if err != nil {
				args.result <- err
			} else {
				args.result <- g.addItem(args.id, args.w)
			}
		}
	}
}

func (g GannoyIndex) walk(root int, node Node, id, tab int) {
	for i := 0; i < tab*2; i++ {
		fmt.Print(" ")
	}
	fmt.Printf("%d [%d] (%d) [nDescendants: %d, v: %v]\n", id, g.maps.getId(id), node.parents[root], node.nDescendants, node.v)
	if !node.isLeaf() {
		for _, child := range node.children {
			g.walk(root, g.nodes.getNode(child), child, tab+1)
		}
	}
}
<|endoftext|>


<commit_before>package ice

import (
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/pion/stun"
	"github.com/pion/transport/vnet"
	"github.com/pion/turn"
)

const (
	stunGatherTimeout = time.Second * 5
)

func (a *Agent) localInterfaces(networkTypes []NetworkType) ([]net.IP, error) {
	ips := []net.IP{}
	ifaces, err := a.net.Interfaces()
	if err != nil {
		return ips, err
	}

	var IPv4Requested, IPv6Requested bool
	for _, typ := range networkTypes {
		if typ.IsIPv4() {
			IPv4Requested = true
		}

		if typ.IsIPv6() {
			IPv6Requested = true
		}
	}

	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}

		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}

		for _, addr := range addrs {
			var ip net.IP
			switch addr := addr.(type) {
			case *net.IPNet:
				ip = addr.IP
			case *net.IPAddr:
				ip = addr.IP

			}
			if ip == nil || ip.IsLoopback() {
				continue
			}

			if ipv4 := ip.To4(); ipv4 == nil {
				if !IPv6Requested {
					continue
				} else if !isSupportedIPv6(ip) {
					continue
				}
			} else if !IPv4Requested {
				continue
			}

			ips = append(ips, ip)
		}
	}
	return ips, nil
}

func (a *Agent) listenUDP(portMax, portMin int, network string, laddr *net.UDPAddr) (vnet.UDPPacketConn, error) {
	if (laddr.Port != 0) || ((portMin == 0) && (portMax == 0)) {
		return a.net.ListenUDP(network, laddr)
	}
	var i, j int
	i = portMin
	if i == 0 {
		i = 1
	}
	j = portMax
	if j == 0 {
		j = 0xFFFF
	}
	for i <= j {
		laddr = &net.UDPAddr{IP: laddr.IP, Port: i}
		c, e := a.net.ListenUDP(network, laddr)
		if e == nil {
			return c, e
		}
		a.log.Debugf("failed to listen %s: %v", laddr.String(), e)
		i++
	}
	return nil, ErrPort
}

// GatherCandidates initiates the trickle based gathering process.
func (a *Agent) GatherCandidates() error {
	gatherErrChan := make(chan error, 1)

	runErr := a.run(func(agent *Agent) {
		if a.gatheringState != GatheringStateNew {
			gatherErrChan <- ErrMultipleGatherAttempted
			return
		} else if a.onCandidateHdlr == nil {
			gatherErrChan <- ErrNoOnCandidateHandler
			return
		}

		go a.gatherCandidates()

		gatherErrChan <- nil
	})
	if runErr != nil {
		return runErr
	}
	return <-gatherErrChan
}

func (a *Agent) gatherCandidates() {
	gatherStateUpdated := make(chan bool)
	if err := a.run(func(agent *Agent) {
		a.gatheringState = GatheringStateGathering
		close(gatherStateUpdated)
	}); err != nil {
		a.log.Warnf("failed to set gatheringState to GatheringStateGathering for gatherCandidates: %v", err)
		return
	}
	<-gatherStateUpdated

	for _, t := range a.candidateTypes {
		switch t {
		case CandidateTypeHost:
			a.gatherCandidatesLocal(a.networkTypes)
		case CandidateTypeServerReflexive:
			a.gatherCandidatesSrflx(a.urls, a.networkTypes)
		case CandidateTypeRelay:
			if err := a.gatherCandidatesRelay(a.urls); err != nil {
				a.log.Errorf("Failed to gather relay candidates: %v\n", err)
			}
		}
	}

	if err := a.run(func(agent *Agent) {
		if a.onCandidateHdlr != nil {
			go a.onCandidateHdlr(nil)
		}
	}); err != nil {
		a.log.Warnf("Failed to run onCandidateHdlr task: %v\n", err)
		return
	}

	if err := a.run(func(agent *Agent) {
		a.gatheringState = GatheringStateComplete
	}); err != nil {
		a.log.Warnf("Failed to update gatheringState: %v\n", err)
		return
	}
}

func (a *Agent) gatherCandidatesLocal(networkTypes []NetworkType) {
	var wg sync.WaitGroup
	defer wg.Wait()

	localIPs, err := a.localInterfaces(networkTypes)
	if err != nil {
		a.log.Warnf("failed to iterate local interfaces, host candidates will not be gathered %s", err)
		return
	}

	wg.Add(len(localIPs) * len(supportedNetworks))
	for _, ip := range localIPs {
		for _, network := range supportedNetworks {
			go func(network string, ip net.IP) {
				defer wg.Done()
				conn, err := a.listenUDP(int(a.portmax), int(a.portmin), network, &net.UDPAddr{IP: ip, Port: 0})
				if err != nil {
					a.log.Warnf("could not listen %s %s\n", network, ip)
					return
				}

				address := ip.String()
				if a.mDNSMode == MulticastDNSModeQueryAndGather {
					address = a.mDNSName
				}

				port := conn.LocalAddr().(*net.UDPAddr).Port

				hostConfig := CandidateHostConfig{
					Network:   network,
					Address:   address,
					Port:      port,
					Component: ComponentRTP,
				}

				c, err := NewCandidateHost(&hostConfig)
				if err != nil {
					a.log.Warnf("Failed to create host candidate: %s %s %d: %v\n", network, ip, port, err)
					return
				}

				if a.mDNSMode == MulticastDNSModeQueryAndGather {
					if err = c.setIP(ip); err != nil {
						a.log.Warnf("Failed to create host candidate: %s %s %d: %v\n", network, ip, port, err)
						return
					}
				}

				if err := a.run(func(agent *Agent) {
					c.start(a, conn)
					a.addCandidate(c)

					if a.onCandidateHdlr != nil {
						go a.onCandidateHdlr(c)
					}
				}); err != nil {
					a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
				}
			}(network, ip)
		}
	}
}

func (a *Agent) gatherCandidatesSrflx(urls []*URL, networkTypes []NetworkType) {
	for _, networkType := range networkTypes {
		network := networkType.String()
		for _, url := range urls {
			if url.Scheme != SchemeTypeSTUN {
				continue
			}

			hostPort := fmt.Sprintf("%s:%d", url.Host, url.Port)
			serverAddr, err := a.net.ResolveUDPAddr(network, hostPort)
			if err != nil {
				a.log.Warnf("failed to resolve stun host: %s: %v", hostPort, err)
				continue
			}

			conn, err := a.listenUDP(int(a.portmax), int(a.portmin), network, &net.UDPAddr{IP: nil, Port: 0})
			if err != nil {
				a.log.Warnf("Failed to listen on %s for %s: %v\n", conn.LocalAddr().String(), serverAddr.String(), err)
				continue
			}

			xoraddr, err := getXORMappedAddr(conn, serverAddr, stunGatherTimeout)
			if err != nil {
				a.log.Warnf("could not get server reflexive address %s %s: %v\n", network, url, err)
				continue
			}

			laddr := conn.LocalAddr().(*net.UDPAddr)
			ip := xoraddr.IP
			port := xoraddr.Port
			relIP := laddr.IP.String()
			relPort := laddr.Port

			srflxConfig := CandidateServerReflexiveConfig{
				Network:   network,
				Address:   ip.String(),
				Port:      port,
				Component: ComponentRTP,
				RelAddr:   relIP,
				RelPort:   relPort,
			}
			c, err := NewCandidateServerReflexive(&srflxConfig)
			if err != nil {
				a.log.Warnf("Failed to create server reflexive candidate: %s %s %d: %v\n", network, ip, port, err)
				continue
			}

			if err := a.run(func(agent *Agent) {
				c.start(a, conn)
				a.addCandidate(c)

				if a.onCandidateHdlr != nil {
					go a.onCandidateHdlr(c)
				}
			}); err != nil {
				a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
			}
		}
	}

}

func (a *Agent) gatherCandidatesRelay(urls []*URL) error {
	network := NetworkTypeUDP4.String() // TODO IPv6
	for _, url := range urls {
		switch {
		case url.Scheme != SchemeTypeTURN:
			continue
		case url.Username == "":
			return ErrUsernameEmpty
		case url.Password == "":
			return ErrPasswordEmpty
		}

		locConn, err := a.net.ListenPacket(network, "0.0.0.0:0")
		if err != nil {
			return err
		}

		client, err := turn.NewClient(&turn.ClientConfig{
			TURNServerAddr: fmt.Sprintf("%s:%d", url.Host, url.Port),
			Conn:           locConn,
			Username:       url.Username,
			Password:       url.Password,
			LoggerFactory:  a.loggerFactory,
			Net:            a.net,
		})
		if err != nil {
			return err
		}

		err = client.Listen()
		if err != nil {
			return err
		}

		relayConn, err := client.Allocate()
		if err != nil {
			return err
		}

		laddr := locConn.LocalAddr().(*net.UDPAddr)
		raddr := relayConn.LocalAddr().(*net.UDPAddr)

		relayConfig := CandidateRelayConfig{
			Network:   network,
			Component: ComponentRTP,
			Address:   raddr.IP.String(),
			Port:      raddr.Port,
			RelAddr:   laddr.IP.String(),
			RelPort:   laddr.Port,
			OnClose: func() error {
				client.Close()
				return locConn.Close()
			},
		}
		candidate, err := NewCandidateRelay(&relayConfig)
		if err != nil {
			a.log.Warnf("Failed to create relay candidate: %s %s: %v\n",
				network, raddr.String(), err)
			continue
		}

		if err := a.run(func(agent *Agent) {
			candidate.start(a, relayConn)
			a.addCandidate(candidate)

			if a.onCandidateHdlr != nil {
				go a.onCandidateHdlr(candidate)
			}
		}); err != nil {
			a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
		}
	}

	return nil
}

// getXORMappedAddr initiates a stun requests to serverAddr using conn, reads the response and returns
// the XORMappedAddress returned by the stun server.
//
// Adapted from stun v0.2.
func getXORMappedAddr(conn net.PacketConn, serverAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error) {
	if deadline > 0 {
		if err := conn.SetReadDeadline(time.Now().Add(deadline)); err != nil {
			return nil, err
		}
	}
	defer func() {
		if deadline > 0 {
			_ = conn.SetReadDeadline(time.Time{})
		}
	}()
	resp, err := stunRequest(
		func(p []byte) (int, error) {
			n, _, errr := conn.ReadFrom(p)
			return n, errr
		},
		func(b []byte) (int, error) {
			return conn.WriteTo(b, serverAddr)
		},
	)
	if err != nil {
		return nil, err
	}
	var addr stun.XORMappedAddress
	if err = addr.GetFrom(resp); err != nil {
		return nil, fmt.Errorf("failed to get XOR-MAPPED-ADDRESS response: %v", err)
	}
	return &addr, nil
}

func stunRequest(read func([]byte) (int, error), write func([]byte) (int, error)) (*stun.Message, error) {
	req, err := stun.Build(stun.BindingRequest, stun.TransactionID)
	if err != nil {
		return nil, err
	}
	if _, err = write(req.Raw); err != nil {
		return nil, err
	}
	const maxMessageSize = 1280
	bs := make([]byte, maxMessageSize)
	n, err := read(bs)
	if err != nil {
		return nil, err
	}
	res := &stun.Message{Raw: bs[:n]}
	if err := res.Decode(); err != nil {
		return nil, err
	}
	return res, nil
}
<commit_msg>Fix connection nil pointer<commit_after>package ice

import (
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/pion/stun"
	"github.com/pion/transport/vnet"
	"github.com/pion/turn"
)

const (
	stunGatherTimeout = time.Second * 5
)

func (a *Agent) localInterfaces(networkTypes []NetworkType) ([]net.IP, error) {
	ips := []net.IP{}
	ifaces, err := a.net.Interfaces()
	if err != nil {
		return ips, err
	}

	var IPv4Requested, IPv6Requested bool
	for _, typ := range networkTypes {
		if typ.IsIPv4() {
			IPv4Requested = true
		}

		if typ.IsIPv6() {
			IPv6Requested = true
		}
	}

	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}

		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}

		for _, addr := range addrs {
			var ip net.IP
			switch addr := addr.(type) {
			case *net.IPNet:
				ip = addr.IP
			case *net.IPAddr:
				ip = addr.IP

			}
			if ip == nil || ip.IsLoopback() {
				continue
			}

			if ipv4 := ip.To4(); ipv4 == nil {
				if !IPv6Requested {
					continue
				} else if !isSupportedIPv6(ip) {
					continue
				}
			} else if !IPv4Requested {
				continue
			}

			ips = append(ips, ip)
		}
	}
	return ips, nil
}

func (a *Agent) listenUDP(portMax, portMin int, network string, laddr *net.UDPAddr) (vnet.UDPPacketConn, error) {
	if (laddr.Port != 0) || ((portMin == 0) && (portMax == 0)) {
		return a.net.ListenUDP(network, laddr)
	}
	var i, j int
	i = portMin
	if i == 0 {
		i = 1
	}
	j = portMax
	if j == 0 {
		j = 0xFFFF
	}
	for i <= j {
		laddr = &net.UDPAddr{IP: laddr.IP, Port: i}
		c, e := a.net.ListenUDP(network, laddr)
		if e == nil {
			return c, e
		}
		a.log.Debugf("failed to listen %s: %v", laddr.String(), e)
		i++
	}
	return nil, ErrPort
}

// GatherCandidates initiates the trickle based gathering process.
func (a *Agent) GatherCandidates() error {
	gatherErrChan := make(chan error, 1)

	runErr := a.run(func(agent *Agent) {
		if a.gatheringState != GatheringStateNew {
			gatherErrChan <- ErrMultipleGatherAttempted
			return
		} else if a.onCandidateHdlr == nil {
			gatherErrChan <- ErrNoOnCandidateHandler
			return
		}

		go a.gatherCandidates()

		gatherErrChan <- nil
	})
	if runErr != nil {
		return runErr
	}
	return <-gatherErrChan
}

func (a *Agent) gatherCandidates() {
	gatherStateUpdated := make(chan bool)
	if err := a.run(func(agent *Agent) {
		a.gatheringState = GatheringStateGathering
		close(gatherStateUpdated)
	}); err != nil {
		a.log.Warnf("failed to set gatheringState to GatheringStateGathering for gatherCandidates: %v", err)
		return
	}
	<-gatherStateUpdated

	for _, t := range a.candidateTypes {
		switch t {
		case CandidateTypeHost:
			a.gatherCandidatesLocal(a.networkTypes)
		case CandidateTypeServerReflexive:
			a.gatherCandidatesSrflx(a.urls, a.networkTypes)
		case CandidateTypeRelay:
			if err := a.gatherCandidatesRelay(a.urls); err != nil {
				a.log.Errorf("Failed to gather relay candidates: %v\n", err)
			}
		}
	}

	if err := a.run(func(agent *Agent) {
		if a.onCandidateHdlr != nil {
			go a.onCandidateHdlr(nil)
		}
	}); err != nil {
		a.log.Warnf("Failed to run onCandidateHdlr task: %v\n", err)
		return
	}

	if err := a.run(func(agent *Agent) {
		a.gatheringState = GatheringStateComplete
	}); err != nil {
		a.log.Warnf("Failed to update gatheringState: %v\n", err)
		return
	}
}

func (a *Agent) gatherCandidatesLocal(networkTypes []NetworkType) {
	var wg sync.WaitGroup
	defer wg.Wait()

	localIPs, err := a.localInterfaces(networkTypes)
	if err != nil {
		a.log.Warnf("failed to iterate local interfaces, host candidates will not be gathered %s", err)
		return
	}

	wg.Add(len(localIPs) * len(supportedNetworks))
	for _, ip := range localIPs {
		for _, network := range supportedNetworks {
			go func(network string, ip net.IP) {
				defer wg.Done()
				conn, err := a.listenUDP(int(a.portmax), int(a.portmin), network, &net.UDPAddr{IP: ip, Port: 0})
				if err != nil {
					a.log.Warnf("could not listen %s %s\n", network, ip)
					return
				}

				address := ip.String()
				if a.mDNSMode == MulticastDNSModeQueryAndGather {
					address = a.mDNSName
				}

				port := conn.LocalAddr().(*net.UDPAddr).Port

				hostConfig := CandidateHostConfig{
					Network:   network,
					Address:   address,
					Port:      port,
					Component: ComponentRTP,
				}

				c, err := NewCandidateHost(&hostConfig)
				if err != nil {
					a.log.Warnf("Failed to create host candidate: %s %s %d: %v\n", network, ip, port, err)
					return
				}

				if a.mDNSMode == MulticastDNSModeQueryAndGather {
					if err = c.setIP(ip); err != nil {
						a.log.Warnf("Failed to create host candidate: %s %s %d: %v\n", network, ip, port, err)
						return
					}
				}

				if err := a.run(func(agent *Agent) {
					c.start(a, conn)
					a.addCandidate(c)

					if a.onCandidateHdlr != nil {
						go a.onCandidateHdlr(c)
					}
				}); err != nil {
					a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
				}
			}(network, ip)
		}
	}
}

func (a *Agent) gatherCandidatesSrflx(urls []*URL, networkTypes []NetworkType) {
	for _, networkType := range networkTypes {
		network := networkType.String()
		for _, url := range urls {
			if url.Scheme != SchemeTypeSTUN {
				continue
			}

			hostPort := fmt.Sprintf("%s:%d", url.Host, url.Port)
			serverAddr, err := a.net.ResolveUDPAddr(network, hostPort)
			if err != nil {
				a.log.Warnf("failed to resolve stun host: %s: %v", hostPort, err)
				continue
			}

			conn, err := a.listenUDP(int(a.portmax), int(a.portmin), network, &net.UDPAddr{IP: nil, Port: 0})
			if err != nil {
				a.log.Warnf("Failed to listen for %s: %v\n", serverAddr.String(), err)
				continue
			}

			xoraddr, err := getXORMappedAddr(conn, serverAddr, stunGatherTimeout)
			if err != nil {
				a.log.Warnf("could not get server reflexive address %s %s: %v\n", network, url, err)
				continue
			}

			laddr := conn.LocalAddr().(*net.UDPAddr)
			ip := xoraddr.IP
			port := xoraddr.Port
			relIP := laddr.IP.String()
			relPort := laddr.Port

			srflxConfig := CandidateServerReflexiveConfig{
				Network:   network,
				Address:   ip.String(),
				Port:      port,
				Component: ComponentRTP,
				RelAddr:   relIP,
				RelPort:   relPort,
			}
			c, err := NewCandidateServerReflexive(&srflxConfig)
			if err != nil {
				a.log.Warnf("Failed to create server reflexive candidate: %s %s %d: %v\n", network, ip, port, err)
				continue
			}

			if err := a.run(func(agent *Agent) {
				c.start(a, conn)
				a.addCandidate(c)

				if a.onCandidateHdlr != nil {
					go a.onCandidateHdlr(c)
				}
			}); err != nil {
				a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
			}
		}
	}

}

func (a *Agent) gatherCandidatesRelay(urls []*URL) error {
	network := NetworkTypeUDP4.String() // TODO IPv6
	for _, url := range urls {
		switch {
		case url.Scheme != SchemeTypeTURN:
			continue
		case url.Username == "":
			return ErrUsernameEmpty
		case url.Password == "":
			return ErrPasswordEmpty
		}

		locConn, err := a.net.ListenPacket(network, "0.0.0.0:0")
		if err != nil {
			return err
		}

		client, err := turn.NewClient(&turn.ClientConfig{
			TURNServerAddr: fmt.Sprintf("%s:%d", url.Host, url.Port),
			Conn:           locConn,
			Username:       url.Username,
			Password:       url.Password,
			LoggerFactory:  a.loggerFactory,
			Net:            a.net,
		})
		if err != nil {
			return err
		}

		err = client.Listen()
		if err != nil {
			return err
		}

		relayConn, err := client.Allocate()
		if err != nil {
			return err
		}

		laddr := locConn.LocalAddr().(*net.UDPAddr)
		raddr := relayConn.LocalAddr().(*net.UDPAddr)

		relayConfig := CandidateRelayConfig{
			Network:   network,
			Component: ComponentRTP,
			Address:   raddr.IP.String(),
			Port:      raddr.Port,
			RelAddr:   laddr.IP.String(),
			RelPort:   laddr.Port,
			OnClose: func() error {
				client.Close()
				return locConn.Close()
			},
		}
		candidate, err := NewCandidateRelay(&relayConfig)
		if err != nil {
			a.log.Warnf("Failed to create relay candidate: %s %s: %v\n",
				network, raddr.String(), err)
			continue
		}

		if err := a.run(func(agent *Agent) {
			candidate.start(a, relayConn)
			a.addCandidate(candidate)

			if a.onCandidateHdlr != nil {
				go a.onCandidateHdlr(candidate)
			}
		}); err != nil {
			a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
		}
	}

	return nil
}

// getXORMappedAddr initiates a stun requests to serverAddr using conn, reads the response and returns
// the XORMappedAddress returned by the stun server.
//
// Adapted from stun v0.2.
func getXORMappedAddr(conn net.PacketConn, serverAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error) {
	if deadline > 0 {
		if err := conn.SetReadDeadline(time.Now().Add(deadline)); err != nil {
			return nil, err
		}
	}
	defer func() {
		if deadline > 0 {
			_ = conn.SetReadDeadline(time.Time{})
		}
	}()
	resp, err := stunRequest(
		func(p []byte) (int, error) {
			n, _, errr := conn.ReadFrom(p)
			return n, errr
		},
		func(b []byte) (int, error) {
			return conn.WriteTo(b, serverAddr)
		},
	)
	if err != nil {
		return nil, err
	}
	var addr stun.XORMappedAddress
	if err = addr.GetFrom(resp); err != nil {
		return nil, fmt.Errorf("failed to get XOR-MAPPED-ADDRESS response: %v", err)
	}
	return &addr, nil
}

func stunRequest(read func([]byte) (int, error), write func([]byte) (int, error)) (*stun.Message, error) {
	req, err := stun.Build(stun.BindingRequest, stun.TransactionID)
	if err != nil {
		return nil, err
	}
	if _, err = write(req.Raw); err != nil {
		return nil, err
	}
	const maxMessageSize = 1280
	bs := make([]byte, maxMessageSize)
	n, err := read(bs)
	if err != nil {
		return nil, err
	}
	res := &stun.Message{Raw: bs[:n]}
	if err := res.Decode(); err != nil {
		return nil, err
	}
	return res, nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package http

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
	"time"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/stats"
	"github.com/apex/log"
)

// Adapter type materializes an http adapter which implements the basic http protocol
type Adapter struct {
	http.Client                    // Adapter is also an http client
	ctx           log.Interface    // Just a logger, no one really cares about him.
	packets       chan PktReq      // Channel used to "transforms" incoming request to something we can handle concurrently
	recipients    []core.Recipient // Known recipient used for broadcast if any
	registrations chan RegReq      // Incoming registrations
	serveMux      *http.ServeMux   // Holds a references to the adapter servemux in order to dynamically define endpoints
}

// Handler defines endpoint-specific handler.
type Handler interface {
	URL() string
	Handle(w http.ResponseWriter, chpkt chan<- PktReq, chreg chan<- RegReq, req *http.Request)
}

// MsgRes are sent through the response channel of a pktReq or regReq
type MsgRes struct {
	StatusCode int    // The http status code to set as an answer
	Content    []byte // The response content.
}

// PktReq are sent through the packets channel when an incoming request arrives
type PktReq struct {
	Packet []byte      // The actual packet that has been parsed
	Chresp chan MsgRes // A response channel waiting for an success or reject confirmation
}

// RegReq are sent through the registration channel when an incoming registration arrives
type RegReq struct {
	Registration core.Registration
	Chresp       chan MsgRes
}

// NewAdapter constructs and allocates a new http adapter
func NewAdapter(net string, recipients []core.Recipient, ctx log.Interface) (*Adapter, error) {
	a := Adapter{
		Client:        http.Client{Timeout: 6 * time.Second},
		ctx:           ctx,
		packets:       make(chan PktReq),
		recipients:    recipients,
		registrations: make(chan RegReq),
		serveMux:      http.NewServeMux(),
	}

	go a.listenRequests(net)

	return &a, nil
}

// Register implements the core.Subscriber interface
func (a *Adapter) Subscribe(r core.Registration) error {
	jsonMarshaler, ok := r.(json.Marshaler)
	if !ok {
		return errors.New(errors.Structural, "Unable to marshal registration")
	}
	httpRecipient, ok := r.Recipient().(Recipient)
	if !ok {
		return errors.New(errors.Structural, "Invalid recipient")
	}

	data, err := jsonMarshaler.MarshalJSON()
	if err != nil {
		return errors.New(errors.Structural, err)
	}
	buf := new(bytes.Buffer)
	buf.Write(data)
	resp, err := a.Post(fmt.Sprintf("http://%s/end-devices", httpRecipient.URL()), "application/json", buf)
	if err != nil {
		return errors.New(errors.Operational, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusOK {
		return errors.New(errors.Operational, "Unable to subscribe")
	}
	return nil
}

// Send implements the core.Adapter interface
func (a *Adapter) Send(p core.Packet, recipients ...core.Recipient) ([]byte, error) {
	stats.MarkMeter("http_adapter.send")
	stats.UpdateHistogram("http_adapter.send_recipients", int64(len(recipients)))

	// Marshal the packet to raw binary data
	data, err := p.MarshalBinary()
	if err != nil {
		a.ctx.WithError(err).Warn("Invalid Packet")
		return nil, errors.New(errors.Structural, err)
	}

	// Try to define a more helpful context
	ctx := a.ctx.WithField("devEUI", p.DevEUI())
	ctx.Debug("Sending Packet")

	// Determine whether it's a broadcast or a direct send
	nb := len(recipients)
	isBroadcast := false
	if nb == 0 {
		// If no recipient was supplied, try with the known one, otherwise quit.
		recipients = a.recipients
		nb = len(recipients)
		isBroadcast = true
		if nb == 0 {
			return nil, errors.New(errors.Structural, "No recipient found")
		}
	}

	// Prepare ground for parrallel http request
	cherr := make(chan error, nb)
	chresp := make(chan []byte, nb)
	wg := sync.WaitGroup{}
	wg.Add(nb)

	// Run each request
	for _, recipient := range recipients {
		go func(rawRecipient core.Recipient) {
			defer wg.Done()

			// Get the actual recipient
			recipient, ok := rawRecipient.(Recipient)
			if !ok {
				ctx.WithField("recipient", rawRecipient).Warn("Unable to interpret recipient as Recipient")
				return
			}
			ctx := ctx.WithField("recipient", recipient.URL())

			// Send request
			ctx.Debugf("%s Request", recipient.Method())
			buf := new(bytes.Buffer)
			buf.Write(data)
			resp, err := a.Post(fmt.Sprintf("http://%s/packets", recipient.URL()), "application/octet-stream", buf)
			if err != nil {
				cherr <- errors.New(errors.Operational, err)
				return
			}
			defer func() {
				// This is needed because the default HTTP client's Transport does not
				// attempt to reuse HTTP/1.0 or HTTP/1.1 TCP connections unless the Body
				// is read to completion and is closed.
				io.Copy(ioutil.Discard, resp.Body)
				resp.Body.Close()
			}()

			// Check response code
			switch resp.StatusCode {
			case http.StatusOK:
				ctx.Debug("Recipient registered for packet")
				data, err := ioutil.ReadAll(resp.Body)
				if err != nil && err != io.EOF {
					cherr <- errors.New(errors.Operational, err)
					return
				}
				chresp <- data
				if isBroadcast { // Generate registration on broadcast
					go func() {
						a.registrations <- RegReq{
							Registration: httpRegistration{
								recipient: rawRecipient,
								devEUI:    p.DevEUI(),
							},
							Chresp: nil,
						}
					}()
				}
			case http.StatusNotFound:
				ctx.Debug("Recipient not interested in packet")
				cherr <- errors.New(errors.Behavioural, "Recipient not interested")
			default:
				cherr <- errors.New(errors.Operational, fmt.Sprintf("Unexpected response from server: %s (%d)", resp.Status, resp.StatusCode))
			}
		}(recipient)
	}

	// Wait for each request to be done
	stats.IncCounter("http_adapter.waiting_for_send")
	wg.Wait()
	stats.DecCounter("http_adapter.waiting_for_send")
	close(cherr)
	close(chresp)

	// Collect errors and see if everything went well
	var errored uint8
	for i := 0; i < len(cherr); i++ {
		err := <-cherr
		if err.(errors.Failure).Nature != errors.Behavioural {
			errored++
			ctx.WithError(err).Warn("POST Failed")
		}
	}

	// Collect response
	if len(chresp) > 1 {
		return nil, errors.New(errors.Behavioural, "Received too many positive answers")
	}

	if len(chresp) == 0 && errored != 0 {
		return nil, errors.New(errors.Operational, "No positive response from recipients but got unexpected answer")
	}

	if len(chresp) == 0 && errored == 0 {
		return nil, errors.New(errors.Behavioural, "No recipient gave a positive answer")
	}

	return <-chresp, nil
}

// GetRecipient implements the core.Adapter interface
func (a *Adapter) GetRecipient(raw []byte) (core.Recipient, error) {
	recipient := new(recipient)
	if err := recipient.UnmarshalBinary(raw); err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	return *recipient, nil
}

// Next implements the core.Adapter interface
func (a *Adapter) Next() ([]byte, core.AckNacker, error) {
	p := <-a.packets
	return p.Packet, httpAckNacker{Chresp: p.Chresp}, nil
}

// NextRegistration implements the core.Adapter interface. Not implemented for this adapter.
//
// See broadcast and pubsub adapters for mechanisms to handle registrations.
func (a *Adapter) NextRegistration() (core.Registration, core.AckNacker, error) {
	r := <-a.registrations
	return r.Registration, regAckNacker{Chresp: r.Chresp}, nil
}

// Bind registers a handler to a specific endpoint
func (a *Adapter) Bind(h Handler) {
	a.ctx.WithField("url", h.URL()).Info("Register new endpoint")
	a.serveMux.HandleFunc(h.URL(), func(w http.ResponseWriter, req *http.Request) {
		a.ctx.WithField("url", h.URL()).Debug("Handle new request")
		h.Handle(w, a.packets, a.registrations, req)
	})
}

// listenRequests handles incoming registration request sent through http to the adapter
func (a *Adapter) listenRequests(net string) {
	server := http.Server{
		Addr:    net,
		Handler: a.serveMux,
	}
	a.ctx.WithField("bind", net).Info("Starting Server")
	err := server.ListenAndServe()
	a.ctx.WithError(err).Warn("HTTP connection lost")
}
<commit_msg>[test/http-adapter] Use of correct http verbs in http adapters<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package http

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
	"time"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/stats"
	"github.com/apex/log"
)

// Adapter type materializes an http adapter which implements the basic http protocol
type Adapter struct {
	http.Client                    // Adapter is also an http client
	ctx           log.Interface    // Just a logger, no one really cares about him.
	packets       chan PktReq      // Channel used to "transforms" incoming request to something we can handle concurrently
	recipients    []core.Recipient // Known recipient used for broadcast if any
	registrations chan RegReq      // Incoming registrations
	serveMux      *http.ServeMux   // Holds a references to the adapter servemux in order to dynamically define endpoints
}

// Handler defines endpoint-specific handler.
type Handler interface {
	URL() string
	Handle(w http.ResponseWriter, chpkt chan<- PktReq, chreg chan<- RegReq, req *http.Request)
}

// MsgRes are sent through the response channel of a pktReq or regReq
type MsgRes struct {
	StatusCode int    // The http status code to set as an answer
	Content    []byte // The response content.
}

// PktReq are sent through the packets channel when an incoming request arrives
type PktReq struct {
	Packet []byte      // The actual packet that has been parsed
	Chresp chan MsgRes // A response channel waiting for an success or reject confirmation
}

// RegReq are sent through the registration channel when an incoming registration arrives
type RegReq struct {
	Registration core.Registration
	Chresp       chan MsgRes
}

// NewAdapter constructs and allocates a new http adapter
func NewAdapter(net string, recipients []core.Recipient, ctx log.Interface) (*Adapter, error) {
	a := Adapter{
		Client:        http.Client{Timeout: 6 * time.Second},
		ctx:           ctx,
		packets:       make(chan PktReq),
		recipients:    recipients,
		registrations: make(chan RegReq),
		serveMux:      http.NewServeMux(),
	}

	go a.listenRequests(net)

	return &a, nil
}

// Register implements the core.Subscriber interface
func (a *Adapter) Subscribe(r core.Registration) error {
	jsonMarshaler, ok := r.(json.Marshaler)
	if !ok {
		return errors.New(errors.Structural, "Unable to marshal registration")
	}
	httpRecipient, ok := r.Recipient().(Recipient)
	if !ok {
		return errors.New(errors.Structural, "Invalid recipient")
	}

	data, err := jsonMarshaler.MarshalJSON()
	if err != nil {
		return errors.New(errors.Structural, err)
	}
	buf := new(bytes.Buffer)
	buf.Write(data)
	req, err := http.NewRequest(httpRecipient.Method(), fmt.Sprintf("http://%s/end-devices", httpRecipient.URL()), buf)
	if err != nil {
		return errors.New(errors.Operational, err)
	}
	req.Header.Add("content-type", "application/json")
	resp, err := a.Do(req)
	if err != nil {
		return errors.New(errors.Operational, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusOK {
		return errors.New(errors.Operational, "Unable to subscribe")
	}
	return nil
}

// Send implements the core.Adapter interface
func (a *Adapter) Send(p core.Packet, recipients ...core.Recipient) ([]byte, error) {
	stats.MarkMeter("http_adapter.send")
	stats.UpdateHistogram("http_adapter.send_recipients", int64(len(recipients)))

	// Marshal the packet to raw binary data
	data, err := p.MarshalBinary()
	if err != nil {
		a.ctx.WithError(err).Warn("Invalid Packet")
		return nil, errors.New(errors.Structural, err)
	}

	// Try to define a more helpful context
	ctx := a.ctx.WithField("devEUI", p.DevEUI())
	ctx.Debug("Sending Packet")

	// Determine whether it's a broadcast or a direct send
	nb := len(recipients)
	isBroadcast := false
	if nb == 0 {
		// If no recipient was supplied, try with the known one, otherwise quit.
		recipients = a.recipients
		nb = len(recipients)
		isBroadcast = true
		if nb == 0 {
			return nil, errors.New(errors.Structural, "No recipient found")
		}
	}

	// Prepare ground for parrallel http request
	cherr := make(chan error, nb)
	chresp := make(chan []byte, nb)
	wg := sync.WaitGroup{}
	wg.Add(nb)

	// Run each request
	for _, recipient := range recipients {
		go func(rawRecipient core.Recipient) {
			defer wg.Done()

			// Get the actual recipient
			recipient, ok := rawRecipient.(Recipient)
			if !ok {
				ctx.WithField("recipient", rawRecipient).Warn("Unable to interpret recipient as Recipient")
				return
			}
			ctx := ctx.WithField("recipient", recipient.URL())

			// Send request
			ctx.Debugf("%s Request", recipient.Method())
			buf := new(bytes.Buffer)
			buf.Write(data)
			req, err := http.NewRequest(recipient.Method(), fmt.Sprintf("http://%s/packets", recipient.URL()), buf)
			if err != nil {
				cherr <- errors.New(errors.Operational, err)
				return
			}
			req.Header.Add("content-type", "application/octet-stream")
			resp, err := a.Do(req)

			if err != nil {
				cherr <- errors.New(errors.Operational, err)
				return
			}
			defer func() {
				// This is needed because the default HTTP client's Transport does not
				// attempt to reuse HTTP/1.0 or HTTP/1.1 TCP connections unless the Body
				// is read to completion and is closed.
				io.Copy(ioutil.Discard, resp.Body)
				resp.Body.Close()
			}()

			// Check response code
			switch resp.StatusCode {
			case http.StatusOK:
				ctx.Debug("Recipient registered for packet")
				data, err := ioutil.ReadAll(resp.Body)
				if err != nil && err != io.EOF {
					cherr <- errors.New(errors.Operational, err)
					return
				}
				chresp <- data
				if isBroadcast { // Generate registration on broadcast
					go func() {
						a.registrations <- RegReq{
							Registration: httpRegistration{
								recipient: rawRecipient,
								devEUI:    p.DevEUI(),
							},
							Chresp: nil,
						}
					}()
				}
			case http.StatusNotFound:
				ctx.Debug("Recipient not interested in packet")
				cherr <- errors.New(errors.Behavioural, "Recipient not interested")
			default:
				cherr <- errors.New(errors.Operational, fmt.Sprintf("Unexpected response from server: %s (%d)", resp.Status, resp.StatusCode))
			}
		}(recipient)
	}

	// Wait for each request to be done
	stats.IncCounter("http_adapter.waiting_for_send")
	wg.Wait()
	stats.DecCounter("http_adapter.waiting_for_send")
	close(cherr)
	close(chresp)

	// Collect errors and see if everything went well
	var errored uint8
	for i := 0; i < len(cherr); i++ {
		err := <-cherr
		if err.(errors.Failure).Nature != errors.Behavioural {
			errored++
			ctx.WithError(err).Warn("POST Failed")
		}
	}

	// Collect response
	if len(chresp) > 1 {
		return nil, errors.New(errors.Behavioural, "Received too many positive answers")
	}

	if len(chresp) == 0 && errored != 0 {
		return nil, errors.New(errors.Operational, "No positive response from recipients but got unexpected answer")
	}

	if len(chresp) == 0 && errored == 0 {
		return nil, errors.New(errors.Behavioural, "No recipient gave a positive answer")
	}

	return <-chresp, nil
}

// GetRecipient implements the core.Adapter interface
func (a *Adapter) GetRecipient(raw []byte) (core.Recipient, error) {
	recipient := new(recipient)
	if err := recipient.UnmarshalBinary(raw); err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	return *recipient, nil
}

// Next implements the core.Adapter interface
func (a *Adapter) Next() ([]byte, core.AckNacker, error) {
	p := <-a.packets
	return p.Packet, httpAckNacker{Chresp: p.Chresp}, nil
}

// NextRegistration implements the core.Adapter interface. Not implemented for this adapter.
//
// See broadcast and pubsub adapters for mechanisms to handle registrations.
func (a *Adapter) NextRegistration() (core.Registration, core.AckNacker, error) {
	r := <-a.registrations
	return r.Registration, regAckNacker{Chresp: r.Chresp}, nil
}

// Bind registers a handler to a specific endpoint
func (a *Adapter) Bind(h Handler) {
	a.ctx.WithField("url", h.URL()).Info("Register new endpoint")
	a.serveMux.HandleFunc(h.URL(), func(w http.ResponseWriter, req *http.Request) {
		a.ctx.WithField("url", h.URL()).Debug("Handle new request")
		h.Handle(w, a.packets, a.registrations, req)
	})
}

// listenRequests handles incoming registration request sent through http to the adapter
func (a *Adapter) listenRequests(net string) {
	server := http.Server{
		Addr:    net,
		Handler: a.serveMux,
	}
	a.ctx.WithField("bind", net).Info("Starting Server")
	err := server.ListenAndServe()
	a.ctx.WithError(err).Warn("HTTP connection lost")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"compress/gzip"
	"flag"
	"fmt"
	"github.com/oschwald/maxminddb-golang"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
)

var reader *maxminddb.Reader
var verbose bool
var format string

type geoIPResult struct {
	Location struct {
		Longitude float64 `maxminddb:"longitude"`
		Latitude  float64 `maxminddb:"latitude"`
	} `maxminddb:"location"`
	City struct {
		Names map[string]string `maxminddb:"names"`
	} `maxminddb:"city"`
	Country struct {
		IsoCode string            `maxminddb:"iso_code"`
		Names   map[string]string `maxminddb:"names"`
	} `maxminddb:"country"`
}

func init() {
	// download database
	path := "/tmp/GeoLite2-City.mmdb"

	_, err := os.Stat(path)
	if os.IsNotExist(err) {
		fmt.Println("Downloading database")

		client := &http.Client{}

		req, err := http.NewRequest("GET", geoLiteURL, nil)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		var resp *http.Response
		if resp, err = client.Do(req); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		defer resp.Body.Close()

		gzf, err := gzip.NewReader(resp.Body)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		defer gzf.Close()

		f, err := os.Create(path)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		defer f.Close()

		_, err = io.Copy(f, gzf)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}

	if verbose {
		fmt.Printf("Using database %s.\n", path)
	}

	reader, err = maxminddb.Open(path)
}

func help() {
	fmt.Println("No ip addresses")
}

func update() error {
	return nil
}

const geoLiteURL = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz"

func main() {
	// initial download
	// update
	// open from cache
	// resolve

	flag.StringVar(&format, "format", "(country) ((city))", "format")
	flag.BoolVar(&verbose, "verbose", false, "verbose")
	flag.Parse()

	fi, err := os.Stdin.Stat()
	if err != nil {
		panic(err)
	}

	var args = flag.Args()

	if fi.Mode()&os.ModeNamedPipe > 0 {
		scanner := bufio.NewScanner(os.Stdin)

		for scanner.Scan() {
			args = append(args, scanner.Text())
		}

		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}

	if len(args) == 0 {
		help()
		os.Exit(1)
	}

	for _, arg := range args {
		var addr net.IP
		addr = net.ParseIP(arg)
		if addr == nil {
			fmt.Printf("%s is not a valid ip address", addr)
			continue
		}

		var result geoIPResult

		var err error
		if err = reader.Lookup(addr, &result); err != nil {
			panic(err)
		}

		var p string
		p = format
		p = strings.Replace(p, "(ip)", addr.String(), -1)
		p = strings.Replace(p, "(country)", result.Country.Names["en"], -1)
		p = strings.Replace(p, "(city)", result.City.Names["en"], -1)
		p = strings.Replace(p, "(lat)", fmt.Sprintf("%f", result.Location.Latitude), -1)
		p = strings.Replace(p, "(long)", fmt.Sprintf("%f", result.Location.Longitude), -1)

		fmt.Print(p)
	}
}
<commit_msg>added newline support in format<commit_after>package main

import (
	"bufio"
	"compress/gzip"
	"flag"
	"fmt"
	"github.com/oschwald/maxminddb-golang"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
)

var reader *maxminddb.Reader
var verbose bool
var format string

type geoIPResult struct {
	Location struct {
		Longitude float64 `maxminddb:"longitude"`
		Latitude  float64 `maxminddb:"latitude"`
	} `maxminddb:"location"`
	City struct {
		Names map[string]string `maxminddb:"names"`
	} `maxminddb:"city"`
	Country struct {
		IsoCode string            `maxminddb:"iso_code"`
		Names   map[string]string `maxminddb:"names"`
	} `maxminddb:"country"`
}

func init() {
	// download database
	path := "/tmp/GeoLite2-City.mmdb"

	_, err := os.Stat(path)
	if os.IsNotExist(err) {
		fmt.Println("Downloading database")

		client := &http.Client{}

		req, err := http.NewRequest("GET", geoLiteURL, nil)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		var resp *http.Response
		if resp, err = client.Do(req); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		defer resp.Body.Close()

		gzf, err := gzip.NewReader(resp.Body)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		defer gzf.Close()

		f, err := os.Create(path)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		defer f.Close()

		_, err = io.Copy(f, gzf)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}

	if verbose {
		fmt.Printf("Using database %s.\n", path)
	}

	reader, err = maxminddb.Open(path)
}

func help() {
	fmt.Println("No ip addresses")
}

func update() error {
	return nil
}

const geoLiteURL = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz"

func main() {
	// initial download
	// update
	// open from cache
	// resolve

	flag.StringVar(&format, "format", "(country) ((city))", "format")
	flag.BoolVar(&verbose, "verbose", false, "verbose")
	flag.Parse()

	format = strings.Replace(format, "\\n", "\n", -1)

	fi, err := os.Stdin.Stat()
	if err != nil {
		panic(err)
	}

	var args = flag.Args()

	if fi.Mode()&os.ModeNamedPipe > 0 {
		scanner := bufio.NewScanner(os.Stdin)

		for scanner.Scan() {
			args = append(args, scanner.Text())
		}

		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}

	if len(args) == 0 {
		help()
		os.Exit(1)
	}

	for _, arg := range args {
		var addr net.IP
		addr = net.ParseIP(arg)
		if addr == nil {
			fmt.Printf("%s is not a valid ip address", addr)
			continue
		}

		var result geoIPResult

		var err error
		if err = reader.Lookup(addr, &result); err != nil {
			panic(err)
		}

		var p string
		p = format
		p = strings.Replace(p, "(ip)", addr.String(), -1)
		p = strings.Replace(p, "(country)", result.Country.Names["en"], -1)
		p = strings.Replace(p, "(city)", result.City.Names["en"], -1)
		p = strings.Replace(p, "(lat)", fmt.Sprintf("%f", result.Location.Latitude), -1)
		p = strings.Replace(p, "(long)", fmt.Sprintf("%f", result.Location.Longitude), -1)

		fmt.Print(p)
	}
}
<|endoftext|>


<commit_before>package meetup

import (
	"errors"
	"reflect"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

var (
	fakeTimeMu sync.Mutex
	fakeTime   = time.Now()
)

func init() {
	now = func() time.Time {
		fakeTimeMu.Lock()
		t := fakeTime
		fakeTimeMu.Unlock()
		return t
	}
}

func advanceTime(d time.Duration) {
	fakeTimeMu.Lock()
	fakeTime = fakeTime.Add(d)
	fakeTimeMu.Unlock()
}

func mustGet(t *testing.T, c *Cache, key string, want interface{}) {
	v, err := c.Get(key)
	if err != nil {
		t.Errorf("Get(%#v) returned unexpected error %v", key, err)
	}
	if !reflect.DeepEqual(v, want) {
		t.Errorf("Get(%#v) = %#v, but wanted %#v", key, v, want)
	}
}

func TestCache(t *testing.T) {
	hits := 0
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			hits++
			return key, nil
		},
	})
	defer c.Close()

	if hits != 0 {
		t.Fatalf("hits != 0 after init")
	}

	mustGet(t, c, "a", "a")
	if hits != 1 {
		t.Fatalf("hits != 1 after first use")
	}

	mustGet(t, c, "b", "b")
	if hits != 2 {
		t.Fatalf("hits = %v after second use", hits)
	}

	mustGet(t, c, "a", "a")
	if hits != 2 {
		t.Fatalf("hits != 2 after third use")
	}
}

func TestExpiry(t *testing.T) {
	var mu sync.Mutex
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			hits++
			mu.Unlock()
			return key, nil
		},
		ExpireAge: time.Second,
	})
	defer c.Close()

	mustGet(t, c, "a", "a")
	mu.Lock()
	if hits != 1 {
		t.Fatalf("hits = %v after first use", hits)
	}
	mu.Unlock()

	advanceTime(2 * time.Second)

	mustGet(t, c, "a", "a")
	mu.Lock()
	if hits != 2 {
		t.Fatalf("hits != 2 after second use")
	}
	mu.Unlock()

	advanceTime(time.Second / 2)

	mustGet(t, c, "a", "a")
	mu.Lock()
	if hits != 2 {
		t.Fatalf("hits = %v after third use", hits)
	}
	mu.Unlock()
}

func TestExpiryUsesStartTime(t *testing.T) {
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			advanceTime(time.Second)
			hits++
			return key, nil
		},
		ExpireAge: time.Second,
	})
	defer c.Close()

	mustGet(t, c, "a", "a")
	if hits != 1 {
		t.Fatalf("hits != 1 after first use")
	}

	mustGet(t, c, "a", "a")
	if hits != 2 {
		t.Fatalf("hits = %v after second use", hits)
	}
}

func TestMeetup(t *testing.T) {
	postGetCheckCh = make(chan struct{})
	defer func() { postGetCheckCh = nil }()

	const concurrency = 1000

	blockGets := newBoolWatcher(true)

	var mu sync.Mutex
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			blockGets.Wait(false)

			mu.Lock()
			c := hits
			hits++
			mu.Unlock()
			return c, nil
		},
	})
	defer c.Close()

	vals := make(chan interface{})
	for i := 0; i < concurrency; i++ {
		go func() {
			v, err := c.Get("a")
			if err != nil {
				t.Errorf(`c.Get("a") returned error %v`, err)
			}
			vals <- v
		}()
	}

	for i := 0; i < concurrency; i++ {
		<-postGetCheckCh
	}

	blockGets.Set(false)

	for i := 0; i < concurrency; i++ {
		v := <-vals
		if !reflect.DeepEqual(v, 0) {
			t.Errorf("Cache did not meet up. Wanted value %v, got %v", 0, v)
		}
	}
}

func TestConcurrencyLimit(t *testing.T) {
	blockGets := newBoolWatcher(true)

	const (
		workers = 1000
		limit   = 3
	)

	var mu sync.Mutex
	conc := 0
	maxConc := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			conc++
			if conc > maxConc {
				maxConc = conc
			}
			mu.Unlock()

			blockGets.Wait(false)

			mu.Lock()
			conc--
			mu.Unlock()

			return key, nil
		},
		Concurrency: limit,
	})
	defer c.Close()

	vals := make(chan interface{})

	for i := 0; i < workers; i++ {
		go func(i int) {
			k := strconv.FormatInt(int64(i), 10)
			v, _ := c.Get(k)
			vals <- v
		}(i)
	}

	blockGets.Set(false)

	for i := 0; i < workers; i++ {
		<-vals
	}

	if conc != 0 {
		t.Errorf("Options.Get still running after Cache.Get returned")
	}

	t.Logf("max concurrency seen was %v", maxConc)
	if maxConc > limit {
		t.Errorf("max concurrency of %v is over limit %v", maxConc, limit)
	}
}

func TestCacheDoesntKeepErrors(t *testing.T) {
	deliberateErr := errors.New("deliberate failure")

	hits := 0
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			hits++
			return nil, deliberateErr
		},
	})
	defer c.Close()

	v, err := c.Get("a")
	if v != nil {
		t.Errorf("Got unexpected value %#v from c.Get", v)
	}
	if err != deliberateErr {
		t.Errorf("Got unexpected error %v from c.Get", err)
	}

	v, err = c.Get("a")
	if v != nil {
		t.Errorf("Got unexpected value %#v from c.Get", v)
	}
	if err != deliberateErr {
		t.Errorf("Got unexpected error %v from c.Get", err)
	}

	if hits != 2 {
		t.Errorf("Hits was %v, not 2", hits)
	}
}

func TestRevalidation(t *testing.T) {
	fillComplete = make(chan struct{})
	defer func() { fillComplete = nil }()

	blockGets := newBoolWatcher(false)
	var getReadyToBlock chan struct{}

	var mu sync.Mutex
	hits := 0
	finished := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			hits++
			myHits := hits
			mu.Unlock()
			if getReadyToBlock != nil {
				getReadyToBlock <- struct{}{}
			}
			blockGets.Wait(false)
			mu.Lock()
			finished++
			mu.Unlock()
			return myHits, nil
		},
		RevalidateAge: time.Second,
	})
	defer c.Close()

	mustGet(t, c, "a", 1)
	<-fillComplete
	if hits != 1 {
		t.Errorf("hits != 1")
	}

	blockGets.Set(true)
	advanceTime(time.Second)

	getReadyToBlock = make(chan struct{}, 1)
	mustGet(t, c, "a", 1) // NB: does not block with background revalidation
	<-getReadyToBlock
	mu.Lock()
	if hits != 2 {
		t.Errorf("hits = %v, wanted 2", hits)
	}
	if finished != 1 {
		t.Errorf("finished != 1")
	}
	mu.Unlock()

	blockGets.Set(false)
	<-fillComplete

	getReadyToBlock = nil
	mustGet(t, c, "a", 2)
}

func TestErrorCaching(t *testing.T) {
	deliberate := errors.New("deliberate failure")

	var mu sync.Mutex
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			hits++
			mu.Unlock()
			return nil, deliberate
		},
		ErrorAge: time.Second,
	})
	defer c.Close()

	_, err := c.Get("a")
	if err != deliberate {
		t.Errorf("Got unexpected error %v from Get", err)
	}

	_, err = c.Get("a")
	if err != deliberate {
		t.Errorf("Got unexpected error %v from Get", err)
	}

	if hits != 1 {
		t.Errorf("error was not cached")
	}

	advanceTime(time.Second)

	deliberate = errors.New("other deliberate failure")
	_, err = c.Get("a")
	if err != deliberate {
		t.Errorf("Got unexpected error %v from Get", err)
	}

	if hits != 2 {
		t.Errorf("hits = %v", hits)
	}
}

func TestClosedGet(t *testing.T) {
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			return key, nil
		},
	})
	c.Close()

	_, err := c.Get("a")
	if err != ErrClosed {
		t.Errorf("Got unexpected error %v from Get", err)
	}
}

func TestGetMeetupCreateRace(t *testing.T) {
	// This test tries to maximize the likelihood of the read lock to write lock
	// transition in Cache.Get noticing that the entry has already been created
	// even though it already transitioned to a write lock.
	//
	// By having multiple workers requesting the same key, then changing that
	// key relatively slowly over time, we should see exactly as many hits as
	// keys, and we should never deadlock.

	var hits uint64
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			atomic.AddUint64(&hits, 1)
			return key, nil
		},
	})
	defer c.Close()

	var keyInt uint64 = 1

	done := make(chan struct{})
	for worker := 0; worker < 5; worker++ {
		go func() {
			for i := 0; i < 1000; i++ {
				for i := 0; i < 10; i++ {
					c.Get(strconv.FormatUint(atomic.LoadUint64(&keyInt), 10))
				}

				newKeyInt := atomic.AddUint64(&keyInt, 1)
				c.Get(strconv.FormatUint(newKeyInt, 10))
			}
			done <- struct{}{}
		}()
	}

	for worker := 0; worker < 5; worker++ {
		<-done
	}

	gotHits := atomic.LoadUint64(&hits)
	gotKeys := atomic.LoadUint64(&keyInt)

	if gotHits != gotKeys {
		t.Errorf("made %v keys, but got %v hits", gotKeys, gotHits)
	}

	t.Logf("key at end was %s", gotKeys)
}
<commit_msg>make TestGetMeetupCreateRace more customizable and slightly stricter<commit_after>package meetup

import (
	"errors"
	"reflect"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

var (
	fakeTimeMu sync.Mutex
	fakeTime   = time.Now()
)

func init() {
	now = func() time.Time {
		fakeTimeMu.Lock()
		t := fakeTime
		fakeTimeMu.Unlock()
		return t
	}
}

func advanceTime(d time.Duration) {
	fakeTimeMu.Lock()
	fakeTime = fakeTime.Add(d)
	fakeTimeMu.Unlock()
}

func mustGet(t *testing.T, c *Cache, key string, want interface{}) {
	v, err := c.Get(key)
	if err != nil {
		t.Errorf("Get(%#v) returned unexpected error %v", key, err)
	}
	if !reflect.DeepEqual(v, want) {
		t.Errorf("Get(%#v) = %#v, but wanted %#v", key, v, want)
	}
}

func TestCache(t *testing.T) {
	hits := 0
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			hits++
			return key, nil
		},
	})
	defer c.Close()

	if hits != 0 {
		t.Fatalf("hits != 0 after init")
	}

	mustGet(t, c, "a", "a")
	if hits != 1 {
		t.Fatalf("hits != 1 after first use")
	}

	mustGet(t, c, "b", "b")
	if hits != 2 {
		t.Fatalf("hits = %v after second use", hits)
	}

	mustGet(t, c, "a", "a")
	if hits != 2 {
		t.Fatalf("hits != 2 after third use")
	}
}

func TestExpiry(t *testing.T) {
	var mu sync.Mutex
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			hits++
			mu.Unlock()
			return key, nil
		},
		ExpireAge: time.Second,
	})
	defer c.Close()

	mustGet(t, c, "a", "a")
	mu.Lock()
	if hits != 1 {
		t.Fatalf("hits = %v after first use", hits)
	}
	mu.Unlock()

	advanceTime(2 * time.Second)

	mustGet(t, c, "a", "a")
	mu.Lock()
	if hits != 2 {
		t.Fatalf("hits != 2 after second use")
	}
	mu.Unlock()

	advanceTime(time.Second / 2)

	mustGet(t, c, "a", "a")
	mu.Lock()
	if hits != 2 {
		t.Fatalf("hits = %v after third use", hits)
	}
	mu.Unlock()
}

func TestExpiryUsesStartTime(t *testing.T) {
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			advanceTime(time.Second)
			hits++
			return key, nil
		},
		ExpireAge: time.Second,
	})
	defer c.Close()

	mustGet(t, c, "a", "a")
	if hits != 1 {
		t.Fatalf("hits != 1 after first use")
	}

	mustGet(t, c, "a", "a")
	if hits != 2 {
		t.Fatalf("hits = %v after second use", hits)
	}
}

func TestMeetup(t *testing.T) {
	postGetCheckCh = make(chan struct{})
	defer func() { postGetCheckCh = nil }()

	const concurrency = 1000

	blockGets := newBoolWatcher(true)

	var mu sync.Mutex
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			blockGets.Wait(false)

			mu.Lock()
			c := hits
			hits++
			mu.Unlock()
			return c, nil
		},
	})
	defer c.Close()

	vals := make(chan interface{})
	for i := 0; i < concurrency; i++ {
		go func() {
			v, err := c.Get("a")
			if err != nil {
				t.Errorf(`c.Get("a") returned error %v`, err)
			}
			vals <- v
		}()
	}

	for i := 0; i < concurrency; i++ {
		<-postGetCheckCh
	}

	blockGets.Set(false)

	for i := 0; i < concurrency; i++ {
		v := <-vals
		if !reflect.DeepEqual(v, 0) {
			t.Errorf("Cache did not meet up. Wanted value %v, got %v", 0, v)
		}
	}
}

func TestConcurrencyLimit(t *testing.T) {
	blockGets := newBoolWatcher(true)

	const (
		workers = 1000
		limit   = 3
	)

	var mu sync.Mutex
	conc := 0
	maxConc := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			conc++
			if conc > maxConc {
				maxConc = conc
			}
			mu.Unlock()

			blockGets.Wait(false)

			mu.Lock()
			conc--
			mu.Unlock()

			return key, nil
		},
		Concurrency: limit,
	})
	defer c.Close()

	vals := make(chan interface{})

	for i := 0; i < workers; i++ {
		go func(i int) {
			k := strconv.FormatInt(int64(i), 10)
			v, _ := c.Get(k)
			vals <- v
		}(i)
	}

	blockGets.Set(false)

	for i := 0; i < workers; i++ {
		<-vals
	}

	if conc != 0 {
		t.Errorf("Options.Get still running after Cache.Get returned")
	}

	t.Logf("max concurrency seen was %v", maxConc)
	if maxConc > limit {
		t.Errorf("max concurrency of %v is over limit %v", maxConc, limit)
	}
}

func TestCacheDoesntKeepErrors(t *testing.T) {
	deliberateErr := errors.New("deliberate failure")

	hits := 0
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			hits++
			return nil, deliberateErr
		},
	})
	defer c.Close()

	v, err := c.Get("a")
	if v != nil {
		t.Errorf("Got unexpected value %#v from c.Get", v)
	}
	if err != deliberateErr {
		t.Errorf("Got unexpected error %v from c.Get", err)
	}

	v, err = c.Get("a")
	if v != nil {
		t.Errorf("Got unexpected value %#v from c.Get", v)
	}
	if err != deliberateErr {
		t.Errorf("Got unexpected error %v from c.Get", err)
	}

	if hits != 2 {
		t.Errorf("Hits was %v, not 2", hits)
	}
}

func TestRevalidation(t *testing.T) {
	fillComplete = make(chan struct{})
	defer func() { fillComplete = nil }()

	blockGets := newBoolWatcher(false)
	var getReadyToBlock chan struct{}

	var mu sync.Mutex
	hits := 0
	finished := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			hits++
			myHits := hits
			mu.Unlock()
			if getReadyToBlock != nil {
				getReadyToBlock <- struct{}{}
			}
			blockGets.Wait(false)
			mu.Lock()
			finished++
			mu.Unlock()
			return myHits, nil
		},
		RevalidateAge: time.Second,
	})
	defer c.Close()

	mustGet(t, c, "a", 1)
	<-fillComplete
	if hits != 1 {
		t.Errorf("hits != 1")
	}

	blockGets.Set(true)
	advanceTime(time.Second)

	getReadyToBlock = make(chan struct{}, 1)
	mustGet(t, c, "a", 1) // NB: does not block with background revalidation
	<-getReadyToBlock
	mu.Lock()
	if hits != 2 {
		t.Errorf("hits = %v, wanted 2", hits)
	}
	if finished != 1 {
		t.Errorf("finished != 1")
	}
	mu.Unlock()

	blockGets.Set(false)
	<-fillComplete

	getReadyToBlock = nil
	mustGet(t, c, "a", 2)
}

func TestErrorCaching(t *testing.T) {
	deliberate := errors.New("deliberate failure")

	var mu sync.Mutex
	hits := 0

	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			mu.Lock()
			hits++
			mu.Unlock()
			return nil, deliberate
		},
		ErrorAge: time.Second,
	})
	defer c.Close()

	_, err := c.Get("a")
	if err != deliberate {
		t.Errorf("Got unexpected error %v from Get", err)
	}

	_, err = c.Get("a")
	if err != deliberate {
		t.Errorf("Got unexpected error %v from Get", err)
	}

	if hits != 1 {
		t.Errorf("error was not cached")
	}

	advanceTime(time.Second)

	deliberate = errors.New("other deliberate failure")
	_, err = c.Get("a")
	if err != deliberate {
		t.Errorf("Got unexpected error %v from Get", err)
	}

	if hits != 2 {
		t.Errorf("hits = %v", hits)
	}
}

func TestClosedGet(t *testing.T) {
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			return key, nil
		},
	})
	c.Close()

	_, err := c.Get("a")
	if err != ErrClosed {
		t.Errorf("Got unexpected error %v from Get", err)
	}
}

func TestGetMeetupCreateRace(t *testing.T) {
	// This test tries to maximize the likelihood of the read lock to write lock
	// transition in Cache.Get noticing that the entry has already been created
	// even though it already transitioned to a write lock.
	//
	// By having multiple workers requesting the same key, then changing that
	// key relatively slowly over time, we should see exactly as many hits as
	// keys, and we should never deadlock.

	const (
		workers         = 5
		iterations      = 1000
		extraGetsPerKey = 2 * workers
	)

	var hits uint64
	c := NewCache(Options{
		Get: func(key string) (interface{}, error) {
			atomic.AddUint64(&hits, 1)
			return key, nil
		},
	})
	defer c.Close()

	var keyInt uint64 = 1

	done := make(chan struct{})
	for worker := 0; worker < workers; worker++ {
		go func() {
			for i := 0; i < iterations; i++ {
				for i := 0; i < extraGetsPerKey; i++ {
					c.Get(strconv.FormatUint(atomic.LoadUint64(&keyInt), 10))
				}

				newKeyInt := atomic.AddUint64(&keyInt, 1)
				c.Get(strconv.FormatUint(newKeyInt, 10))
			}
			done <- struct{}{}
		}()
	}

	for worker := 0; worker < workers; worker++ {
		<-done
	}

	gotHits := atomic.LoadUint64(&hits)
	gotKeys := atomic.LoadUint64(&keyInt)

	if gotHits != gotKeys {
		t.Errorf("made %v keys, but got %v hits", gotKeys, gotHits)
	}

	if gotKeys != iterations*workers+1 {
		t.Errorf("created %v keys but wanted %v", gotKeys, iterations*workers+1)
	}

	t.Logf("key at end was %v", gotKeys)
}
<|endoftext|>