<commit_before>package client

import (
	"net/http"
	"log"
	"github.com/SchweizerischeBundesbahnen/openshift-monitoring/models"
	"time"
	"strings"
	"os/exec"
	"bytes"
	"net"
	"crypto/tls"
)

const (
	deamonDNSEndpoint = "deamon.ose-mon-a.endpoints.cluster.local"
	deamonDNSServiceA = "deamon.ose-mon-a.svc.cluster.local"
	deamonDNSServiceB = "deamon.ose-mon-b.svc.cluster.local"
	deamonDNSServiceC = "deamon.ose-mon-c.svc.cluster.local"
	deamonDNSPod = "deamon"
	kubernetesIP = "172.30.0.1"
)

func startChecks(dc *models.DeamonClient, checks *models.Checks) {
	tickExt := time.Tick(time.Duration(checks.CheckInterval) * time.Millisecond)
	tickInt := time.Tick(5 * time.Second)

	log.Println("starting checks")

	go func() {
		for {
			select {
			case <-dc.Quit:
				log.Println("stopped checks")
				return
			case <-tickInt:
				if (checks.MasterApiCheck) {
					go checkMasterApis(dc, checks.MasterApiUrls)
				}
			case <-tickExt:
				if (checks.DnsCheck) {
					go checkDnsNslookupOnKubernetes(dc)

					if (dc.Deamon.IsNode()) {
						go checkDnsServiceNode(dc)
					}

					if (dc.Deamon.IsPod()) {
						go checkDnsInPod(dc)
					}
				}

				if (checks.HttpChecks) {
					if (dc.Deamon.IsNode() || (dc.Deamon.IsPod() && strings.HasSuffix(dc.Deamon.Namespace, "a"))) {
						go checkPodHttpAtoB(dc)
						go checkPodHttpAtoC(dc)
					}

					go checkHttpHaProxy(dc, checks.DeamonPublicUrl)
				}
			}
		}
	}()
}

func stopChecks(dc *models.DeamonClient) {
	dc.Quit <- true
}

func checkDnsNslookupOnKubernetes(dc *models.DeamonClient) {
	handleCheckStarted(dc)
	isOk := false
	var msg string

	cmd := exec.Command("nslookup", deamonDNSEndpoint, kubernetesIP)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		isOk = false
		log.Println("error with nslookup: ", err)
		msg = "DNS resolution via nslookup & kubernetes failed."
	}

	stdOut := out.String()

	if (strings.Contains(stdOut, "Server") && strings.Count(stdOut, "Address") >= 2 && strings.Contains(stdOut, "Name")) {
		isOk = true
	} else {
		msg += "NsLookup had wrong output"
	}

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.DNS_NSLOOKUP_KUBERNETES, IsOk: isOk, Message: msg}
}

func checkDnsServiceNode(dc *models.DeamonClient) {
	handleCheckStarted(dc)
	isOk := false
	var msg string

	ips := getIpsForName(deamonDNSServiceA)

	if (ips == nil) {
		isOk = false
		msg = "Failed to lookup ip on node (dnsmasq) for name " + deamonDNSServiceA
	}

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.DNS_SERVICE_NODE, IsOk: isOk, Message: msg}
}

func checkDnsInPod(dc *models.DeamonClient) {
	handleCheckStarted(dc)
	isOk := false
	var msg string

	ips := getIpsForName(deamonDNSPod)

	if (ips == nil) {
		isOk = false
		msg = "Failed to lookup ip in pod for name " + deamonDNSPod
	} else {
		isOk = true
	}

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.DNS_SERVICE_POD, IsOk: isOk, Message: msg}
}

func getIpsForName(n string) []net.IP {
	ips, err := net.LookupIP(n)
	if (err != nil) {
		log.Println("failed to lookup ip for name ", n)
		return nil
	}
	return ips
}

func checkMasterApis(dc *models.DeamonClient, urls string) {
	handleCheckStarted(dc)
	urlArr := strings.Split(urls, ",")

	oneApiOk := false
	var msg string
	for _, u := range urlArr {
		if (checkHttp(u)) {
			oneApiOk = true
		} else {
			msg += u + " is not reachable. ";
		}
	}

	handleCheckFinished(dc, oneApiOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.MASTER_API_CHECK, IsOk: oneApiOk, Message: msg}
}

func checkHttp(toCall string) bool {
	if (strings.HasPrefix(toCall, "https")) {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
		client := &http.Client{Transport: tr}
		_, err := client.Get(toCall)
		if (err != nil) {
			log.Println("error in http check: ", err)
		}
		return err == nil
	} else {
		_, err := http.Get(toCall)
		if (err != nil) {
			log.Println("error in http check: ", err)
		}
		return err == nil
	}
}

func checkPodHttpAtoB(dc *models.DeamonClient) {
	// This should fail as we do not have access to this project
	handleCheckStarted(dc)
	var msg string

	isOk := !checkHttp(deamonDNSServiceB + ":8090/hello")

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.HTTP_POD_SERVICE_A_B, IsOk: isOk, Message: msg}
}

func checkPodHttpAtoC(dc *models.DeamonClient) {
	// This should work as we joined this projects
	handleCheckStarted(dc)
	var msg string

	isOk := checkHttp(deamonDNSServiceC + ":8090/hello")

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.HTTP_POD_SERVICE_A_C, IsOk: isOk, Message: msg}
}

func checkHttpHaProxy(dc *models.DeamonClient, publicUrl string) {
	handleCheckStarted(dc)
	var msg string

	isOk := checkHttp(publicUrl + ":80/hello")

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.HTTP_HAPROXY, IsOk: isOk, Message: msg}
}
<commit_msg>http / https handling<commit_after>package client

import (
	"net/http"
	"log"
	"github.com/SchweizerischeBundesbahnen/openshift-monitoring/models"
	"time"
	"strings"
	"os/exec"
	"bytes"
	"net"
	"crypto/tls"
)

const (
	deamonDNSEndpoint = "deamon.ose-mon-a.endpoints.cluster.local"
	deamonDNSServiceA = "deamon.ose-mon-a.svc.cluster.local"
	deamonDNSServiceB = "deamon.ose-mon-b.svc.cluster.local"
	deamonDNSServiceC = "deamon.ose-mon-c.svc.cluster.local"
	deamonDNSPod = "deamon"
	kubernetesIP = "172.30.0.1"
)

func startChecks(dc *models.DeamonClient, checks *models.Checks) {
	tickExt := time.Tick(time.Duration(checks.CheckInterval) * time.Millisecond)
	tickInt := time.Tick(5 * time.Second)

	log.Println("starting checks")

	go func() {
		for {
			select {
			case <-dc.Quit:
				log.Println("stopped checks")
				return
			case <-tickInt:
				if (checks.MasterApiCheck) {
					go checkMasterApis(dc, checks.MasterApiUrls)
				}
			case <-tickExt:
				if (checks.DnsCheck) {
					go checkDnsNslookupOnKubernetes(dc)

					if (dc.Deamon.IsNode()) {
						go checkDnsServiceNode(dc)
					}

					if (dc.Deamon.IsPod()) {
						go checkDnsInPod(dc)
					}
				}

				if (checks.HttpChecks) {
					if (dc.Deamon.IsNode() || (dc.Deamon.IsPod() && strings.HasSuffix(dc.Deamon.Namespace, "a"))) {
						go checkPodHttpAtoB(dc)
						go checkPodHttpAtoC(dc)
					}

					go checkHttpHaProxy(dc, checks.DeamonPublicUrl)
				}
			}
		}
	}()
}

func stopChecks(dc *models.DeamonClient) {
	dc.Quit <- true
}

func checkDnsNslookupOnKubernetes(dc *models.DeamonClient) {
	handleCheckStarted(dc)
	isOk := false
	var msg string

	cmd := exec.Command("nslookup", deamonDNSEndpoint, kubernetesIP)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		isOk = false
		log.Println("error with nslookup: ", err)
		msg = "DNS resolution via nslookup & kubernetes failed."
	}

	stdOut := out.String()

	if (strings.Contains(stdOut, "Server") && strings.Count(stdOut, "Address") >= 2 && strings.Contains(stdOut, "Name")) {
		isOk = true
	} else {
		msg += "NsLookup had wrong output"
	}

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.DNS_NSLOOKUP_KUBERNETES, IsOk: isOk, Message: msg}
}

func checkDnsServiceNode(dc *models.DeamonClient) {
	handleCheckStarted(dc)
	isOk := false
	var msg string

	ips := getIpsForName(deamonDNSServiceA)

	if (ips == nil) {
		isOk = false
		msg = "Failed to lookup ip on node (dnsmasq) for name " + deamonDNSServiceA
	}

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.DNS_SERVICE_NODE, IsOk: isOk, Message: msg}
}

func checkDnsInPod(dc *models.DeamonClient) {
	handleCheckStarted(dc)
	isOk := false
	var msg string

	ips := getIpsForName(deamonDNSPod)

	if (ips == nil) {
		isOk = false
		msg = "Failed to lookup ip in pod for name " + deamonDNSPod
	} else {
		isOk = true
	}

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.DNS_SERVICE_POD, IsOk: isOk, Message: msg}
}

func getIpsForName(n string) []net.IP {
	ips, err := net.LookupIP(n)
	if (err != nil) {
		log.Println("failed to lookup ip for name ", n)
		return nil
	}
	return ips
}

func checkMasterApis(dc *models.DeamonClient, urls string) {
	handleCheckStarted(dc)
	urlArr := strings.Split(urls, ",")

	oneApiOk := false
	var msg string
	for _, u := range urlArr {
		if (checkHttp(u)) {
			oneApiOk = true
		} else {
			msg += u + " is not reachable. ";
		}
	}

	handleCheckFinished(dc, oneApiOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.MASTER_API_CHECK, IsOk: oneApiOk, Message: msg}
}

func checkHttp(toCall string) bool {
	if (strings.HasPrefix(toCall, "https")) {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
		client := &http.Client{Transport: tr}
		_, err := client.Get(toCall)
		if (err != nil) {
			log.Println("error in http check: ", err)
		}
		return err == nil
	} else {
		_, err := http.Get(toCall)
		if (err != nil) {
			log.Println("error in http check: ", err)
		}
		return err == nil
	}
}

func checkPodHttpAtoB(dc *models.DeamonClient) {
	// This should fail as we do not have access to this project
	handleCheckStarted(dc)
	var msg string

	isOk := !checkHttp("http://" + deamonDNSServiceB + ":8090/hello")

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.HTTP_POD_SERVICE_A_B, IsOk: isOk, Message: msg}
}

func checkPodHttpAtoC(dc *models.DeamonClient) {
	// This should work as we joined this projects
	handleCheckStarted(dc)
	var msg string

	isOk := checkHttp("http://" + deamonDNSServiceC + ":8090/hello")

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.HTTP_POD_SERVICE_A_C, IsOk: isOk, Message: msg}
}

func checkHttpHaProxy(dc *models.DeamonClient, publicUrl string) {
	handleCheckStarted(dc)
	var msg string

	isOk := checkHttp(publicUrl + ":80/hello")

	handleCheckFinished(dc, isOk)

	// Tell the hub about it
	dc.ToHub <- models.CheckResult{Type: models.HTTP_HAPROXY, IsOk: isOk, Message: msg}
}
<|endoftext|>


<commit_before>/*
RpcClient for Go RPC Servers
Copyright (C) 2012-2014 ITsysCOM GmbH

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package rpcclient

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	//"log/syslog"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	JSON_RPC       = "json"
	JSON_HTTP      = "http_jsonrpc"
	GOB_RPC        = "gob"
	INTERNAL_RPC   = "*internal"
	POOL_FIRST     = "first"
	POOL_RANDOM    = "random"
	POOL_NEXT      = "next"
	POOL_BROADCAST = "broadcast"
)

var (
	ErrReqUnsynchronized       = errors.New("REQ_UNSYNCHRONIZED")
	ErrUnsupporteServiceMethod = errors.New("UNSUPPORTED_SERVICE_METHOD")
	ErrWrongArgsType           = errors.New("WRONG_ARGS_TYPE")
	ErrWrongReplyType          = errors.New("WRONG_REPLY_TYPE")
	ErrDisconnected            = errors.New("DISCONNECTED")
	ErrReplyTimeout            = errors.New("REPLY_TIMEOUT")
	//logger                     *syslog.Writer
)

func init() {
	//logger, _ = syslog.New(syslog.LOG_INFO, "RPCClient") // If we need to report anything to syslog
}

// successive Fibonacci numbers.
func Fib() func() time.Duration {
	a, b := 0, 1
	return func() time.Duration {
		a, b = b, a+b
		return time.Duration(a) * time.Second
	}
}

func NewRpcClient(transport, addr string, connectAttempts, reconnects int, connTimeout, replyTimeout time.Duration, codec string, internalConn RpcClientConnection) (*RpcClient, error) {
	var err error
	rpcClient := &RpcClient{transport: transport, address: addr, reconnects: reconnects,
		connTimeout: connTimeout, replyTimeout: replyTimeout, codec: codec, connection: internalConn, connMux: new(sync.Mutex)}
	delay := Fib()
	for i := 0; i < connectAttempts; i++ {
		err = rpcClient.connect()
		if err == nil { //Connected so no need to reiterate
			break
		}
		time.Sleep(delay())
	}
	return rpcClient, err
}

type RpcClient struct {
	transport    string
	address      string
	reconnects   int
	connTimeout  time.Duration
	replyTimeout time.Duration
	codec        string // JSON_RPC or GOB_RPC
	connection   RpcClientConnection
	connMux      *sync.Mutex
}

func (self *RpcClient) connect() (err error) {
	self.connMux.Lock()
	defer self.connMux.Unlock()
	if self.codec == INTERNAL_RPC {
		return nil
	} else if self.codec == JSON_HTTP {
		self.connection = &HttpJsonRpcClient{httpClient: new(http.Client), url: self.address}
		return
	}
	// RPC compliant connections here, manually create connection to timeout
	netconn, err := net.DialTimeout(self.transport, self.address, self.connTimeout)
	if err != nil {
		return err
	}
	if self.codec == JSON_RPC {
		self.connection = jsonrpc.NewClient(netconn)
	} else {
		self.connection = rpc.NewClient(netconn)
	}
	if err != nil {
		self.connection = nil // So we don't wrap nil into the interface
	}
	return
}

func (self *RpcClient) reconnect() (err error) {
	if self.codec == JSON_HTTP { // http client has automatic reconnects in place
		return self.connect()
	}
	i := 0
	delay := Fib()
	for {
		if self.reconnects != -1 && i >= self.reconnects { // Maximum reconnects reached, -1 for infinite reconnects
			break
		}
		if err = self.connect(); err == nil { // No error on connect, succcess
			return nil
		}
		i++
		time.Sleep(delay()) // Cound not reconnect, retry
	}
	return errors.New("RECONNECT_FAIL")
}

func (self *RpcClient) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	if args == nil {
		return fmt.Errorf("nil rpc in argument method: %s in: %v out: %v", serviceMethod, args, reply)
	}
	if self.connection == nil {
		err = ErrDisconnected
	} else {
		errChan := make(chan error, 1)
		go func() {
			errChan <- self.connection.Call(serviceMethod, args, reply)
		}()
		select {
		case err = <-errChan:
		case <-time.After(self.replyTimeout):
			err = ErrReplyTimeout
		}
	}
	if isNetworkError(err) && err != ErrReplyTimeout && self.reconnects != 0 { // ReplyTimeout should not reconnect since it creates loop
		if errReconnect := self.reconnect(); errReconnect != nil {
			return err
		} else { // Run command after reconnect
			return self.Call(serviceMethod, args, reply)
		}
	}
	return err
}

// Connection used in RpcClient, as interface so we can combine the rpc.RpcClient with http one or websocket
type RpcClientConnection interface {
	Call(string, interface{}, interface{}) error
}

// Response received for
type JsonRpcResponse struct {
	Id     uint64
	Result *json.RawMessage
	Error  interface{}
}

type HttpJsonRpcClient struct {
	httpClient *http.Client
	id         uint64
	url        string
}

func (self *HttpJsonRpcClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
	self.id += 1
	id := self.id
	data, err := json.Marshal(map[string]interface{}{
		"method": serviceMethod,
		"id":     self.id,
		"params": [1]interface{}{args},
	})
	if err != nil {
		return err
	}
	resp, err := self.httpClient.Post(self.url, "application/json", ioutil.NopCloser(strings.NewReader(string(data)))) // Closer so we automatically have close after response
	if err != nil {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var jsonRsp JsonRpcResponse
	err = json.Unmarshal(body, &jsonRsp)
	if err != nil {
		return err
	}
	if jsonRsp.Id != id {
		return ErrReqUnsynchronized
	}
	if jsonRsp.Error != nil || jsonRsp.Result == nil {
		x, ok := jsonRsp.Error.(string)
		if !ok {
			return fmt.Errorf("invalid error %v", jsonRsp.Error)
		}
		if x == "" {
			x = "unspecified error"
		}
		return errors.New(x)
	}
	return json.Unmarshal(*jsonRsp.Result, reply)
}

type RpcClientPool struct {
	transmissionType string
	connections      []RpcClientConnection
	counter          int
	replyTimeout     time.Duration
}

func NewRpcClientPool(transmissionType string, replyTimeout time.Duration) *RpcClientPool {
	return &RpcClientPool{transmissionType: transmissionType, replyTimeout: replyTimeout}
}

func (pool *RpcClientPool) AddClient(rcc RpcClientConnection) {
	if rcc != nil {
		pool.connections = append(pool.connections, rcc)
	}
}

func (pool *RpcClientPool) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	switch pool.transmissionType {
	case POOL_BROADCAST:
		replyChan := make(chan *rpcReplyError, len(pool.connections))
		for _, rc := range pool.connections {
			go func(conn RpcClientConnection) {
				// make a new pointer of the same type
				rpl := reflect.New(reflect.TypeOf(reflect.ValueOf(reply).Elem().Interface()))
				err := conn.Call(serviceMethod, args, rpl.Interface())
				if !isNetworkError(err) {
					replyChan <- &rpcReplyError{reply: rpl.Interface(), err: err}
				}
			}(rc)
		}
		//get first response with timeout
		var re *rpcReplyError
		select {
		case re = <-replyChan:
		case <-time.After(pool.replyTimeout):
			return ErrReplyTimeout
		}
		// put received value in the orig reply
		reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(re.reply).Elem())
		return re.err
	case POOL_FIRST:
		for _, rc := range pool.connections {
			err = rc.Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_NEXT:
		ln := len(pool.connections)
		rrIndexes := roundIndex(int(math.Mod(float64(pool.counter), float64(ln))), ln)
		pool.counter++
		for _, index := range rrIndexes {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_RANDOM:
		rand.Seed(time.Now().UnixNano())
		randomIndex := rand.Perm(len(pool.connections))
		for _, index := range randomIndex {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	}
	return
}

type rpcReplyError struct {
	reply interface{}
	err   error
}

// generates round robin indexes for a slice of length max
// starting from index start
func roundIndex(start, max int) []int {
	if start < 0 {
		start = 0
	}
	result := make([]int, max)
	for i := 0; i < max; i++ {
		if start+i < max {
			result[i] = start + i
		} else {
			result[i] = int(math.Abs(float64(max - (start + i))))
		}
	}
	return result
}

func isNetworkError(err error) bool {
	if err == nil {
		return false
	}
	if operr, ok := err.(*net.OpError); ok && strings.HasSuffix(operr.Err.Error(), syscall.ECONNRESET.Error()) { // connection reset
		return true
	}
	return err == rpc.ErrShutdown ||
		err == ErrReqUnsynchronized ||
		err == ErrDisconnected ||
		err == ErrReplyTimeout ||
		strings.HasPrefix(err.Error(), "rpc: can't find service")
}
<commit_msg>Empty connection on errors<commit_after>/*
RpcClient for Go RPC Servers
Copyright (C) ITsysCOM GmbH

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package rpcclient

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	//"log/syslog"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	JSON_RPC       = "json"
	JSON_HTTP      = "http_jsonrpc"
	GOB_RPC        = "gob"
	INTERNAL_RPC   = "*internal"
	POOL_FIRST     = "first"
	POOL_RANDOM    = "random"
	POOL_NEXT      = "next"
	POOL_BROADCAST = "broadcast"
)

var (
	ErrReqUnsynchronized       = errors.New("REQ_UNSYNCHRONIZED")
	ErrUnsupporteServiceMethod = errors.New("UNSUPPORTED_SERVICE_METHOD")
	ErrWrongArgsType           = errors.New("WRONG_ARGS_TYPE")
	ErrWrongReplyType          = errors.New("WRONG_REPLY_TYPE")
	ErrDisconnected            = errors.New("DISCONNECTED")
	ErrReplyTimeout            = errors.New("REPLY_TIMEOUT")
	//logger                     *syslog.Writer
)

func init() {
	//logger, _ = syslog.New(syslog.LOG_INFO, "RPCClient") // If we need to report anything to syslog
}

// successive Fibonacci numbers.
func Fib() func() time.Duration {
	a, b := 0, 1
	return func() time.Duration {
		a, b = b, a+b
		return time.Duration(a) * time.Second
	}
}

func NewRpcClient(transport, addr string, connectAttempts, reconnects int, connTimeout, replyTimeout time.Duration, codec string, internalConn RpcClientConnection) (*RpcClient, error) {
	var err error
	rpcClient := &RpcClient{transport: transport, address: addr, reconnects: reconnects,
		connTimeout: connTimeout, replyTimeout: replyTimeout, codec: codec, connection: internalConn, connMux: new(sync.Mutex)}
	delay := Fib()
	for i := 0; i < connectAttempts; i++ {
		err = rpcClient.connect()
		if err == nil { //Connected so no need to reiterate
			break
		}
		time.Sleep(delay())
	}
	return rpcClient, err
}

type RpcClient struct {
	transport    string
	address      string
	reconnects   int
	connTimeout  time.Duration
	replyTimeout time.Duration
	codec        string // JSON_RPC or GOB_RPC
	connection   RpcClientConnection
	connMux      *sync.Mutex
}

func (self *RpcClient) connect() (err error) {
	self.connMux.Lock()
	defer self.connMux.Unlock()
	if self.codec == INTERNAL_RPC {
		return nil
	} else if self.codec == JSON_HTTP {
		self.connection = &HttpJsonRpcClient{httpClient: new(http.Client), url: self.address}
		return
	}
	// RPC compliant connections here, manually create connection to timeout
	netconn, err := net.DialTimeout(self.transport, self.address, self.connTimeout)
	if err != nil {
		self.connection = nil // So we don't wrap nil into the interface
		return err
	}
	if self.codec == JSON_RPC {
		self.connection = jsonrpc.NewClient(netconn)
	} else {
		self.connection = rpc.NewClient(netconn)
	}
	return
}

func (self *RpcClient) reconnect() (err error) {
	if self.codec == JSON_HTTP { // http client has automatic reconnects in place
		return self.connect()
	}
	i := 0
	delay := Fib()
	for {
		if self.reconnects != -1 && i >= self.reconnects { // Maximum reconnects reached, -1 for infinite reconnects
			break
		}
		if err = self.connect(); err == nil { // No error on connect, succcess
			return nil
		}
		i++
		time.Sleep(delay()) // Cound not reconnect, retry
	}
	return errors.New("RECONNECT_FAIL")
}

func (self *RpcClient) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	if args == nil {
		return fmt.Errorf("nil rpc in argument method: %s in: %v out: %v", serviceMethod, args, reply)
	}
	if self.connection == nil {
		err = ErrDisconnected
	} else {
		errChan := make(chan error, 1)
		go func() {
			errChan <- self.connection.Call(serviceMethod, args, reply)
		}()
		select {
		case err = <-errChan:
		case <-time.After(self.replyTimeout):
			err = ErrReplyTimeout
		}
	}
	if isNetworkError(err) && err != ErrReplyTimeout && self.reconnects != 0 { // ReplyTimeout should not reconnect since it creates loop
		if errReconnect := self.reconnect(); errReconnect != nil {
			return err
		} else { // Run command after reconnect
			return self.Call(serviceMethod, args, reply)
		}
	}
	return err
}

// Connection used in RpcClient, as interface so we can combine the rpc.RpcClient with http one or websocket
type RpcClientConnection interface {
	Call(string, interface{}, interface{}) error
}

// Response received for
type JsonRpcResponse struct {
	Id     uint64
	Result *json.RawMessage
	Error  interface{}
}

type HttpJsonRpcClient struct {
	httpClient *http.Client
	id         uint64
	url        string
}

func (self *HttpJsonRpcClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
	self.id += 1
	id := self.id
	data, err := json.Marshal(map[string]interface{}{
		"method": serviceMethod,
		"id":     self.id,
		"params": [1]interface{}{args},
	})
	if err != nil {
		return err
	}
	resp, err := self.httpClient.Post(self.url, "application/json", ioutil.NopCloser(strings.NewReader(string(data)))) // Closer so we automatically have close after response
	if err != nil {
		return err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	var jsonRsp JsonRpcResponse
	err = json.Unmarshal(body, &jsonRsp)
	if err != nil {
		return err
	}
	if jsonRsp.Id != id {
		return ErrReqUnsynchronized
	}
	if jsonRsp.Error != nil || jsonRsp.Result == nil {
		x, ok := jsonRsp.Error.(string)
		if !ok {
			return fmt.Errorf("invalid error %v", jsonRsp.Error)
		}
		if x == "" {
			x = "unspecified error"
		}
		return errors.New(x)
	}
	return json.Unmarshal(*jsonRsp.Result, reply)
}

type RpcClientPool struct {
	transmissionType string
	connections      []RpcClientConnection
	counter          int
	replyTimeout     time.Duration
}

func NewRpcClientPool(transmissionType string, replyTimeout time.Duration) *RpcClientPool {
	return &RpcClientPool{transmissionType: transmissionType, replyTimeout: replyTimeout}
}

func (pool *RpcClientPool) AddClient(rcc RpcClientConnection) {
	if rcc != nil {
		pool.connections = append(pool.connections, rcc)
	}
}

func (pool *RpcClientPool) Call(serviceMethod string, args interface{}, reply interface{}) (err error) {
	switch pool.transmissionType {
	case POOL_BROADCAST:
		replyChan := make(chan *rpcReplyError, len(pool.connections))
		for _, rc := range pool.connections {
			go func(conn RpcClientConnection) {
				// make a new pointer of the same type
				rpl := reflect.New(reflect.TypeOf(reflect.ValueOf(reply).Elem().Interface()))
				err := conn.Call(serviceMethod, args, rpl.Interface())
				if !isNetworkError(err) {
					replyChan <- &rpcReplyError{reply: rpl.Interface(), err: err}
				}
			}(rc)
		}
		//get first response with timeout
		var re *rpcReplyError
		select {
		case re = <-replyChan:
		case <-time.After(pool.replyTimeout):
			return ErrReplyTimeout
		}
		// put received value in the orig reply
		reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(re.reply).Elem())
		return re.err
	case POOL_FIRST:
		for _, rc := range pool.connections {
			err = rc.Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_NEXT:
		ln := len(pool.connections)
		rrIndexes := roundIndex(int(math.Mod(float64(pool.counter), float64(ln))), ln)
		pool.counter++
		for _, index := range rrIndexes {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	case POOL_RANDOM:
		rand.Seed(time.Now().UnixNano())
		randomIndex := rand.Perm(len(pool.connections))
		for _, index := range randomIndex {
			err = pool.connections[index].Call(serviceMethod, args, reply)
			if isNetworkError(err) {
				continue
			}
			return
		}
	}
	return
}

type rpcReplyError struct {
	reply interface{}
	err   error
}

// generates round robin indexes for a slice of length max
// starting from index start
func roundIndex(start, max int) []int {
	if start < 0 {
		start = 0
	}
	result := make([]int, max)
	for i := 0; i < max; i++ {
		if start+i < max {
			result[i] = start + i
		} else {
			result[i] = int(math.Abs(float64(max - (start + i))))
		}
	}
	return result
}

func isNetworkError(err error) bool {
	if err == nil {
		return false
	}
	if operr, ok := err.(*net.OpError); ok && strings.HasSuffix(operr.Err.Error(), syscall.ECONNRESET.Error()) { // connection reset
		return true
	}
	return err == rpc.ErrShutdown ||
		err == ErrReqUnsynchronized ||
		err == ErrDisconnected ||
		err == ErrReplyTimeout ||
		strings.HasPrefix(err.Error(), "rpc: can't find service")
}
<|endoftext|>


<commit_before>package sarama

import (
	"testing"
	"time"
)

func initOffsetManager(t *testing.T) (om OffsetManager,
	testClient Client, broker, coordinator *MockBroker) {

	config := NewConfig()
	config.Metadata.Retry.Max = 1
	config.Consumer.Offsets.CommitInterval = 1 * time.Millisecond
	config.Version = V0_9_0_0

	broker = NewMockBroker(t, 1)
	coordinator = NewMockBroker(t, 2)

	seedMeta := new(MetadataResponse)
	seedMeta.AddBroker(coordinator.Addr(), coordinator.BrokerID())
	seedMeta.AddTopicPartition("my_topic", 0, 1, []int32{}, []int32{}, ErrNoError)
	seedMeta.AddTopicPartition("my_topic", 1, 1, []int32{}, []int32{}, ErrNoError)
	broker.Returns(seedMeta)

	var err error
	testClient, err = NewClient([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   coordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: coordinator.Port(),
	})

	om, err = NewOffsetManagerFromClient("group", testClient)
	if err != nil {
		t.Fatal(err)
	}

	return om, testClient, broker, coordinator
}

func initPartitionOffsetManager(t *testing.T, om OffsetManager,
	coordinator *MockBroker, initialOffset int64, metadata string) PartitionOffsetManager {

	fetchResponse := new(OffsetFetchResponse)
	fetchResponse.AddBlock("my_topic", 0, &OffsetFetchResponseBlock{
		Err:      ErrNoError,
		Offset:   initialOffset,
		Metadata: metadata,
	})
	coordinator.Returns(fetchResponse)

	pom, err := om.ManagePartition("my_topic", 0)
	if err != nil {
		t.Fatal(err)
	}

	return pom
}

func TestNewOffsetManager(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	seedBroker.Returns(new(MetadataResponse))
	defer seedBroker.Close()

	testClient, err := NewClient([]string{seedBroker.Addr()}, nil)
	if err != nil {
		t.Fatal(err)
	}

	om, err := NewOffsetManagerFromClient("group", testClient)
	if err != nil {
		t.Error(err)
	}
	safeClose(t, om)
	safeClose(t, testClient)

	_, err = NewOffsetManagerFromClient("group", testClient)
	if err != ErrClosedClient {
		t.Errorf("Error expected for closed client; actual value: %v", err)
	}
}

// Test recovery from ErrNotCoordinatorForConsumer
// on first fetchInitialOffset call
func TestOffsetManagerFetchInitialFail(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)

	// Error on first fetchInitialOffset call
	responseBlock := OffsetFetchResponseBlock{
		Err:      ErrNotCoordinatorForConsumer,
		Offset:   5,
		Metadata: "test_meta",
	}

	fetchResponse := new(OffsetFetchResponse)
	fetchResponse.AddBlock("my_topic", 0, &responseBlock)
	coordinator.Returns(fetchResponse)

	// Refresh coordinator
	newCoordinator := NewMockBroker(t, 3)
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	// Second fetchInitialOffset call is fine
	fetchResponse2 := new(OffsetFetchResponse)
	responseBlock2 := responseBlock
	responseBlock2.Err = ErrNoError
	fetchResponse2.AddBlock("my_topic", 0, &responseBlock2)
	newCoordinator.Returns(fetchResponse2)

	pom, err := om.ManagePartition("my_topic", 0)
	if err != nil {
		t.Error(err)
	}

	broker.Close()
	coordinator.Close()
	newCoordinator.Close()
	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
}

// Test fetchInitialOffset retry on ErrOffsetsLoadInProgress
func TestOffsetManagerFetchInitialLoadInProgress(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)

	// Error on first fetchInitialOffset call
	responseBlock := OffsetFetchResponseBlock{
		Err:      ErrOffsetsLoadInProgress,
		Offset:   5,
		Metadata: "test_meta",
	}

	fetchResponse := new(OffsetFetchResponse)
	fetchResponse.AddBlock("my_topic", 0, &responseBlock)
	coordinator.Returns(fetchResponse)

	// Second fetchInitialOffset call is fine
	fetchResponse2 := new(OffsetFetchResponse)
	responseBlock2 := responseBlock
	responseBlock2.Err = ErrNoError
	fetchResponse2.AddBlock("my_topic", 0, &responseBlock2)
	coordinator.Returns(fetchResponse2)

	pom, err := om.ManagePartition("my_topic", 0)
	if err != nil {
		t.Error(err)
	}

	broker.Close()
	coordinator.Close()
	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
}

func TestPartitionOffsetManagerInitialOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	testClient.Config().Consumer.Offsets.Initial = OffsetOldest

	// Kafka returns -1 if no offset has been stored for this partition yet.
	pom := initPartitionOffsetManager(t, om, coordinator, -1, "")

	offset, meta := pom.NextOffset()
	if offset != OffsetOldest {
		t.Errorf("Expected offset 5. Actual: %v", offset)
	}
	if meta != "" {
		t.Errorf("Expected metadata to be empty. Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	broker.Close()
	coordinator.Close()
	safeClose(t, testClient)
}

func TestPartitionOffsetManagerNextOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "test_meta")

	offset, meta := pom.NextOffset()
	if offset != 5 {
		t.Errorf("Expected offset 5. Actual: %v", offset)
	}
	if meta != "test_meta" {
		t.Errorf("Expected metadata \"test_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	broker.Close()
	coordinator.Close()
	safeClose(t, testClient)
}

func TestPartitionOffsetManagerResetOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	coordinator.Returns(ocResponse)

	expected := int64(1)
	pom.ResetOffset(expected, "modified_meta")
	actual, meta := pom.NextOffset()

	if actual != expected {
		t.Errorf("Expected offset %v. Actual: %v", expected, actual)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerResetOffsetWithRetention(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	testClient.Config().Consumer.Offsets.Retention = time.Hour

	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	handler := func(req *request) (res encoder) {
		if req.body.version() != 2 {
			t.Errorf("Expected to be using version 2. Actual: %v", req.body.version())
		}
		offsetCommitRequest := req.body.(*OffsetCommitRequest)
		if offsetCommitRequest.RetentionTime != (60 * 60 * 1000) {
			t.Errorf("Expected an hour retention time. Actual: %v", offsetCommitRequest.RetentionTime)
		}
		return ocResponse
	}
	coordinator.setHandler(handler)

	expected := int64(1)
	pom.ResetOffset(expected, "modified_meta")
	actual, meta := pom.NextOffset()

	if actual != expected {
		t.Errorf("Expected offset %v. Actual: %v", expected, actual)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerMarkOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	coordinator.Returns(ocResponse)

	pom.MarkOffset(100, "modified_meta")
	offset, meta := pom.NextOffset()

	if offset != 100 {
		t.Errorf("Expected offset 100. Actual: %v", offset)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerMarkOffsetWithRetention(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	testClient.Config().Consumer.Offsets.Retention = time.Hour

	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	handler := func(req *request) (res encoder) {
		if req.body.version() != 2 {
			t.Errorf("Expected to be using version 2. Actual: %v", req.body.version())
		}
		offsetCommitRequest := req.body.(*OffsetCommitRequest)
		if offsetCommitRequest.RetentionTime != (60 * 60 * 1000) {
			t.Errorf("Expected an hour retention time. Actual: %v", offsetCommitRequest.RetentionTime)
		}
		return ocResponse
	}
	coordinator.setHandler(handler)

	pom.MarkOffset(100, "modified_meta")
	offset, meta := pom.NextOffset()

	if offset != 100 {
		t.Errorf("Expected offset 100. Actual: %v", offset)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerCommitErr(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "meta")

	// Error on one partition
	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrOffsetOutOfRange)
	ocResponse.AddError("my_topic", 1, ErrNoError)
	coordinator.Returns(ocResponse)

	newCoordinator := NewMockBroker(t, 3)

	// For RefreshCoordinator()
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	// Nothing in response.Errors at all
	ocResponse2 := new(OffsetCommitResponse)
	newCoordinator.Returns(ocResponse2)

	// No error, no need to refresh coordinator

	// Error on the wrong partition for this pom
	ocResponse3 := new(OffsetCommitResponse)
	ocResponse3.AddError("my_topic", 1, ErrNoError)
	newCoordinator.Returns(ocResponse3)

	// No error, no need to refresh coordinator

	// ErrUnknownTopicOrPartition/ErrNotLeaderForPartition/ErrLeaderNotAvailable block
	ocResponse4 := new(OffsetCommitResponse)
	ocResponse4.AddError("my_topic", 0, ErrUnknownTopicOrPartition)
	newCoordinator.Returns(ocResponse4)

	// For RefreshCoordinator()
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	// Normal error response
	ocResponse5 := new(OffsetCommitResponse)
	ocResponse5.AddError("my_topic", 0, ErrNoError)
	newCoordinator.Returns(ocResponse5)

	pom.MarkOffset(100, "modified_meta")

	err := pom.Close()
	if err != nil {
		t.Error(err)
	}

	broker.Close()
	coordinator.Close()
	newCoordinator.Close()
	safeClose(t, om)
	safeClose(t, testClient)
}

// Test of recovery from abort
func TestAbortPartitionOffsetManager(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "meta")

	// this triggers an error in the CommitOffset request,
	// which leads to the abort call
	coordinator.Close()

	// Response to refresh coordinator request
	newCoordinator := NewMockBroker(t, 3)
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	newCoordinator.Returns(ocResponse)

	pom.MarkOffset(100, "modified_meta")

	safeClose(t, pom)
	safeClose(t, om)
	broker.Close()
	safeClose(t, testClient)
}
<commit_msg>Fix test race<commit_after>package sarama

import (
	"testing"
	"time"
)

func initOffsetManager(t *testing.T, retention time.Duration) (om OffsetManager,
	testClient Client, broker, coordinator *MockBroker) {

	config := NewConfig()
	config.Metadata.Retry.Max = 1
	config.Consumer.Offsets.CommitInterval = 1 * time.Millisecond
	config.Version = V0_9_0_0
	if retention > 0 {
		config.Consumer.Offsets.Retention = retention
	}

	broker = NewMockBroker(t, 1)
	coordinator = NewMockBroker(t, 2)

	seedMeta := new(MetadataResponse)
	seedMeta.AddBroker(coordinator.Addr(), coordinator.BrokerID())
	seedMeta.AddTopicPartition("my_topic", 0, 1, []int32{}, []int32{}, ErrNoError)
	seedMeta.AddTopicPartition("my_topic", 1, 1, []int32{}, []int32{}, ErrNoError)
	broker.Returns(seedMeta)

	var err error
	testClient, err = NewClient([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   coordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: coordinator.Port(),
	})

	om, err = NewOffsetManagerFromClient("group", testClient)
	if err != nil {
		t.Fatal(err)
	}

	return om, testClient, broker, coordinator
}

func initPartitionOffsetManager(t *testing.T, om OffsetManager,
	coordinator *MockBroker, initialOffset int64, metadata string) PartitionOffsetManager {

	fetchResponse := new(OffsetFetchResponse)
	fetchResponse.AddBlock("my_topic", 0, &OffsetFetchResponseBlock{
		Err:      ErrNoError,
		Offset:   initialOffset,
		Metadata: metadata,
	})
	coordinator.Returns(fetchResponse)

	pom, err := om.ManagePartition("my_topic", 0)
	if err != nil {
		t.Fatal(err)
	}

	return pom
}

func TestNewOffsetManager(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	seedBroker.Returns(new(MetadataResponse))
	defer seedBroker.Close()

	testClient, err := NewClient([]string{seedBroker.Addr()}, nil)
	if err != nil {
		t.Fatal(err)
	}

	om, err := NewOffsetManagerFromClient("group", testClient)
	if err != nil {
		t.Error(err)
	}
	safeClose(t, om)
	safeClose(t, testClient)

	_, err = NewOffsetManagerFromClient("group", testClient)
	if err != ErrClosedClient {
		t.Errorf("Error expected for closed client; actual value: %v", err)
	}
}

// Test recovery from ErrNotCoordinatorForConsumer
// on first fetchInitialOffset call
func TestOffsetManagerFetchInitialFail(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)

	// Error on first fetchInitialOffset call
	responseBlock := OffsetFetchResponseBlock{
		Err:      ErrNotCoordinatorForConsumer,
		Offset:   5,
		Metadata: "test_meta",
	}

	fetchResponse := new(OffsetFetchResponse)
	fetchResponse.AddBlock("my_topic", 0, &responseBlock)
	coordinator.Returns(fetchResponse)

	// Refresh coordinator
	newCoordinator := NewMockBroker(t, 3)
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	// Second fetchInitialOffset call is fine
	fetchResponse2 := new(OffsetFetchResponse)
	responseBlock2 := responseBlock
	responseBlock2.Err = ErrNoError
	fetchResponse2.AddBlock("my_topic", 0, &responseBlock2)
	newCoordinator.Returns(fetchResponse2)

	pom, err := om.ManagePartition("my_topic", 0)
	if err != nil {
		t.Error(err)
	}

	broker.Close()
	coordinator.Close()
	newCoordinator.Close()
	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
}

// Test fetchInitialOffset retry on ErrOffsetsLoadInProgress
func TestOffsetManagerFetchInitialLoadInProgress(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)

	// Error on first fetchInitialOffset call
	responseBlock := OffsetFetchResponseBlock{
		Err:      ErrOffsetsLoadInProgress,
		Offset:   5,
		Metadata: "test_meta",
	}

	fetchResponse := new(OffsetFetchResponse)
	fetchResponse.AddBlock("my_topic", 0, &responseBlock)
	coordinator.Returns(fetchResponse)

	// Second fetchInitialOffset call is fine
	fetchResponse2 := new(OffsetFetchResponse)
	responseBlock2 := responseBlock
	responseBlock2.Err = ErrNoError
	fetchResponse2.AddBlock("my_topic", 0, &responseBlock2)
	coordinator.Returns(fetchResponse2)

	pom, err := om.ManagePartition("my_topic", 0)
	if err != nil {
		t.Error(err)
	}

	broker.Close()
	coordinator.Close()
	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
}

func TestPartitionOffsetManagerInitialOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)
	testClient.Config().Consumer.Offsets.Initial = OffsetOldest

	// Kafka returns -1 if no offset has been stored for this partition yet.
	pom := initPartitionOffsetManager(t, om, coordinator, -1, "")

	offset, meta := pom.NextOffset()
	if offset != OffsetOldest {
		t.Errorf("Expected offset 5. Actual: %v", offset)
	}
	if meta != "" {
		t.Errorf("Expected metadata to be empty. Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	broker.Close()
	coordinator.Close()
	safeClose(t, testClient)
}

func TestPartitionOffsetManagerNextOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "test_meta")

	offset, meta := pom.NextOffset()
	if offset != 5 {
		t.Errorf("Expected offset 5. Actual: %v", offset)
	}
	if meta != "test_meta" {
		t.Errorf("Expected metadata \"test_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	broker.Close()
	coordinator.Close()
	safeClose(t, testClient)
}

func TestPartitionOffsetManagerResetOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	coordinator.Returns(ocResponse)

	expected := int64(1)
	pom.ResetOffset(expected, "modified_meta")
	actual, meta := pom.NextOffset()

	if actual != expected {
		t.Errorf("Expected offset %v. Actual: %v", expected, actual)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerResetOffsetWithRetention(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, time.Hour)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	handler := func(req *request) (res encoder) {
		if req.body.version() != 2 {
			t.Errorf("Expected to be using version 2. Actual: %v", req.body.version())
		}
		offsetCommitRequest := req.body.(*OffsetCommitRequest)
		if offsetCommitRequest.RetentionTime != (60 * 60 * 1000) {
			t.Errorf("Expected an hour retention time. Actual: %v", offsetCommitRequest.RetentionTime)
		}
		return ocResponse
	}
	coordinator.setHandler(handler)

	expected := int64(1)
	pom.ResetOffset(expected, "modified_meta")
	actual, meta := pom.NextOffset()

	if actual != expected {
		t.Errorf("Expected offset %v. Actual: %v", expected, actual)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerMarkOffset(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	coordinator.Returns(ocResponse)

	pom.MarkOffset(100, "modified_meta")
	offset, meta := pom.NextOffset()

	if offset != 100 {
		t.Errorf("Expected offset 100. Actual: %v", offset)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerMarkOffsetWithRetention(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, time.Hour)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "original_meta")

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	handler := func(req *request) (res encoder) {
		if req.body.version() != 2 {
			t.Errorf("Expected to be using version 2. Actual: %v", req.body.version())
		}
		offsetCommitRequest := req.body.(*OffsetCommitRequest)
		if offsetCommitRequest.RetentionTime != (60 * 60 * 1000) {
			t.Errorf("Expected an hour retention time. Actual: %v", offsetCommitRequest.RetentionTime)
		}
		return ocResponse
	}
	coordinator.setHandler(handler)

	pom.MarkOffset(100, "modified_meta")
	offset, meta := pom.NextOffset()

	if offset != 100 {
		t.Errorf("Expected offset 100. Actual: %v", offset)
	}
	if meta != "modified_meta" {
		t.Errorf("Expected metadata \"modified_meta\". Actual: %q", meta)
	}

	safeClose(t, pom)
	safeClose(t, om)
	safeClose(t, testClient)
	broker.Close()
	coordinator.Close()
}

func TestPartitionOffsetManagerCommitErr(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "meta")

	// Error on one partition
	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrOffsetOutOfRange)
	ocResponse.AddError("my_topic", 1, ErrNoError)
	coordinator.Returns(ocResponse)

	newCoordinator := NewMockBroker(t, 3)

	// For RefreshCoordinator()
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	// Nothing in response.Errors at all
	ocResponse2 := new(OffsetCommitResponse)
	newCoordinator.Returns(ocResponse2)

	// No error, no need to refresh coordinator

	// Error on the wrong partition for this pom
	ocResponse3 := new(OffsetCommitResponse)
	ocResponse3.AddError("my_topic", 1, ErrNoError)
	newCoordinator.Returns(ocResponse3)

	// No error, no need to refresh coordinator

	// ErrUnknownTopicOrPartition/ErrNotLeaderForPartition/ErrLeaderNotAvailable block
	ocResponse4 := new(OffsetCommitResponse)
	ocResponse4.AddError("my_topic", 0, ErrUnknownTopicOrPartition)
	newCoordinator.Returns(ocResponse4)

	// For RefreshCoordinator()
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	// Normal error response
	ocResponse5 := new(OffsetCommitResponse)
	ocResponse5.AddError("my_topic", 0, ErrNoError)
	newCoordinator.Returns(ocResponse5)

	pom.MarkOffset(100, "modified_meta")

	err := pom.Close()
	if err != nil {
		t.Error(err)
	}

	broker.Close()
	coordinator.Close()
	newCoordinator.Close()
	safeClose(t, om)
	safeClose(t, testClient)
}

// Test of recovery from abort
func TestAbortPartitionOffsetManager(t *testing.T) {
	om, testClient, broker, coordinator := initOffsetManager(t, 0)
	pom := initPartitionOffsetManager(t, om, coordinator, 5, "meta")

	// this triggers an error in the CommitOffset request,
	// which leads to the abort call
	coordinator.Close()

	// Response to refresh coordinator request
	newCoordinator := NewMockBroker(t, 3)
	broker.Returns(&ConsumerMetadataResponse{
		CoordinatorID:   newCoordinator.BrokerID(),
		CoordinatorHost: "127.0.0.1",
		CoordinatorPort: newCoordinator.Port(),
	})

	ocResponse := new(OffsetCommitResponse)
	ocResponse.AddError("my_topic", 0, ErrNoError)
	newCoordinator.Returns(ocResponse)

	pom.MarkOffset(100, "modified_meta")

	safeClose(t, pom)
	safeClose(t, om)
	broker.Close()
	safeClose(t, testClient)
}
<|endoftext|>


<commit_before>package rpg2d

import (
	"time"

	"github.com/ghthor/engine/rpg2d/entity"
	"github.com/ghthor/engine/rpg2d/quad"
	"github.com/ghthor/engine/sim"
	"github.com/ghthor/engine/sim/stime"
)

// An interface the user will implement to resolve
// an entity from an actor. This is user defined
// because the user is creating the entities and actors
// that will be added to the simulation. This allows the
// user to define how the actors state is stored,
// aka database design/interaction.
type EntityResolver interface {
	EntityForActor(sim.Actor) entity.Entity
}

// A SimulationDef used to configure a simulation
// to define the how the simulation will behave.
type SimulationDef struct {
	// The target FPS for the simulation to calculate at
	FPS int

	// Initial World State
	Now      stime.Time
	QuadTree quad.Quad

	// User defined to resolve an entity from and actor
	EntityResolver EntityResolver

	// User defined input application phase
	InputPhaseHandler quad.InputPhaseHandler

	// User defined the narrow phase
	NarrowPhaseHandler quad.NarrowPhaseHandler
}

type initialWorldState struct {
	now      stime.Time
	quadTree quad.Quad
}

type simSettings struct {
	fps int

	EntityResolver
	quad.InputPhaseHandler
	quad.NarrowPhaseHandler
}

// An implementation of engine/sim.RunningSimulation
type runningSimulation struct {
	//---- Communication
	// These channels are used by the public api
	// to add and remove actors. They are 1way
	// send only channels. The requests contains
	// a send only and recieve only channel so the
	// the public api call will be atomic action
	// and the caller can assume without a doubt
	// that the actor is now added or removed
	// from the simulation.
	addActor    chan<- addActorReq
	removeActor chan<- removeActorReq

	// This channel is used by the public api to request
	// that the simulation is halted. You send a 1way
	// send only channel into the game loop so the public
	// api can wait and be notified that the go routine
	// has returned and is no longer running.
	requestHalt chan<- chan<- struct{}
}

// Communication object used to atomicly add a new actor to the sim
type addActorReq struct {
	toBeAdded chan sim.Actor
	wasAdded  chan sim.Actor
}

// Implement engine/sim.RunningSimulation
func (s runningSimulation) ConnectActor(a sim.Actor) error {
	ch := make(chan sim.Actor)

	// Create an add request
	actor := addActorReq{ch, ch}

	// Send the add request to the game loop
	s.addActor <- actor

	// Send the actor to be added to the game loop
	actor.toBeAdded <- a

	// Wait for the add request to be successfully completed
	a = <-actor.wasAdded
	return nil
}

// Communication object used to atomicly remove an actor from the sim
type removeActorReq struct {
	toBeRemoved chan sim.Actor
	wasRemoved  chan sim.Actor
}

// Implement engine/sim.RunningSimulation
func (s runningSimulation) RemoveActor(a sim.Actor) error {
	ch := make(chan sim.Actor)

	// Create a remove request
	actor := removeActorReq{ch, ch}

	// Send the remove request to the game loop
	s.removeActor <- actor

	// Send the actor to be removed to the game loop
	actor.toBeRemoved <- a

	// Wait for the remove request to be successfully completed
	a = <-actor.wasRemoved
	return nil
}

// Implement engine/sim.RunningSimulation
func (s runningSimulation) Halt() (sim.HaltedSimulation, error) {
	wasHalted := make(chan struct{})

	// Send a request to the game loop to halt
	s.requestHalt <- wasHalted

	// Wait for the halt request to be successfully completed
	<-wasHalted

	return haltedSimulation{}, nil
}

// Implement engine/sim.UnstartedSimulation
func (s SimulationDef) Begin() (sim.RunningSimulation, error) {
	initialState := initialWorldState{
		now:      s.Now,
		quadTree: s.QuadTree,
	}

	settings := simSettings{
		s.FPS,

		s.EntityResolver,
		s.InputPhaseHandler,
		s.NarrowPhaseHandler,
	}

	rs := &runningSimulation{}

	// Starts 2 go routines and returns
	// The ticker and the engine communication kernel
	rs.startLoop(initialState, settings)

	return rs, nil
}

// Prepares a closure and executes it as a go routine
// Calling this function will create 2 infinte looping
// go routines. One for the clock ticker and one for
// calulating the next world state and adding/removing actors.
// This method has a pointer recv because it MUST set the
// addActor and removeActor communication channels used
// by the public api to request adding & removing actors
func (s *runningSimulation) startLoop(initialState initialWorldState, settings simSettings) {
	//---- Create all the communication channels

	// Make the 2way channels that will be used to make
	// add and remove actor requests to the go routine game loop
	addCh := make(chan addActorReq)
	removeCh := make(chan removeActorReq)

	// Set the 1way send chanels used by the public api
	s.addActor = addCh
	s.removeActor = removeCh

	// Set the 1way recieve channels used by the game loop
	var addReq <-chan addActorReq
	var removeReq <-chan removeActorReq

	addReq = addCh
	removeReq = removeCh

	// Make channel to be used to by the public api to
	// request that the simulation be halted
	haltCh := make(chan chan<- struct{})

	// Set the 1way send channel used by the public api
	s.requestHalt = haltCh

	// Set the 1way recieve channel used by the game loop
	var haltReq <-chan chan<- struct{}
	haltReq = haltCh

	quadTree := initialState.quadTree
	clock := stime.Clock(initialState.now)

	//---- User provided actor to entity resolver
	entityResolver := settings.EntityResolver

	//---- User provided input application phase
	inputPhase := settings.InputPhaseHandler

	//---- User provided narrow phase
	narrowPhase := settings.NarrowPhaseHandler

	runPhase := func(q quad.Quad, t stime.Time) quad.Quad {
		return quad.RunPhasesOn(q, inputPhase, narrowPhase, t)
	}

	// Start the Clock
	ticker := time.NewTicker(time.Duration(1000/settings.fps) * time.Millisecond)

	// Start the simulation server
	go func() {
		var hasHalted chan<- struct{}

	gameLoop:
		for {
			// Prioritized select for ticker.C and haltReq
			select {
			case <-ticker.C:
				goto tick

			case hasHalted = <-haltReq:
				break gameLoop
			default:
			}

			select {
			case <-ticker.C:
				goto tick

			case actor := <-addReq:
				// a is the new sim.Actor{} to be inserted into the sim
				a := <-actor.toBeAdded

				e := entityResolver.EntityForActor(a)
				quadTree = quadTree.Insert(e)

				// signal that the operation was a success
				actor.wasAdded <- a

			case actor := <-removeReq:
				// a is the new sim.Actor{} to be removed from the sim
				a := <-actor.toBeRemoved

				// TODO removed the actor from the simulation

				// signal that the operation was a success
				actor.wasRemoved <- a

			case hasHalted = <-haltReq:
				break gameLoop
			}

		tick:
			clock = clock.Tick()
			quadTree = runPhase(quadTree, clock.Now())
			// TODO quadTree to state
			// TODO send state
			continue
		}

		// We're done with cleanup and going to exit
		hasHalted <- struct{}{}
	}()
}

type haltedSimulation struct{}
<commit_msg>Renamed closure that executes all phases for a frame<commit_after>package rpg2d

import (
	"time"

	"github.com/ghthor/engine/rpg2d/entity"
	"github.com/ghthor/engine/rpg2d/quad"
	"github.com/ghthor/engine/sim"
	"github.com/ghthor/engine/sim/stime"
)

// An interface the user will implement to resolve
// an entity from an actor. This is user defined
// because the user is creating the entities and actors
// that will be added to the simulation. This allows the
// user to define how the actors state is stored,
// aka database design/interaction.
type EntityResolver interface {
	EntityForActor(sim.Actor) entity.Entity
}

// A SimulationDef used to configure a simulation
// to define the how the simulation will behave.
type SimulationDef struct {
	// The target FPS for the simulation to calculate at
	FPS int

	// Initial World State
	Now      stime.Time
	QuadTree quad.Quad

	// User defined to resolve an entity from and actor
	EntityResolver EntityResolver

	// User defined input application phase
	InputPhaseHandler quad.InputPhaseHandler

	// User defined the narrow phase
	NarrowPhaseHandler quad.NarrowPhaseHandler
}

type initialWorldState struct {
	now      stime.Time
	quadTree quad.Quad
}

type simSettings struct {
	fps int

	EntityResolver
	quad.InputPhaseHandler
	quad.NarrowPhaseHandler
}

// An implementation of engine/sim.RunningSimulation
type runningSimulation struct {
	//---- Communication
	// These channels are used by the public api
	// to add and remove actors. They are 1way
	// send only channels. The requests contains
	// a send only and recieve only channel so the
	// the public api call will be atomic action
	// and the caller can assume without a doubt
	// that the actor is now added or removed
	// from the simulation.
	addActor    chan<- addActorReq
	removeActor chan<- removeActorReq

	// This channel is used by the public api to request
	// that the simulation is halted. You send a 1way
	// send only channel into the game loop so the public
	// api can wait and be notified that the go routine
	// has returned and is no longer running.
	requestHalt chan<- chan<- struct{}
}

// Communication object used to atomicly add a new actor to the sim
type addActorReq struct {
	toBeAdded chan sim.Actor
	wasAdded  chan sim.Actor
}

// Implement engine/sim.RunningSimulation
func (s runningSimulation) ConnectActor(a sim.Actor) error {
	ch := make(chan sim.Actor)

	// Create an add request
	actor := addActorReq{ch, ch}

	// Send the add request to the game loop
	s.addActor <- actor

	// Send the actor to be added to the game loop
	actor.toBeAdded <- a

	// Wait for the add request to be successfully completed
	a = <-actor.wasAdded
	return nil
}

// Communication object used to atomicly remove an actor from the sim
type removeActorReq struct {
	toBeRemoved chan sim.Actor
	wasRemoved  chan sim.Actor
}

// Implement engine/sim.RunningSimulation
func (s runningSimulation) RemoveActor(a sim.Actor) error {
	ch := make(chan sim.Actor)

	// Create a remove request
	actor := removeActorReq{ch, ch}

	// Send the remove request to the game loop
	s.removeActor <- actor

	// Send the actor to be removed to the game loop
	actor.toBeRemoved <- a

	// Wait for the remove request to be successfully completed
	a = <-actor.wasRemoved
	return nil
}

// Implement engine/sim.RunningSimulation
func (s runningSimulation) Halt() (sim.HaltedSimulation, error) {
	wasHalted := make(chan struct{})

	// Send a request to the game loop to halt
	s.requestHalt <- wasHalted

	// Wait for the halt request to be successfully completed
	<-wasHalted

	return haltedSimulation{}, nil
}

// Implement engine/sim.UnstartedSimulation
func (s SimulationDef) Begin() (sim.RunningSimulation, error) {
	initialState := initialWorldState{
		now:      s.Now,
		quadTree: s.QuadTree,
	}

	settings := simSettings{
		s.FPS,

		s.EntityResolver,
		s.InputPhaseHandler,
		s.NarrowPhaseHandler,
	}

	rs := &runningSimulation{}

	// Starts 2 go routines and returns
	// The ticker and the engine communication kernel
	rs.startLoop(initialState, settings)

	return rs, nil
}

// Prepares a closure and executes it as a go routine
// Calling this function will create 2 infinte looping
// go routines. One for the clock ticker and one for
// calulating the next world state and adding/removing actors.
// This method has a pointer recv because it MUST set the
// addActor and removeActor communication channels used
// by the public api to request adding & removing actors
func (s *runningSimulation) startLoop(initialState initialWorldState, settings simSettings) {
	//---- Create all the communication channels

	// Make the 2way channels that will be used to make
	// add and remove actor requests to the go routine game loop
	addCh := make(chan addActorReq)
	removeCh := make(chan removeActorReq)

	// Set the 1way send chanels used by the public api
	s.addActor = addCh
	s.removeActor = removeCh

	// Set the 1way recieve channels used by the game loop
	var addReq <-chan addActorReq
	var removeReq <-chan removeActorReq

	addReq = addCh
	removeReq = removeCh

	// Make channel to be used to by the public api to
	// request that the simulation be halted
	haltCh := make(chan chan<- struct{})

	// Set the 1way send channel used by the public api
	s.requestHalt = haltCh

	// Set the 1way recieve channel used by the game loop
	var haltReq <-chan chan<- struct{}
	haltReq = haltCh

	quadTree := initialState.quadTree
	clock := stime.Clock(initialState.now)

	//---- User provided actor to entity resolver
	entityResolver := settings.EntityResolver

	//---- User provided input application phase
	inputPhase := settings.InputPhaseHandler

	//---- User provided narrow phase
	narrowPhase := settings.NarrowPhaseHandler

	runFrame := func(q quad.Quad, t stime.Time) quad.Quad {
		return quad.RunPhasesOn(q, inputPhase, narrowPhase, t)
	}

	// Start the Clock
	ticker := time.NewTicker(time.Duration(1000/settings.fps) * time.Millisecond)

	// Start the simulation server
	go func() {
		var hasHalted chan<- struct{}

	gameLoop:
		for {
			// Prioritized select for ticker.C and haltReq
			select {
			case <-ticker.C:
				goto tick

			case hasHalted = <-haltReq:
				break gameLoop
			default:
			}

			select {
			case <-ticker.C:
				goto tick

			case actor := <-addReq:
				// a is the new sim.Actor{} to be inserted into the sim
				a := <-actor.toBeAdded

				e := entityResolver.EntityForActor(a)
				quadTree = quadTree.Insert(e)

				// signal that the operation was a success
				actor.wasAdded <- a

			case actor := <-removeReq:
				// a is the new sim.Actor{} to be removed from the sim
				a := <-actor.toBeRemoved

				// TODO removed the actor from the simulation

				// signal that the operation was a success
				actor.wasRemoved <- a

			case hasHalted = <-haltReq:
				break gameLoop
			}

		tick:
			clock = clock.Tick()
			quadTree = runFrame(quadTree, clock.Now())
			// TODO quadTree to state
			// TODO send state
			continue
		}

		// We're done with cleanup and going to exit
		hasHalted <- struct{}{}
	}()
}

type haltedSimulation struct{}
<|endoftext|>


<commit_before>// The subnets package provides a subnet pool from which networks may be dynamically acquired or
// statically reserved.
package subnets

import (
	"fmt"
	"math"
	"net"
	"sync"

	"github.com/pivotal-golang/lager"
)

type Network struct {
	Subnet *net.IPNet
	IP     net.IP
}

//go:generate counterfeiter -o fake_subnet_pool/fake_pool.go . Pool
type Pool interface {
	// Allocates an IP address and associates it with a subnet. The subnet is selected by the given SubnetSelector.
	// The IP address is selected by the given IPSelector.
	// Returns a subnet, an IP address, and a boolean which is true if and only if this is the
	// first IP address to be associated with this subnet.
	// If either selector fails, an error is returned.
	Acquire(lager.Logger, SubnetSelector, IPSelector) (*net.IPNet, net.IP, error)

	// Releases an IP address associated with an allocated subnet. If the subnet has no other IP
	// addresses associated with it, it is deallocated.
	// Returns a boolean which is true if and only if the subnet was deallocated.
	// Returns an error if the given combination is not already in the pool.
	Release(*net.IPNet, net.IP) error

	// Remove an IP address so it appears to be associated with the given subnet.
	Remove(*net.IPNet, net.IP) error

	// Returns the number of /30 subnets which can be Acquired by a DynamicSubnetSelector.
	Capacity() int
}

type pool struct {
	allocated    map[string][]net.IP // net.IPNet.String +> seq net.IP
	dynamicRange *net.IPNet
	mu           sync.Mutex
}

//go:generate counterfeiter . SubnetSelector

// SubnetSelector is a strategy for selecting a subnet.
type SubnetSelector interface {
	// Returns a subnet based on a dynamic range and some existing statically-allocated
	// subnets. If no suitable subnet can be found, returns an error.
	SelectSubnet(dynamic *net.IPNet, existing []*net.IPNet) (*net.IPNet, error)
}

//go:generate counterfeiter . IPSelector

// IPSelector is a strategy for selecting an IP address in a subnet.
type IPSelector interface {
	// Returns an IP address in the given subnet which is not one of the given existing
	// IP addresses. If no such IP address can be found, returns an error.
	SelectIP(subnet *net.IPNet, existing []net.IP) (net.IP, error)
}

func NewPool(ipNet *net.IPNet) Pool {
	return &pool{dynamicRange: ipNet, allocated: make(map[string][]net.IP)}
}

// Acquire uses the given subnet and IP selectors to request a subnet, container IP address combination
// from the pool.
func (p *pool) Acquire(log lager.Logger, sn SubnetSelector, i IPSelector) (subnet *net.IPNet, ip net.IP, err error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if subnet, err = sn.SelectSubnet(p.dynamicRange, existingSubnets(p.allocated)); err != nil {
		return nil, nil, err
	}

	ips := p.allocated[subnet.String()]
	existingIPs := append(ips, NetworkIP(subnet), GatewayIP(subnet), BroadcastIP(subnet))
	if ip, err = i.SelectIP(subnet, existingIPs); err != nil {
		return nil, nil, err
	}

	p.allocated[subnet.String()] = append(ips, ip)
	return subnet, ip, err
}

// Recover re-allocates a given subnet and ip address combination in the pool. It returns
// an error if the combination is already allocated.
func (p *pool) Remove(subnet *net.IPNet, ip net.IP) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if ip == nil {
		return ErrIpCannotBeNil
	}

	for _, existing := range p.allocated[subnet.String()] {
		if existing.Equal(ip) {
			return ErrOverlapsExistingSubnet
		}
	}

	p.allocated[subnet.String()] = append(p.allocated[subnet.String()], ip)
	return nil
}

func (p *pool) Release(subnet *net.IPNet, ip net.IP) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	subnetString := subnet.String()
	ips := p.allocated[subnetString]

	if i, found := indexOf(ips, ip); found {
		if reducedIps, empty := removeIPAtIndex(ips, i); empty {
			delete(p.allocated, subnetString)
		} else {
			p.allocated[subnetString] = reducedIps
		}

		return nil
	}

	return ErrReleasedUnallocatedSubnet
}

// Capacity returns the number of /30 subnets that can be allocated
// from the pool's dynamic allocation range.
func (m *pool) Capacity() int {
	masked, total := m.dynamicRange.Mask.Size()
	return int(math.Pow(2, float64(total-masked)) / 4)
}

// Returns the gateway IP of a given subnet, which is always the maximum valid IP
func GatewayIP(subnet *net.IPNet) net.IP {
	return next(subnet.IP)
}

// Returns the network IP of a subnet.
func NetworkIP(subnet *net.IPNet) net.IP {
	return subnet.IP
}

// Returns the broadcast IP of a subnet.
func BroadcastIP(subnet *net.IPNet) net.IP {
	return max(subnet)
}

// returns the keys in the given map whose values are non-empty slices
func existingSubnets(m map[string][]net.IP) (result []*net.IPNet) {
	for k, v := range m {
		if len(v) > 0 {
			_, ipn, err := net.ParseCIDR(k)
			if err != nil {
				panic(fmt.Sprintf("failed to parse a CIDR in the subnet pool: %s", err))
			}

			result = append(result, ipn)
		}
	}

	return result
}

func indexOf(a []net.IP, w net.IP) (int, bool) {
	for i, v := range a {
		if v.Equal(w) {
			return i, true
		}
	}

	return -1, false
}

// removeAtIndex removes from a slice at the given index,
// and returns the new slice and boolean, true iff the new slice is empty.
func removeIPAtIndex(ips []net.IP, i int) ([]net.IP, bool) {
	l := len(ips)
	ips[i] = ips[l-1]
	ips = ips[:l-1]
	return ips, l == 1
}
<commit_msg>Remove unused Network data structure<commit_after>// The subnets package provides a subnet pool from which networks may be dynamically acquired or
// statically reserved.
package subnets

import (
	"fmt"
	"math"
	"net"
	"sync"

	"github.com/pivotal-golang/lager"
)

//go:generate counterfeiter -o fake_subnet_pool/fake_pool.go . Pool
type Pool interface {
	// Allocates an IP address and associates it with a subnet. The subnet is selected by the given SubnetSelector.
	// The IP address is selected by the given IPSelector.
	// Returns a subnet, an IP address, and a boolean which is true if and only if this is the
	// first IP address to be associated with this subnet.
	// If either selector fails, an error is returned.
	Acquire(lager.Logger, SubnetSelector, IPSelector) (*net.IPNet, net.IP, error)

	// Releases an IP address associated with an allocated subnet. If the subnet has no other IP
	// addresses associated with it, it is deallocated.
	// Returns a boolean which is true if and only if the subnet was deallocated.
	// Returns an error if the given combination is not already in the pool.
	Release(*net.IPNet, net.IP) error

	// Remove an IP address so it appears to be associated with the given subnet.
	Remove(*net.IPNet, net.IP) error

	// Returns the number of /30 subnets which can be Acquired by a DynamicSubnetSelector.
	Capacity() int
}

type pool struct {
	allocated    map[string][]net.IP // net.IPNet.String +> seq net.IP
	dynamicRange *net.IPNet
	mu           sync.Mutex
}

//go:generate counterfeiter . SubnetSelector

// SubnetSelector is a strategy for selecting a subnet.
type SubnetSelector interface {
	// Returns a subnet based on a dynamic range and some existing statically-allocated
	// subnets. If no suitable subnet can be found, returns an error.
	SelectSubnet(dynamic *net.IPNet, existing []*net.IPNet) (*net.IPNet, error)
}

//go:generate counterfeiter . IPSelector

// IPSelector is a strategy for selecting an IP address in a subnet.
type IPSelector interface {
	// Returns an IP address in the given subnet which is not one of the given existing
	// IP addresses. If no such IP address can be found, returns an error.
	SelectIP(subnet *net.IPNet, existing []net.IP) (net.IP, error)
}

func NewPool(ipNet *net.IPNet) Pool {
	return &pool{dynamicRange: ipNet, allocated: make(map[string][]net.IP)}
}

// Acquire uses the given subnet and IP selectors to request a subnet, container IP address combination
// from the pool.
func (p *pool) Acquire(log lager.Logger, sn SubnetSelector, i IPSelector) (subnet *net.IPNet, ip net.IP, err error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if subnet, err = sn.SelectSubnet(p.dynamicRange, existingSubnets(p.allocated)); err != nil {
		return nil, nil, err
	}

	ips := p.allocated[subnet.String()]
	existingIPs := append(ips, NetworkIP(subnet), GatewayIP(subnet), BroadcastIP(subnet))
	if ip, err = i.SelectIP(subnet, existingIPs); err != nil {
		return nil, nil, err
	}

	p.allocated[subnet.String()] = append(ips, ip)
	return subnet, ip, err
}

// Recover re-allocates a given subnet and ip address combination in the pool. It returns
// an error if the combination is already allocated.
func (p *pool) Remove(subnet *net.IPNet, ip net.IP) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if ip == nil {
		return ErrIpCannotBeNil
	}

	for _, existing := range p.allocated[subnet.String()] {
		if existing.Equal(ip) {
			return ErrOverlapsExistingSubnet
		}
	}

	p.allocated[subnet.String()] = append(p.allocated[subnet.String()], ip)
	return nil
}

func (p *pool) Release(subnet *net.IPNet, ip net.IP) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	subnetString := subnet.String()
	ips := p.allocated[subnetString]

	if i, found := indexOf(ips, ip); found {
		if reducedIps, empty := removeIPAtIndex(ips, i); empty {
			delete(p.allocated, subnetString)
		} else {
			p.allocated[subnetString] = reducedIps
		}

		return nil
	}

	return ErrReleasedUnallocatedSubnet
}

// Capacity returns the number of /30 subnets that can be allocated
// from the pool's dynamic allocation range.
func (m *pool) Capacity() int {
	masked, total := m.dynamicRange.Mask.Size()
	return int(math.Pow(2, float64(total-masked)) / 4)
}

// Returns the gateway IP of a given subnet, which is always the maximum valid IP
func GatewayIP(subnet *net.IPNet) net.IP {
	return next(subnet.IP)
}

// Returns the network IP of a subnet.
func NetworkIP(subnet *net.IPNet) net.IP {
	return subnet.IP
}

// Returns the broadcast IP of a subnet.
func BroadcastIP(subnet *net.IPNet) net.IP {
	return max(subnet)
}

// returns the keys in the given map whose values are non-empty slices
func existingSubnets(m map[string][]net.IP) (result []*net.IPNet) {
	for k, v := range m {
		if len(v) > 0 {
			_, ipn, err := net.ParseCIDR(k)
			if err != nil {
				panic(fmt.Sprintf("failed to parse a CIDR in the subnet pool: %s", err))
			}

			result = append(result, ipn)
		}
	}

	return result
}

func indexOf(a []net.IP, w net.IP) (int, bool) {
	for i, v := range a {
		if v.Equal(w) {
			return i, true
		}
	}

	return -1, false
}

// removeAtIndex removes from a slice at the given index,
// and returns the new slice and boolean, true iff the new slice is empty.
func removeIPAtIndex(ips []net.IP, i int) ([]net.IP, bool) {
	l := len(ips)
	ips[i] = ips[l-1]
	ips = ips[:l-1]
	return ips, l == 1
}
<|endoftext|>


<commit_before>package run

import (
	"github.com/workanator/go-floc.v2"
	"unidata/lib.coflow/coflow/flow"
)

const locWhile = "While"

/*
While repeats running the job while the condition is met.

Summary:
	- Run jobs in goroutines : NO
	- Wait all jobs finish   : YES
	- Run order              : SEQUENCE

Diagram:
                    YES
    +-------[JOB]<------+
    |                   |
    V                   | NO
  ----(CONDITION MET?)--+---->
*/
func While(predicate floc.Predicate, job floc.Job) floc.Job {
	return func(ctx floc.Context, ctrl floc.Control) error {
		for predicate(ctx) && !ctrl.IsFinished() {
			err := job(ctx, ctrl)
			if handledErr := handleResult(ctrl, err, locWhile); handledErr != nil {
				return handledErr
			}
		}

		return nil
	}
}
<commit_msg>Add While<commit_after>package run

import (
	"github.com/workanator/go-floc.v2"
)

const locWhile = "While"

/*
While repeats running the job while the condition is met.

Summary:
	- Run jobs in goroutines : NO
	- Wait all jobs finish   : YES
	- Run order              : SEQUENCE

Diagram:
                    YES
    +-------[JOB]<------+
    |                   |
    V                   | NO
  ----(CONDITION MET?)--+---->
*/
func While(predicate floc.Predicate, job floc.Job) floc.Job {
	return func(ctx floc.Context, ctrl floc.Control) error {
		for predicate(ctx) && !ctrl.IsFinished() {
			err := job(ctx, ctrl)
			if handledErr := handleResult(ctrl, err, locWhile); handledErr != nil {
				return handledErr
			}
		}

		return nil
	}
}
<|endoftext|>


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

package selector

import (
	"sync"

	"github.com/sirupsen/logrus"

	"github.com/networkservicemesh/networkservicemesh/controlplane/pkg/apis/local/connection"
	"github.com/networkservicemesh/networkservicemesh/controlplane/pkg/apis/registry"
)

type matchSelector struct {
	sync.Mutex
	roundRobin Selector
}

// NewMatchSelector creates a new
func NewMatchSelector() Selector {
	return &matchSelector{
		roundRobin: NewRoundRobinSelector(),
	}
}

// isSubset checks if B is a subset of A. TODO: reconsider this as a part of "tools"
func isSubset(A, B map[string]string) bool {
	if len(A) < len(B) {
		return false
	}
	for k, v := range B {
		if A[k] != v {
			return false
		}
	}
	return true
}

func (m *matchSelector) matchEndpoint(nsLabels map[string]string, ns *registry.NetworkService, networkServiceEndpoints []*registry.NetworkServiceEndpoint) *registry.NetworkServiceEndpoint {
	logrus.Infof("Matching ednpoint for labels %v", nsLabels)
	//Iterate through the matches
	for _, match := range ns.GetMatches() {
		// All match source selector labels should be present in the requested labels map
		if !isSubset(nsLabels, match.GetSourceSelector()) {
			continue
		}

		nseCandidates := []*registry.NetworkServiceEndpoint{}
		// Check all Destinations in that match
		for _, destination := range match.GetRoutes() {
			// Each NSE should be matched against that destination
			for _, nse := range networkServiceEndpoints {
				if isSubset(nse.GetLabels(), destination.GetDestinationSelector()) {
					nseCandidates = append(nseCandidates, nse)
				}
			}
		}

		if len(nseCandidates) > 0 {
			// We found candidates. Use RoundRobin to select one
			return m.roundRobin.SelectEndpoint(nil, ns, nseCandidates)
		}
	}
	return nil
}

func (m *matchSelector) SelectEndpoint(requestConnection *connection.Connection, ns *registry.NetworkService, networkServiceEndpoints []*registry.NetworkServiceEndpoint) *registry.NetworkServiceEndpoint {
	logrus.Infof("Selecting endpoint for %s with %d matches.", requestConnection.GetNetworkService(), len(ns.GetMatches()))
	if len(ns.GetMatches()) == 0 {
		return m.roundRobin.SelectEndpoint(nil, ns, networkServiceEndpoints)
	}

	return m.matchEndpoint(requestConnection.GetLabels(), ns, networkServiceEndpoints)
}
<commit_msg>Fix typo in match_selector.go (#1534)<commit_after>// Copyright 2018 VMware, Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package selector

import (
	"sync"

	"github.com/sirupsen/logrus"

	"github.com/networkservicemesh/networkservicemesh/controlplane/pkg/apis/local/connection"
	"github.com/networkservicemesh/networkservicemesh/controlplane/pkg/apis/registry"
)

type matchSelector struct {
	sync.Mutex
	roundRobin Selector
}

// NewMatchSelector creates a new
func NewMatchSelector() Selector {
	return &matchSelector{
		roundRobin: NewRoundRobinSelector(),
	}
}

// isSubset checks if B is a subset of A. TODO: reconsider this as a part of "tools"
func isSubset(A, B map[string]string) bool {
	if len(A) < len(B) {
		return false
	}
	for k, v := range B {
		if A[k] != v {
			return false
		}
	}
	return true
}

func (m *matchSelector) matchEndpoint(nsLabels map[string]string, ns *registry.NetworkService, networkServiceEndpoints []*registry.NetworkServiceEndpoint) *registry.NetworkServiceEndpoint {
	logrus.Infof("Matching endpoint for labels %v", nsLabels)
	//Iterate through the matches
	for _, match := range ns.GetMatches() {
		// All match source selector labels should be present in the requested labels map
		if !isSubset(nsLabels, match.GetSourceSelector()) {
			continue
		}

		nseCandidates := []*registry.NetworkServiceEndpoint{}
		// Check all Destinations in that match
		for _, destination := range match.GetRoutes() {
			// Each NSE should be matched against that destination
			for _, nse := range networkServiceEndpoints {
				if isSubset(nse.GetLabels(), destination.GetDestinationSelector()) {
					nseCandidates = append(nseCandidates, nse)
				}
			}
		}

		if len(nseCandidates) > 0 {
			// We found candidates. Use RoundRobin to select one
			return m.roundRobin.SelectEndpoint(nil, ns, nseCandidates)
		}
	}
	return nil
}

func (m *matchSelector) SelectEndpoint(requestConnection *connection.Connection, ns *registry.NetworkService, networkServiceEndpoints []*registry.NetworkServiceEndpoint) *registry.NetworkServiceEndpoint {
	logrus.Infof("Selecting endpoint for %s with %d matches.", requestConnection.GetNetworkService(), len(ns.GetMatches()))
	if len(ns.GetMatches()) == 0 {
		return m.roundRobin.SelectEndpoint(nil, ns, networkServiceEndpoints)
	}

	return m.matchEndpoint(requestConnection.GetLabels(), ns, networkServiceEndpoints)
}
<|endoftext|>


<commit_before>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The suffixarray package implements substring search in logarithmic time
// using an in-memory suffix array.
//
// Example use:
//
//	// create index for some data
//	index := suffixarray.New(data)
//
//	// lookup byte slice s
//	offsets1 := index.Lookup(s, -1) // the list of all indices where s occurs in data
//	offsets2 := index.Lookup(s, 3)  // the list of at most 3 indices where s occurs in data
//
package suffixarray

import (
	"bytes"
	"container/vector"
	"sort"
)

// BUG(gri): For larger data (10MB) which contains very long (say 100000)
// contiguous sequences of identical bytes, index creation time will be extremely slow.

// TODO(gri): Use a more sophisticated algorithm to create the suffix array.

// Index implements a suffix array for fast substring search.
type Index struct {
	data []byte
	sa   []int // suffix array for data
}

// New creates a new Index for data.
// Index creation time is approximately O(N*log(N)) for N = len(data).
//
func New(data []byte) *Index {
	sa := make([]int, len(data))
	for i, _ := range sa {
		sa[i] = i
	}
	x := &Index{data, sa}
	sort.Sort((*index)(x))
	return x
}

func (x *Index) at(i int) []byte {
	return x.data[x.sa[i]:]
}

// Binary search according to "A Method of Programming", E.W. Dijkstra.
func (x *Index) search(s []byte) int {
	i, j := 0, len(x.sa)
	// i < j for non-empty x
	for i+1 < j {
		// 0 <= i < j <= len(x.sa) && (x.at(i) <= s < x.at(j) || (s is not in x))
		h := i + (j-i)/2 // i < h < j
		if bytes.Compare(x.at(h), s) <= 0 {
			i = h
		} else { // s < x.at(h)
			j = h
		}
	}
	// i+1 == j for non-empty x
	return i
}

// Lookup returns an unsorted list of at most n indices where the byte string s
// occurs in the indexed data. If n < 0, all occurrences are returned.
// The result is nil if s is empty, s is not found, or n == 0.
// Lookup time is O((log(N) + len(result))*len(s)) where N is the
// size of the indexed data.
//
func (x *Index) Lookup(s []byte, n int) []int {
	var res vector.IntVector

	if len(s) > 0 && n != 0 {
		// find matching suffix index i
		i := x.search(s)
		// x.at(i) <= s < x.at(i+1)

		// ignore the first suffix if it is < s
		if i < len(x.sa) && bytes.Compare(x.at(i), s) < 0 {
			i++
		}

		// collect the following suffixes with matching prefixes
		for (n < 0 || len(res) < n) && i < len(x.sa) && bytes.HasPrefix(x.at(i), s) {
			res.Push(x.sa[i])
			i++
		}
	}

	return res
}

// index is used to hide the sort.Interface
type index Index

func (x *index) Len() int           { return len(x.sa) }
func (x *index) Less(i, j int) bool { return bytes.Compare(x.at(i), x.at(j)) < 0 }
func (x *index) Swap(i, j int)      { x.sa[i], x.sa[j] = x.sa[j], x.sa[i] }
func (a *index) at(i int) []byte    { return a.data[a.sa[i]:] }
<commit_msg>index/suffixarray: use sort.Search<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The suffixarray package implements substring search in logarithmic time
// using an in-memory suffix array.
//
// Example use:
//
//	// create index for some data
//	index := suffixarray.New(data)
//
//	// lookup byte slice s
//	offsets1 := index.Lookup(s, -1) // the list of all indices where s occurs in data
//	offsets2 := index.Lookup(s, 3)  // the list of at most 3 indices where s occurs in data
//
package suffixarray

import (
	"bytes"
	"container/vector"
	"sort"
)

// BUG(gri): For larger data (10MB) which contains very long (say 100000)
// contiguous sequences of identical bytes, index creation time will be extremely slow.

// TODO(gri): Use a more sophisticated algorithm to create the suffix array.

// Index implements a suffix array for fast substring search.
type Index struct {
	data []byte
	sa   []int // suffix array for data
}

// New creates a new Index for data.
// Index creation time is approximately O(N*log(N)) for N = len(data).
//
func New(data []byte) *Index {
	sa := make([]int, len(data))
	for i, _ := range sa {
		sa[i] = i
	}
	x := &Index{data, sa}
	sort.Sort((*index)(x))
	return x
}

func (x *Index) at(i int) []byte {
	return x.data[x.sa[i]:]
}

func (x *Index) search(s []byte) int {
	return sort.Search(len(x.sa), func(i int) bool { return bytes.Compare(x.at(i), s) >= 0 })
}

// Lookup returns an unsorted list of at most n indices where the byte string s
// occurs in the indexed data. If n < 0, all occurrences are returned.
// The result is nil if s is empty, s is not found, or n == 0.
// Lookup time is O((log(N) + len(result))*len(s)) where N is the
// size of the indexed data.
//
func (x *Index) Lookup(s []byte, n int) []int {
	var res vector.IntVector

	if len(s) > 0 && n != 0 {
		// find matching suffix index i
		i := x.search(s)
		// x.at(i-1) < s <= x.at(i)

		// collect the following suffixes with matching prefixes
		for (n < 0 || len(res) < n) && i < len(x.sa) && bytes.HasPrefix(x.at(i), s) {
			res.Push(x.sa[i])
			i++
		}
	}

	return res
}

// index is used to hide the sort.Interface
type index Index

func (x *index) Len() int           { return len(x.sa) }
func (x *index) Less(i, j int) bool { return bytes.Compare(x.at(i), x.at(j)) < 0 }
func (x *index) Swap(i, j int)      { x.sa[i], x.sa[j] = x.sa[j], x.sa[i] }
func (a *index) at(i int) []byte    { return a.data[a.sa[i]:] }
<|endoftext|>


<commit_before>/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func isCousins(root *TreeNode, x int, y int) bool {
    if root == nil {
        return false
    }
    parentX, depthX := findNode(root, nil, 0, x)
    parentY, depthY := findNode(root, nil, 0, y)
    if parentX != parentY && depthX == depthY {
        return true
    }else {
        return false
    }

}

func findNode(root *TreeNode, parent *TreeNode, depth int, val int) (retp *TreeNode, retd int) {
    if root.Val == val {
        retp = parent
        retd = depth
    }else{
        if root.Left != nil {
            retp, retd = findNode(root.Left, root, depth+1, val)

        }
        if retp != nil {
            return retp, retd
        }
        if root.Right != nil {
            retp, retd = findNode(root.Right, root, depth+1, val)
        }
    }
    return retp, retd
}
<commit_msg>Cousins in Binary Tree<commit_after>/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func isCousins(root *TreeNode, x int, y int) bool {
    if root == nil {
        return false
    }
    parentX, depthX := findNode(root, nil, 0, x)
    parentY, depthY := findNode(root, nil, 0, y)
    if parentX != parentY && depthX == depthY {
        return true
    }else {
        return false
    }

}

func findNode(root *TreeNode, parent *TreeNode, depth int, val int) (retp *TreeNode, retd int) {
    if root.Val == val {
        retp = parent
        retd = depth
    }else{
        if root.Left != nil {
            retp, retd = findNode(root.Left, root, depth+1, val)

        }
        if retp != nil {
            return retp, retd
        }
        if root.Right != nil {
            retp, retd = findNode(root.Right, root, depth+1, val)
        }
    }
    return retp, retd
}

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func isCousins(root *TreeNode, x int, y int) bool {
    var depthx, depthy int
    var prex, prey *TreeNode
    dfs(root, nil, 0, x, y, &depthx, &depthy, &prex, &prey)
    if depthx == depthy && prex != prey {
        return true
    }
    return false
}

func dfs(root *TreeNode, pre *TreeNode, depth int, x int, y int, depthx *int, depthy *int, prex **TreeNode, prey **TreeNode) {
    if root == nil {
        return
    }
    if root.Val == x {
        *depthx = depth
        *prex = pre
        return
    }else if root.Val == y {
        *depthy = depth
        *prey = pre
        return
    }else {
        dfs(root.Left, root, depth+1, x, y, depthx, depthy, prex, prey)
        dfs(root.Right, root, depth+1, x, y, depthx, depthy, prex, prey)
    }
}
<|endoftext|>