<commit_before>// kontrolclient implements a kite.Client for interacting with Kontrol kite.
package kite

import (
	"errors"
	"fmt"
	"math/rand"
	"net/url"
	"os"
	"sync"
	"time"

	"github.com/koding/kite/dnode"
	"github.com/koding/kite/protocol"
)

const (
	kontrolRetryDuration = 10 * time.Second
	proxyRetryDuration   = 10 * time.Second

	// kontrolConnectTimeout is the timeout for connecting to Kontrol in
	// TellKontrol-like methods.
	kontrolConnectTimeout = 10 * time.Second
)

// Returned from GetKites when query matches no kites.
var ErrNoKitesAvailable = errors.New("no kites availabile")

// kontrolClient is a kite for registering and querying Kites from Kontrol.
type kontrolClient struct {
	*Client
	sync.Mutex // protects Client

	// used for synchronizing methods that needs to be called after
	// successful connection or/and registiration to kontrol.
	onceConnected   sync.Once
	onceRegistered  sync.Once
	readyConnected  chan struct{}
	readyRegistered chan struct{}

	// lastRegisteredURL stores the Kite url what was send/registered
	// succesfully to kontrol
	lastRegisteredURL *url.URL

	// registerChan registers the url's it receives from the channel to Kontrol
	registerChan chan *url.URL
}

type registerResult struct {
	URL *url.URL
}

// SetupKontrolClient setups and prepares a the kontrol instance. It connects
// to kontrol and reconnects again if there is any disconnections. This method
// is called internally whenever a kontrol client specific action is taking.
// However if you wish to connect earlier you may call this method.
func (k *Kite) SetupKontrolClient() error {
	if k.kontrol.Client != nil {
		return nil // already prepared
	}

	if k.Config.KontrolURL == "" {
		return errors.New("no kontrol URL given in config")
	}

	client := k.NewClient(k.Config.KontrolURL)
	client.Kite = protocol.Kite{Name: "kontrol"} // for logging purposes
	client.Auth = &Auth{
		Type: "kiteKey",
		Key:  k.Config.KiteKey,
	}

	k.kontrol.Lock()
	k.kontrol.Client = client
	k.kontrol.Unlock()

	k.kontrol.OnConnect(func() {
		k.Log.Info("Connected to Kontrol ")

		// try to re-register on connect
		if k.kontrol.lastRegisteredURL != nil {
			select {
			case k.kontrol.registerChan <- k.kontrol.lastRegisteredURL:
			default:
			}
		}

		// signal all other methods that are listening on this channel, that we
		// are connected to kontrol.
		k.kontrol.onceConnected.Do(func() { close(k.kontrol.readyConnected) })
	})

	k.kontrol.OnDisconnect(func() {
		k.Log.Warning("Disconnected from Kontrol.")
	})

	// non blocking, is going to reconnect if the connection goes down.
	if _, err := k.kontrol.DialForever(); err != nil {
		return err
	}

	return nil
}

// GetKites returns the list of Kites matching the query. The returned list
// contains Ready to connect Client instances. The caller must connect
// with Client.Dial() before using each Kite. An error is returned when no
// kites are available.
func (k *Kite) GetKites(query *protocol.KontrolQuery) ([]*Client, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return nil, err
	}

	clients, err := k.getKites(protocol.GetKitesArgs{Query: query})
	if err != nil {
		return nil, err
	}

	if len(clients) == 0 {
		return nil, ErrNoKitesAvailable
	}

	return clients, nil
}

// used internally for GetKites() and WatchKites()
func (k *Kite) getKites(args protocol.GetKitesArgs) ([]*Client, error) {
	<-k.kontrol.readyConnected

	response, err := k.kontrol.TellWithTimeout("getKites", 4*time.Second, args)
	if err != nil {
		return nil, err
	}

	var result = new(protocol.GetKitesResult)
	err = response.Unmarshal(&result)
	if err != nil {
		return nil, err
	}

	clients := make([]*Client, len(result.Kites))
	for i, currentKite := range result.Kites {
		auth := &Auth{
			Type: "token",
			Key:  currentKite.Token,
		}

		clients[i] = k.NewClient(currentKite.URL)
		clients[i].Kite = currentKite.Kite
		clients[i].Auth = auth
	}

	// Renew tokens
	for _, r := range clients {
		token, err := NewTokenRenewer(r, k)
		if err != nil {
			k.Log.Error("Error in token. Token will not be renewed when it expires: %s", err.Error())
			continue
		}
		token.RenewWhenExpires()
	}

	return clients, nil
}

// GetToken is used to get a new token for a single Kite.
func (k *Kite) GetToken(kite *protocol.Kite) (string, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return "", err
	}

	<-k.kontrol.readyConnected

	result, err := k.kontrol.TellWithTimeout("getToken", 4*time.Second, kite)
	if err != nil {
		return "", err
	}

	var tkn string
	err = result.Unmarshal(&tkn)
	if err != nil {
		return "", err
	}

	return tkn, nil
}

// GetKey is used to get a new public key from kontrol if the current one is
// invalidated. The key is also replaced in memory and every request is going
// to use it. This means even if kite.key contains the old key, the kite itself
// uses the new one.
func (k *Kite) GetKey() (string, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return "", err
	}

	<-k.kontrol.readyConnected

	result, err := k.kontrol.TellWithTimeout("getKey", 4*time.Second)
	if err != nil {
		return "", err
	}

	var key string
	err = result.Unmarshal(&key)
	if err != nil {
		return "", err
	}

	k.Config.KontrolKey = key
	return key, nil
}

// NewKeyRenewer renews the internal key every given interval
func (k *Kite) NewKeyRenewer(interval time.Duration) {
	ticker := time.NewTicker(interval)
	for _ = range ticker.C {
		_, err := k.GetKey()
		if err != nil {
			k.Log.Warning("Key renew failed: %s", err)
		}
	}
}

// KontrolReadyNotify returns a channel that is closed when a successful
// registiration to kontrol is done.
func (k *Kite) KontrolReadyNotify() chan struct{} {
	return k.kontrol.readyRegistered
}

// signalReady is an internal method to notify that a sucessful registiration
// is done.
func (k *Kite) signalReady() {
	k.kontrol.onceRegistered.Do(func() { close(k.kontrol.readyRegistered) })
}

// RegisterForever is equilavent to Register(), but it tries to re-register if
// there is a disconnection. The returned error is for the first register
// attempt. It returns nil if ReadNotify() is ready and it's registered
// succesfull.
func (k *Kite) RegisterForever(kiteURL *url.URL) error {
	errs := make(chan error, 1)
	go func() {
		for u := range k.kontrol.registerChan {
			_, err := k.Register(u)
			if err == nil {
				k.kontrol.lastRegisteredURL = u
				k.signalReady()
				continue
			}

			select {
			case errs <- err:
			default:
			}

			k.Log.Error("Cannot register to Kontrol: %s Will retry after %d seconds",
				err, kontrolRetryDuration/time.Second)

			time.AfterFunc(kontrolRetryDuration, func() {
				select {
				case k.kontrol.registerChan <- u:
				default:
				}
			})
		}
	}()

	// don't block if there the given url is nil
	if kiteURL == nil {
		return nil
	}

	// initiate a registiration if a url is given
	k.kontrol.registerChan <- kiteURL

	select {
	case <-k.KontrolReadyNotify():
		return nil
	case err := <-errs:
		return err
	}
}

// Register registers current Kite to Kontrol. After registration other Kites
// can find it via GetKites() or WatchKites() method.  This method does not
// handle the reconnection case. If you want to keep registered to kontrol, use
// RegisterForever().
func (k *Kite) Register(kiteURL *url.URL) (*registerResult, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return nil, err
	}

	<-k.kontrol.readyConnected

	args := protocol.RegisterArgs{
		URL: kiteURL.String(),
	}

	k.Log.Info("Registering to kontrol with URL: %s", kiteURL.String())

	response, err := k.kontrol.TellWithTimeout("register", 4*time.Second, args)
	if err != nil {
		return nil, err
	}

	var rr protocol.RegisterResult
	err = response.Unmarshal(&rr)
	if err != nil {
		return nil, err
	}

	k.Log.Info("Registered to kontrol with URL: %s and Kite query: %s",
		rr.URL, k.Kite())

	parsed, err := url.Parse(rr.URL)
	if err != nil {
		k.Log.Error("Cannot parse registered URL: %s", err.Error())
	}

	// we also received a new public key (means the old one was invalidated).
	// Use it now.
	if rr.PublicKey != "" {
		k.Config.KontrolKey = rr.PublicKey
	}

	return &registerResult{parsed}, nil
}

// RegisterToTunnel finds a tunnel proxy kite by asking kontrol then registers
// itselfs on proxy. On error, retries forever. On every successfull
// registration, it sends the proxied URL to the registerChan channel. There is
// no register URL needed because the Tunnel Proxy automatically gets the IP
// from tunneling. This is a blocking function.
func (k *Kite) RegisterToTunnel() {
	query := &protocol.KontrolQuery{
		Username:    k.Config.KontrolUser,
		Environment: k.Config.Environment,
		Name:        "tunnelproxy",
	}

	k.RegisterToProxy(nil, query)
}

// RegisterToProxy is just like RegisterForever but registers the given URL
// to kontrol over a kite-proxy. A Kiteproxy is a reverseproxy that can be used
// for SSL termination or handling hundreds of kites behind a single. This is a
// blocking function.
func (k *Kite) RegisterToProxy(registerURL *url.URL, query *protocol.KontrolQuery) {
	go k.RegisterForever(nil)

	for {
		var proxyKite *Client

		// The proxy kite to connect can be overriden with the
		// environmental variable "KITE_PROXY_URL". If it is not set
		// we will ask Kontrol for available Proxy kites.
		// As an authentication informain kiteKey method will be used,
		// so be careful when using this feature.
		kiteProxyURL := os.Getenv("KITE_PROXY_URL")
		if kiteProxyURL != "" {
			proxyKite = k.NewClient(kiteProxyURL)
			proxyKite.Auth = &Auth{
				Type: "kiteKey",
				Key:  k.Config.KiteKey,
			}
		} else {
			kites, err := k.GetKites(query)
			if err != nil {
				k.Log.Error("Cannot get Proxy kites from Kontrol: %s", err.Error())
				time.Sleep(proxyRetryDuration)
				continue
			}

			// If more than one one Proxy Kite is available pick one randomly.
			// It does not matter which one we connect.
			proxyKite = kites[rand.Int()%len(kites)]
		}

		// Notify us on disconnect
		disconnect := make(chan bool, 1)
		proxyKite.OnDisconnect(func() {
			select {
			case disconnect <- true:
			default:
			}
		})

		proxyURL, err := k.registerToProxyKite(proxyKite, registerURL)
		if err != nil {
			time.Sleep(proxyRetryDuration)
			continue
		}

		k.kontrol.registerChan <- proxyURL

		// Block until disconnect from Proxy Kite.
		<-disconnect
	}
}

// registerToProxyKite dials the proxy kite and calls register method then
// returns the reverse-proxy URL.
func (k *Kite) registerToProxyKite(c *Client, kiteURL *url.URL) (*url.URL, error) {
	err := c.Dial()
	if err != nil {
		k.Log.Error("Cannot connect to Proxy kite: %s", err.Error())
		return nil, err
	}

	// Disconnect from Proxy Kite if error happens while registering.
	defer func() {
		if err != nil {
			c.Close()
		}
	}()

	// do not panic if we call Tell method below
	if kiteURL == nil {
		kiteURL = &url.URL{}
	}

	// this could be tunnelproxy or reverseproxy. Tunnelproxy doesn't need an
	// URL however Reverseproxy needs one.
	result, err := c.TellWithTimeout("register", 4*time.Second, kiteURL.String())
	if err != nil {
		k.Log.Error("Proxy register error: %s", err.Error())
		return nil, err
	}

	proxyURL, err := result.String()
	if err != nil {
		k.Log.Error("Proxy register result error: %s", err.Error())
		return nil, err
	}

	parsed, err := url.Parse(proxyURL)
	if err != nil {
		k.Log.Error("Cannot parse Proxy URL: %s", err.Error())
		return nil, err
	}

	return parsed, nil
}

// TellKontrolWithTimeout is a lower level function for communicating directly with
// kontrol. Like GetKites and GetToken, this automatically sets up and connects to
// kontrol as needed.
func (k *Kite) TellKontrolWithTimeout(method string, timeout time.Duration, args ...interface{}) (result *dnode.Partial, err error) {
	if err := k.SetupKontrolClient(); err != nil {
		return nil, err
	}

	// Wait for readyConnect, or timeout
	select {
	case <-time.After(kontrolConnectTimeout):
		return nil, &Error{
			Type: "timeout",
			Message: fmt.Sprintf(
				"Timed out registering to kontrol for %s method after %s",
				method, kontrolConnectTimeout,
			),
		}
	case <-k.kontrol.readyConnected:
	}

	return k.kontrol.TellWithTimeout(method, timeout, args...)
}
<commit_msg>kontrolclient: fix race on lastRegisteredURL<commit_after>// kontrolclient implements a kite.Client for interacting with Kontrol kite.
package kite

import (
	"errors"
	"fmt"
	"math/rand"
	"net/url"
	"os"
	"sync"
	"time"

	"github.com/koding/kite/dnode"
	"github.com/koding/kite/protocol"
)

const (
	kontrolRetryDuration = 10 * time.Second
	proxyRetryDuration   = 10 * time.Second

	// kontrolConnectTimeout is the timeout for connecting to Kontrol in
	// TellKontrol-like methods.
	kontrolConnectTimeout = 10 * time.Second
)

// Returned from GetKites when query matches no kites.
var ErrNoKitesAvailable = errors.New("no kites availabile")

// kontrolClient is a kite for registering and querying Kites from Kontrol.
type kontrolClient struct {
	*Client
	sync.Mutex // protects Client

	// used for synchronizing methods that needs to be called after
	// successful connection or/and registiration to kontrol.
	onceConnected   sync.Once
	onceRegistered  sync.Once
	readyConnected  chan struct{}
	readyRegistered chan struct{}

	// lastRegisteredURL stores the Kite url what was send/registered
	// succesfully to kontrol
	lastRegisteredURL *url.URL

	// registerChan registers the url's it receives from the channel to Kontrol
	registerChan chan *url.URL
}

type registerResult struct {
	URL *url.URL
}

// SetupKontrolClient setups and prepares a the kontrol instance. It connects
// to kontrol and reconnects again if there is any disconnections. This method
// is called internally whenever a kontrol client specific action is taking.
// However if you wish to connect earlier you may call this method.
func (k *Kite) SetupKontrolClient() error {
	if k.kontrol.Client != nil {
		return nil // already prepared
	}

	if k.Config.KontrolURL == "" {
		return errors.New("no kontrol URL given in config")
	}

	client := k.NewClient(k.Config.KontrolURL)
	client.Kite = protocol.Kite{Name: "kontrol"} // for logging purposes
	client.Auth = &Auth{
		Type: "kiteKey",
		Key:  k.Config.KiteKey,
	}

	k.kontrol.Lock()
	k.kontrol.Client = client
	k.kontrol.Unlock()

	k.kontrol.OnConnect(func() {
		k.Log.Info("Connected to Kontrol ")

		// try to re-register on connect
		k.kontrol.Lock()
		if k.kontrol.lastRegisteredURL != nil {
			select {
			case k.kontrol.registerChan <- k.kontrol.lastRegisteredURL:
			default:
			}
		}
		k.kontrol.Unlock()

		// signal all other methods that are listening on this channel, that we
		// are connected to kontrol.
		k.kontrol.onceConnected.Do(func() { close(k.kontrol.readyConnected) })
	})

	k.kontrol.OnDisconnect(func() {
		k.Log.Warning("Disconnected from Kontrol.")
	})

	// non blocking, is going to reconnect if the connection goes down.
	if _, err := k.kontrol.DialForever(); err != nil {
		return err
	}

	return nil
}

// GetKites returns the list of Kites matching the query. The returned list
// contains Ready to connect Client instances. The caller must connect
// with Client.Dial() before using each Kite. An error is returned when no
// kites are available.
func (k *Kite) GetKites(query *protocol.KontrolQuery) ([]*Client, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return nil, err
	}

	clients, err := k.getKites(protocol.GetKitesArgs{Query: query})
	if err != nil {
		return nil, err
	}

	if len(clients) == 0 {
		return nil, ErrNoKitesAvailable
	}

	return clients, nil
}

// used internally for GetKites() and WatchKites()
func (k *Kite) getKites(args protocol.GetKitesArgs) ([]*Client, error) {
	<-k.kontrol.readyConnected

	response, err := k.kontrol.TellWithTimeout("getKites", 4*time.Second, args)
	if err != nil {
		return nil, err
	}

	var result = new(protocol.GetKitesResult)
	err = response.Unmarshal(&result)
	if err != nil {
		return nil, err
	}

	clients := make([]*Client, len(result.Kites))
	for i, currentKite := range result.Kites {
		auth := &Auth{
			Type: "token",
			Key:  currentKite.Token,
		}

		clients[i] = k.NewClient(currentKite.URL)
		clients[i].Kite = currentKite.Kite
		clients[i].Auth = auth
	}

	// Renew tokens
	for _, r := range clients {
		token, err := NewTokenRenewer(r, k)
		if err != nil {
			k.Log.Error("Error in token. Token will not be renewed when it expires: %s", err.Error())
			continue
		}
		token.RenewWhenExpires()
	}

	return clients, nil
}

// GetToken is used to get a new token for a single Kite.
func (k *Kite) GetToken(kite *protocol.Kite) (string, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return "", err
	}

	<-k.kontrol.readyConnected

	result, err := k.kontrol.TellWithTimeout("getToken", 4*time.Second, kite)
	if err != nil {
		return "", err
	}

	var tkn string
	err = result.Unmarshal(&tkn)
	if err != nil {
		return "", err
	}

	return tkn, nil
}

// GetKey is used to get a new public key from kontrol if the current one is
// invalidated. The key is also replaced in memory and every request is going
// to use it. This means even if kite.key contains the old key, the kite itself
// uses the new one.
func (k *Kite) GetKey() (string, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return "", err
	}

	<-k.kontrol.readyConnected

	result, err := k.kontrol.TellWithTimeout("getKey", 4*time.Second)
	if err != nil {
		return "", err
	}

	var key string
	err = result.Unmarshal(&key)
	if err != nil {
		return "", err
	}

	k.Config.KontrolKey = key
	return key, nil
}

// NewKeyRenewer renews the internal key every given interval
func (k *Kite) NewKeyRenewer(interval time.Duration) {
	ticker := time.NewTicker(interval)
	for _ = range ticker.C {
		_, err := k.GetKey()
		if err != nil {
			k.Log.Warning("Key renew failed: %s", err)
		}
	}
}

// KontrolReadyNotify returns a channel that is closed when a successful
// registiration to kontrol is done.
func (k *Kite) KontrolReadyNotify() chan struct{} {
	return k.kontrol.readyRegistered
}

// signalReady is an internal method to notify that a sucessful registiration
// is done.
func (k *Kite) signalReady() {
	k.kontrol.onceRegistered.Do(func() { close(k.kontrol.readyRegistered) })
}

// RegisterForever is equilavent to Register(), but it tries to re-register if
// there is a disconnection. The returned error is for the first register
// attempt. It returns nil if ReadNotify() is ready and it's registered
// succesfull.
func (k *Kite) RegisterForever(kiteURL *url.URL) error {
	errs := make(chan error, 1)
	go func() {
		for u := range k.kontrol.registerChan {
			_, err := k.Register(u)
			if err == nil {
				k.kontrol.Lock()
				k.kontrol.lastRegisteredURL = u
				k.kontrol.Unlock()
				k.signalReady()
				continue
			}

			select {
			case errs <- err:
			default:
			}

			k.Log.Error("Cannot register to Kontrol: %s Will retry after %d seconds",
				err, kontrolRetryDuration/time.Second)

			time.AfterFunc(kontrolRetryDuration, func() {
				select {
				case k.kontrol.registerChan <- u:
				default:
				}
			})
		}
	}()

	// don't block if there the given url is nil
	if kiteURL == nil {
		return nil
	}

	// initiate a registiration if a url is given
	k.kontrol.registerChan <- kiteURL

	select {
	case <-k.KontrolReadyNotify():
		return nil
	case err := <-errs:
		return err
	}
}

// Register registers current Kite to Kontrol. After registration other Kites
// can find it via GetKites() or WatchKites() method.  This method does not
// handle the reconnection case. If you want to keep registered to kontrol, use
// RegisterForever().
func (k *Kite) Register(kiteURL *url.URL) (*registerResult, error) {
	if err := k.SetupKontrolClient(); err != nil {
		return nil, err
	}

	<-k.kontrol.readyConnected

	args := protocol.RegisterArgs{
		URL: kiteURL.String(),
	}

	k.Log.Info("Registering to kontrol with URL: %s", kiteURL.String())

	response, err := k.kontrol.TellWithTimeout("register", 4*time.Second, args)
	if err != nil {
		return nil, err
	}

	var rr protocol.RegisterResult
	err = response.Unmarshal(&rr)
	if err != nil {
		return nil, err
	}

	k.Log.Info("Registered to kontrol with URL: %s and Kite query: %s",
		rr.URL, k.Kite())

	parsed, err := url.Parse(rr.URL)
	if err != nil {
		k.Log.Error("Cannot parse registered URL: %s", err.Error())
	}

	// we also received a new public key (means the old one was invalidated).
	// Use it now.
	if rr.PublicKey != "" {
		k.Config.KontrolKey = rr.PublicKey
	}

	return &registerResult{parsed}, nil
}

// RegisterToTunnel finds a tunnel proxy kite by asking kontrol then registers
// itselfs on proxy. On error, retries forever. On every successfull
// registration, it sends the proxied URL to the registerChan channel. There is
// no register URL needed because the Tunnel Proxy automatically gets the IP
// from tunneling. This is a blocking function.
func (k *Kite) RegisterToTunnel() {
	query := &protocol.KontrolQuery{
		Username:    k.Config.KontrolUser,
		Environment: k.Config.Environment,
		Name:        "tunnelproxy",
	}

	k.RegisterToProxy(nil, query)
}

// RegisterToProxy is just like RegisterForever but registers the given URL
// to kontrol over a kite-proxy. A Kiteproxy is a reverseproxy that can be used
// for SSL termination or handling hundreds of kites behind a single. This is a
// blocking function.
func (k *Kite) RegisterToProxy(registerURL *url.URL, query *protocol.KontrolQuery) {
	go k.RegisterForever(nil)

	for {
		var proxyKite *Client

		// The proxy kite to connect can be overriden with the
		// environmental variable "KITE_PROXY_URL". If it is not set
		// we will ask Kontrol for available Proxy kites.
		// As an authentication informain kiteKey method will be used,
		// so be careful when using this feature.
		kiteProxyURL := os.Getenv("KITE_PROXY_URL")
		if kiteProxyURL != "" {
			proxyKite = k.NewClient(kiteProxyURL)
			proxyKite.Auth = &Auth{
				Type: "kiteKey",
				Key:  k.Config.KiteKey,
			}
		} else {
			kites, err := k.GetKites(query)
			if err != nil {
				k.Log.Error("Cannot get Proxy kites from Kontrol: %s", err.Error())
				time.Sleep(proxyRetryDuration)
				continue
			}

			// If more than one one Proxy Kite is available pick one randomly.
			// It does not matter which one we connect.
			proxyKite = kites[rand.Int()%len(kites)]
		}

		// Notify us on disconnect
		disconnect := make(chan bool, 1)
		proxyKite.OnDisconnect(func() {
			select {
			case disconnect <- true:
			default:
			}
		})

		proxyURL, err := k.registerToProxyKite(proxyKite, registerURL)
		if err != nil {
			time.Sleep(proxyRetryDuration)
			continue
		}

		k.kontrol.registerChan <- proxyURL

		// Block until disconnect from Proxy Kite.
		<-disconnect
	}
}

// registerToProxyKite dials the proxy kite and calls register method then
// returns the reverse-proxy URL.
func (k *Kite) registerToProxyKite(c *Client, kiteURL *url.URL) (*url.URL, error) {
	err := c.Dial()
	if err != nil {
		k.Log.Error("Cannot connect to Proxy kite: %s", err.Error())
		return nil, err
	}

	// Disconnect from Proxy Kite if error happens while registering.
	defer func() {
		if err != nil {
			c.Close()
		}
	}()

	// do not panic if we call Tell method below
	if kiteURL == nil {
		kiteURL = &url.URL{}
	}

	// this could be tunnelproxy or reverseproxy. Tunnelproxy doesn't need an
	// URL however Reverseproxy needs one.
	result, err := c.TellWithTimeout("register", 4*time.Second, kiteURL.String())
	if err != nil {
		k.Log.Error("Proxy register error: %s", err.Error())
		return nil, err
	}

	proxyURL, err := result.String()
	if err != nil {
		k.Log.Error("Proxy register result error: %s", err.Error())
		return nil, err
	}

	parsed, err := url.Parse(proxyURL)
	if err != nil {
		k.Log.Error("Cannot parse Proxy URL: %s", err.Error())
		return nil, err
	}

	return parsed, nil
}

// TellKontrolWithTimeout is a lower level function for communicating directly with
// kontrol. Like GetKites and GetToken, this automatically sets up and connects to
// kontrol as needed.
func (k *Kite) TellKontrolWithTimeout(method string, timeout time.Duration, args ...interface{}) (result *dnode.Partial, err error) {
	if err := k.SetupKontrolClient(); err != nil {
		return nil, err
	}

	// Wait for readyConnect, or timeout
	select {
	case <-time.After(kontrolConnectTimeout):
		return nil, &Error{
			Type: "timeout",
			Message: fmt.Sprintf(
				"Timed out registering to kontrol for %s method after %s",
				method, kontrolConnectTimeout,
			),
		}
	case <-k.kontrol.readyConnected:
	}

	return k.kontrol.TellWithTimeout(method, timeout, args...)
}
<|endoftext|>


<commit_before>package kuzzle

import (
  "github.com/kuzzleio/sdk-go/connection"
  "github.com/kuzzleio/sdk-go/types"
  "encoding/json"
  "github.com/satori/go.uuid"
  "sync"
)

type IKuzzle interface {
  Query(types.KuzzleRequest, chan<- types.KuzzleResponse, *types.Options)
}

type Kuzzle struct {
  Host   string
  socket connection.Connection
  State  *int

  wasConnected bool
  lastUrl      string
  message      chan []byte
  mu           *sync.Mutex
  jwtToken     string
}

// Kuzzle constructor
func NewKuzzle(c connection.Connection, options *types.Options) (*Kuzzle, error) {
  var err error

  if options == nil {
    options = types.DefaultOptions()
  }

  k := &Kuzzle{
    mu:     &sync.Mutex{},
    socket: c,
  }

  k.State = k.socket.GetState()

  if options.Connect == types.Auto {
    err = k.Connect()
  }

  return k, err
}

// Connects to a Kuzzle instance using the provided host and port.
func (k *Kuzzle) Connect() error {
  wasConnected, err := k.socket.Connect()
  if err == nil {
    //if k.lastUrl != k.Host {
    //  k.wasConnected = false
    //  k.lastUrl = k.Host
    //}

    if wasConnected {
      if k.jwtToken != "" {
        // todo avoid import cycle (kuzzle)
        //go func() {
        //	res, err := kuzzle.CheckToken(k, k.jwtToken)
        //
        //	if err != nil {
        //		k.jwtToken = ""
        //		k.emitEvent(event.JwtTokenExpired, nil)
        //		k.Reconnect()
        //		return
        //	}
        //
        //	if !res.Valid {
        //		k.jwtToken = ""
        //		k.emitEvent(event.JwtTokenExpired, nil)
        //	}
        //	k.Reconnect()
        //}()
      }
    }
    return nil
  }

  return err
}

// This is a low-level method, exposed to allow advanced SDK users to bypass high-level methods.
func (k Kuzzle) Query(query types.KuzzleRequest, options *types.Options, responseChannel chan<- types.KuzzleResponse) {
  requestId := uuid.NewV4().String()

  query.RequestId = requestId

  type body struct{}

  if query.Body == nil {
    query.Body = make(map[string]interface{})
  }

  jsonRequest, err := json.Marshal(query)
  if err != nil {
    responseChannel <- types.KuzzleResponse{Error: types.MessageError{Message: err.Error()}}
    return
  }

  err = k.socket.Send(jsonRequest, options, responseChannel, requestId)
  if err != nil {
    responseChannel <- types.KuzzleResponse{Error: types.MessageError{Message: err.Error()}}
    return
  }
}

func (k Kuzzle) GetOfflineQueue() *[]types.QueryObject {
  return k.socket.GetOfflineQueue()
}

// Disconnect from Kuzzle and invalidate this instance.
// Does not fire a disconnected event.
func (k *Kuzzle) Disconnect() error {
  err := k.socket.Close()

  if err != nil {
    return err
  }
  k.wasConnected = false

  return nil
}
<commit_msg>remove disconnect method from kuzzle file<commit_after>package kuzzle

import (
  "github.com/kuzzleio/sdk-go/connection"
  "github.com/kuzzleio/sdk-go/types"
  "encoding/json"
  "github.com/satori/go.uuid"
  "sync"
)

type IKuzzle interface {
  Query(types.KuzzleRequest, chan<- types.KuzzleResponse, *types.Options)
}

type Kuzzle struct {
  Host   string
  socket connection.Connection
  State  *int

  wasConnected bool
  lastUrl      string
  message      chan []byte
  mu           *sync.Mutex
  jwtToken     string
}

// Kuzzle constructor
func NewKuzzle(c connection.Connection, options *types.Options) (*Kuzzle, error) {
  var err error

  if options == nil {
    options = types.DefaultOptions()
  }

  k := &Kuzzle{
    mu:     &sync.Mutex{},
    socket: c,
  }

  k.State = k.socket.GetState()

  if options.Connect == types.Auto {
    err = k.Connect()
  }

  return k, err
}

// Connects to a Kuzzle instance using the provided host and port.
func (k *Kuzzle) Connect() error {
  wasConnected, err := k.socket.Connect()
  if err == nil {
    //if k.lastUrl != k.Host {
    //  k.wasConnected = false
    //  k.lastUrl = k.Host
    //}

    if wasConnected {
      if k.jwtToken != "" {
        // todo avoid import cycle (kuzzle)
        //go func() {
        //	res, err := kuzzle.CheckToken(k, k.jwtToken)
        //
        //	if err != nil {
        //		k.jwtToken = ""
        //		k.emitEvent(event.JwtTokenExpired, nil)
        //		k.Reconnect()
        //		return
        //	}
        //
        //	if !res.Valid {
        //		k.jwtToken = ""
        //		k.emitEvent(event.JwtTokenExpired, nil)
        //	}
        //	k.Reconnect()
        //}()
      }
    }
    return nil
  }

  return err
}

// This is a low-level method, exposed to allow advanced SDK users to bypass high-level methods.
func (k Kuzzle) Query(query types.KuzzleRequest, options *types.Options, responseChannel chan<- types.KuzzleResponse) {
  requestId := uuid.NewV4().String()

  query.RequestId = requestId

  type body struct{}

  if query.Body == nil {
    query.Body = make(map[string]interface{})
  }

  jsonRequest, err := json.Marshal(query)
  if err != nil {
    responseChannel <- types.KuzzleResponse{Error: types.MessageError{Message: err.Error()}}
    return
  }

  err = k.socket.Send(jsonRequest, options, responseChannel, requestId)
  if err != nil {
    responseChannel <- types.KuzzleResponse{Error: types.MessageError{Message: err.Error()}}
    return
  }
}

func (k Kuzzle) GetOfflineQueue() *[]types.QueryObject {
  return k.socket.GetOfflineQueue()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gophergala2016/goad/Godeps/_workspace/src/github.com/aws/aws-sdk-go/aws"
	"github.com/gophergala2016/goad/helpers"
	"github.com/gophergala2016/goad/queue"
)

func main() {

	var (
		address          string
		sqsurl           string
		concurrencycount int
		maxRequestCount  int
		timeout          string
		frequency        string
		awsregion        string
		queueRegion      string
		requestMethod    string
		requestBody      string
		requestHeaders   helpers.StringsliceFlag
	)

	flag.StringVar(&address, "u", "", "URL to load test (required)")
	flag.StringVar(&requestMethod, "m", "GET", "HTTP method")
	flag.StringVar(&requestBody, "b", "", "HTTP request body")
	flag.StringVar(&awsregion, "r", "", "AWS region to run in")
	flag.StringVar(&queueRegion, "q", "", "Queue region")
	flag.StringVar(&sqsurl, "s", "", "sqsUrl")
	flag.StringVar(&timeout, "t", "15s", "request timeout in seconds")
	flag.StringVar(&frequency, "f", "15s", "Reporting frequency in seconds")

	flag.IntVar(&concurrencycount, "c", 10, "number of concurrent requests")
	flag.IntVar(&maxRequestCount, "n", 1000, "number of total requests to make")

	flag.Var(&requestHeaders, "H", "List of headers")
	flag.Parse()

	clientTimeout, _ := time.ParseDuration(timeout)
	fmt.Printf("Using a timeout of %s\n", clientTimeout)
	reportingFrequency, _ := time.ParseDuration(frequency)
	fmt.Printf("Using a reporting frequency of %s\n", reportingFrequency)

	client := &http.Client{}
	client.Timeout = clientTimeout

	fmt.Printf("Will spawn %d workers making %d requests to %s\n", concurrencycount, maxRequestCount, address)
	runLoadTest(client, sqsurl, address, maxRequestCount, concurrencycount, awsregion, reportingFrequency, queueRegion, requestMethod, requestBody, requestHeaders)
}

type RequestResult struct {
	Time             int64  `json:"time"`
	Host             string `json:"host"`
	Type             string `json:"type"`
	Status           int    `json:"status"`
	ElapsedFirstByte int64  `json:"elapsed-first-byte"`
	ElapsedLastByte  int64  `json:"elapsed-last-byte"`
	Elapsed          int64  `json:"elapsed"`
	Bytes            int    `json:"bytes"`
	Timeout          bool   `json:"timeout"`
	ConnectionError  bool   `json:"connection-error"`
	State            string `json:"state"`
}

func runLoadTest(client *http.Client, sqsurl string, url string, totalRequests int, concurrencycount int, awsregion string, reportingFrequency time.Duration, queueRegion string, requestMethod string, requestBody string, requestHeaders []string) {
	awsConfig := aws.NewConfig().WithRegion(queueRegion)
	sqsAdaptor := queue.NewSQSAdaptor(awsConfig, sqsurl)
	//sqsAdaptor := queue.NewDummyAdaptor(sqsurl)
	jobs := make(chan struct{}, totalRequests)
	ch := make(chan RequestResult, totalRequests)
	var wg sync.WaitGroup
	loadTestStartTime := time.Now()
	var requestsSoFar int
	for i := 0; i < totalRequests; i++ {
		jobs <- struct{}{}
	}
	close(jobs)
	fmt.Print("Spawning workers…")
	for i := 0; i < concurrencycount; i++ {
		wg.Add(1)
		go fetch(loadTestStartTime, client, url, totalRequests, jobs, ch, &wg, awsregion, requestMethod, requestBody, requestHeaders)
		fmt.Print(".")
	}
	fmt.Println(" done.\nWaiting for results…")

	ticker := time.NewTicker(reportingFrequency)
	quit := make(chan struct{})
	quitting := false

	for requestsSoFar < totalRequests && !quitting {
		i := 0

		var timeToFirstTotal int64
		var requestTimeTotal int64
		totBytesRead := 0
		statuses := make(map[string]int)
		var firstRequestTime int64
		var lastRequestTime int64
		var slowest int64
		var fastest int64
		var totalTimedOut int
		var totalConnectionError int

		resetStats := false
		for requestsSoFar < totalRequests && !quitting && !resetStats {
			aggregate := false
			select {
			case r := <-ch:
				i++
				requestsSoFar++
				if requestsSoFar%10 == 0 || requestsSoFar == totalRequests {
					fmt.Printf("\r%.2f%% done (%d requests out of %d)", (float64(requestsSoFar)/float64(totalRequests))*100.0, requestsSoFar, totalRequests)
				}
				if firstRequestTime == 0 {
					firstRequestTime = r.Time
				}

				lastRequestTime = r.Time

				if r.Timeout {
					totalTimedOut++
					continue
				}
				if r.ConnectionError {
					totalConnectionError++
					continue
				}

				if r.ElapsedLastByte > slowest {
					slowest = r.ElapsedLastByte
				}
				if fastest == 0 {
					fastest = r.ElapsedLastByte
				} else {
					if r.ElapsedLastByte < fastest {
						fastest = r.ElapsedLastByte
					}
				}

				timeToFirstTotal += r.ElapsedFirstByte
				totBytesRead += r.Bytes
				statusStr := strconv.Itoa(r.Status)
				_, ok := statuses[statusStr]
				if !ok {
					statuses[statusStr] = 1
				} else {
					statuses[statusStr]++
				}
				requestTimeTotal += r.Elapsed
				if requestsSoFar == totalRequests {
					quitting = true
				}
			case <-ticker.C:
				if i == 0 {
					continue
				}
				aggregate = true
			case <-quit:
				ticker.Stop()
				quitting = true
			}
			if aggregate || quitting {
				durationNanoSeconds := lastRequestTime - firstRequestTime
				durationSeconds := float32(durationNanoSeconds) / float32(1000000000)
				var reqPerSec float32
				var kbPerSec float32
				if durationSeconds > 0 {
					reqPerSec = float32(i) / durationSeconds
					kbPerSec = (float32(totBytesRead) / durationSeconds) / 1024.0
				} else {
					reqPerSec = 0
					kbPerSec = 0
				}

				fatalError := ""
				if (totalTimedOut + totalConnectionError) > i/2 {
					fatalError = "Over 50% of requests failed, aborting"
					quitting = true
				}
				aggData := queue.AggData{
					i,
					totalTimedOut,
					totalConnectionError,
					timeToFirstTotal / int64(i),
					totBytesRead,
					statuses,
					requestTimeTotal / int64(i),
					reqPerSec,
					kbPerSec,
					slowest,
					fastest,
					awsregion,
					fatalError,
				}
				sqsAdaptor.SendResult(aggData)
				resetStats = true
			}
		}
	}
	fmt.Printf("\nYay🎈  - %d requests completed\n", requestsSoFar)

}

func fetch(loadTestStartTime time.Time, client *http.Client, address string, requestcount int, jobs <-chan struct{}, ch chan RequestResult, wg *sync.WaitGroup, awsregion string, requestMethod string, requestBody string, requestHeaders []string) {
	defer wg.Done()
	for _ = range jobs {
		start := time.Now()
		req, err := http.NewRequest(requestMethod, address, bytes.NewBufferString(requestBody))
		if err != nil {
			fmt.Println("Error creating the HTTP request:", err)
		}
		req.Header.Add("User-Agent", "Mozilla/5.0 (compatible; Goad/1.0; +https://goad.io)")
		req.Header.Add("Accept-Encoding", "gzip")
		for _, v := range requestHeaders {
			header := strings.Split(v, ":")
			req.Header.Add(strings.Trim(header[0], " "), strings.Trim(header[1], " "))
		}

		response, err := client.Do(req)
		var status string
		var elapsedFirstByte time.Duration
		var elapsedLastByte time.Duration
		var elapsed time.Duration
		var statusCode int
		var bytesRead int
		buf := []byte(" ")
		timedOut := false
		connectionError := false
		isRedirect := err != nil && strings.Contains(err.Error(), "redirect")
		if err != nil && !isRedirect {
			status = fmt.Sprintf("ERROR: %s\n", err)
			switch err := err.(type) {
			case *url.Error:
				if err, ok := err.Err.(net.Error); ok && err.Timeout() {
					timedOut = true
				}
			case net.Error:
				if err.Timeout() {
					timedOut = true
				}
			}

			if !timedOut {
				connectionError = true
			}
		} else {
			statusCode = response.StatusCode
			elapsedFirstByte = time.Since(start)
			if !isRedirect {
				_, err = response.Body.Read(buf)
				firstByteRead := true
				if err != nil {
					status = fmt.Sprintf("reading first byte failed: %s\n", err)
					firstByteRead = false
				}
				body, err := ioutil.ReadAll(response.Body)
				if firstByteRead {
					bytesRead = len(body) + 1
				}
				elapsedLastByte = time.Since(start)
				if err != nil {
					// todo: detect timeout here as well
					status = fmt.Sprintf("reading response body failed: %s\n", err)
					connectionError = true
				} else {
					status = "Success"
				}
			} else {
				status = "Redirect"
			}
			response.Body.Close()

			elapsed = time.Since(start)
		}
		//fmt.Printf("Request end: %d, elapsed: %d\n", time.Now().Sub(loadTestStartTime).Nanoseconds(), elapsed.Nanoseconds())
		result := RequestResult{
			start.Sub(loadTestStartTime).Nanoseconds(),
			req.URL.Host,
			req.Method,
			statusCode,
			elapsed.Nanoseconds(),
			elapsedFirstByte.Nanoseconds(),
			elapsedLastByte.Nanoseconds(),
			bytesRead,
			timedOut,
			connectionError,
			status,
		}
		ch <- result
	}
}
<commit_msg>Don't proceed with request if an error occurs<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gophergala2016/goad/Godeps/_workspace/src/github.com/aws/aws-sdk-go/aws"
	"github.com/gophergala2016/goad/helpers"
	"github.com/gophergala2016/goad/queue"
)

func main() {

	var (
		address          string
		sqsurl           string
		concurrencycount int
		maxRequestCount  int
		timeout          string
		frequency        string
		awsregion        string
		queueRegion      string
		requestMethod    string
		requestBody      string
		requestHeaders   helpers.StringsliceFlag
	)

	flag.StringVar(&address, "u", "", "URL to load test (required)")
	flag.StringVar(&requestMethod, "m", "GET", "HTTP method")
	flag.StringVar(&requestBody, "b", "", "HTTP request body")
	flag.StringVar(&awsregion, "r", "", "AWS region to run in")
	flag.StringVar(&queueRegion, "q", "", "Queue region")
	flag.StringVar(&sqsurl, "s", "", "sqsUrl")
	flag.StringVar(&timeout, "t", "15s", "request timeout in seconds")
	flag.StringVar(&frequency, "f", "15s", "Reporting frequency in seconds")

	flag.IntVar(&concurrencycount, "c", 10, "number of concurrent requests")
	flag.IntVar(&maxRequestCount, "n", 1000, "number of total requests to make")

	flag.Var(&requestHeaders, "H", "List of headers")
	flag.Parse()

	clientTimeout, _ := time.ParseDuration(timeout)
	fmt.Printf("Using a timeout of %s\n", clientTimeout)
	reportingFrequency, _ := time.ParseDuration(frequency)
	fmt.Printf("Using a reporting frequency of %s\n", reportingFrequency)

	client := &http.Client{}
	client.Timeout = clientTimeout

	fmt.Printf("Will spawn %d workers making %d requests to %s\n", concurrencycount, maxRequestCount, address)
	runLoadTest(client, sqsurl, address, maxRequestCount, concurrencycount, awsregion, reportingFrequency, queueRegion, requestMethod, requestBody, requestHeaders)
}

type RequestResult struct {
	Time             int64  `json:"time"`
	Host             string `json:"host"`
	Type             string `json:"type"`
	Status           int    `json:"status"`
	ElapsedFirstByte int64  `json:"elapsed-first-byte"`
	ElapsedLastByte  int64  `json:"elapsed-last-byte"`
	Elapsed          int64  `json:"elapsed"`
	Bytes            int    `json:"bytes"`
	Timeout          bool   `json:"timeout"`
	ConnectionError  bool   `json:"connection-error"`
	State            string `json:"state"`
}

func runLoadTest(client *http.Client, sqsurl string, url string, totalRequests int, concurrencycount int, awsregion string, reportingFrequency time.Duration, queueRegion string, requestMethod string, requestBody string, requestHeaders []string) {
	awsConfig := aws.NewConfig().WithRegion(queueRegion)
	sqsAdaptor := queue.NewSQSAdaptor(awsConfig, sqsurl)
	//sqsAdaptor := queue.NewDummyAdaptor(sqsurl)
	jobs := make(chan struct{}, totalRequests)
	ch := make(chan RequestResult, totalRequests)
	var wg sync.WaitGroup
	loadTestStartTime := time.Now()
	var requestsSoFar int
	for i := 0; i < totalRequests; i++ {
		jobs <- struct{}{}
	}
	close(jobs)
	fmt.Print("Spawning workers…")
	for i := 0; i < concurrencycount; i++ {
		wg.Add(1)
		go fetch(loadTestStartTime, client, url, totalRequests, jobs, ch, &wg, awsregion, requestMethod, requestBody, requestHeaders)
		fmt.Print(".")
	}
	fmt.Println(" done.\nWaiting for results…")

	ticker := time.NewTicker(reportingFrequency)
	quit := make(chan struct{})
	quitting := false

	for requestsSoFar < totalRequests && !quitting {
		i := 0

		var timeToFirstTotal int64
		var requestTimeTotal int64
		totBytesRead := 0
		statuses := make(map[string]int)
		var firstRequestTime int64
		var lastRequestTime int64
		var slowest int64
		var fastest int64
		var totalTimedOut int
		var totalConnectionError int

		resetStats := false
		for requestsSoFar < totalRequests && !quitting && !resetStats {
			aggregate := false
			select {
			case r := <-ch:
				i++
				requestsSoFar++
				if requestsSoFar%10 == 0 || requestsSoFar == totalRequests {
					fmt.Printf("\r%.2f%% done (%d requests out of %d)", (float64(requestsSoFar)/float64(totalRequests))*100.0, requestsSoFar, totalRequests)
				}
				if firstRequestTime == 0 {
					firstRequestTime = r.Time
				}

				lastRequestTime = r.Time

				if r.Timeout {
					totalTimedOut++
					continue
				}
				if r.ConnectionError {
					totalConnectionError++
					continue
				}

				if r.ElapsedLastByte > slowest {
					slowest = r.ElapsedLastByte
				}
				if fastest == 0 {
					fastest = r.ElapsedLastByte
				} else {
					if r.ElapsedLastByte < fastest {
						fastest = r.ElapsedLastByte
					}
				}

				timeToFirstTotal += r.ElapsedFirstByte
				totBytesRead += r.Bytes
				statusStr := strconv.Itoa(r.Status)
				_, ok := statuses[statusStr]
				if !ok {
					statuses[statusStr] = 1
				} else {
					statuses[statusStr]++
				}
				requestTimeTotal += r.Elapsed
				if requestsSoFar == totalRequests {
					quitting = true
				}
			case <-ticker.C:
				if i == 0 {
					continue
				}
				aggregate = true
			case <-quit:
				ticker.Stop()
				quitting = true
			}
			if aggregate || quitting {
				durationNanoSeconds := lastRequestTime - firstRequestTime
				durationSeconds := float32(durationNanoSeconds) / float32(1000000000)
				var reqPerSec float32
				var kbPerSec float32
				if durationSeconds > 0 {
					reqPerSec = float32(i) / durationSeconds
					kbPerSec = (float32(totBytesRead) / durationSeconds) / 1024.0
				} else {
					reqPerSec = 0
					kbPerSec = 0
				}

				fatalError := ""
				if (totalTimedOut + totalConnectionError) > i/2 {
					fatalError = "Over 50% of requests failed, aborting"
					quitting = true
				}
				aggData := queue.AggData{
					i,
					totalTimedOut,
					totalConnectionError,
					timeToFirstTotal / int64(i),
					totBytesRead,
					statuses,
					requestTimeTotal / int64(i),
					reqPerSec,
					kbPerSec,
					slowest,
					fastest,
					awsregion,
					fatalError,
				}
				sqsAdaptor.SendResult(aggData)
				resetStats = true
			}
		}
	}
	fmt.Printf("\nYay🎈  - %d requests completed\n", requestsSoFar)

}

func fetch(loadTestStartTime time.Time, client *http.Client, address string, requestcount int, jobs <-chan struct{}, ch chan RequestResult, wg *sync.WaitGroup, awsregion string, requestMethod string, requestBody string, requestHeaders []string) {
	defer wg.Done()
	for _ = range jobs {
		start := time.Now()
		req, err := http.NewRequest(requestMethod, address, bytes.NewBufferString(requestBody))
		if err != nil {
			fmt.Println("Error creating the HTTP request:", err)
			return
		}
		req.Header.Add("User-Agent", "Mozilla/5.0 (compatible; Goad/1.0; +https://goad.io)")
		req.Header.Add("Accept-Encoding", "gzip")
		for _, v := range requestHeaders {
			header := strings.Split(v, ":")
			req.Header.Add(strings.Trim(header[0], " "), strings.Trim(header[1], " "))
		}

		response, err := client.Do(req)
		var status string
		var elapsedFirstByte time.Duration
		var elapsedLastByte time.Duration
		var elapsed time.Duration
		var statusCode int
		var bytesRead int
		buf := []byte(" ")
		timedOut := false
		connectionError := false
		isRedirect := err != nil && strings.Contains(err.Error(), "redirect")
		if err != nil && !isRedirect {
			status = fmt.Sprintf("ERROR: %s\n", err)
			switch err := err.(type) {
			case *url.Error:
				if err, ok := err.Err.(net.Error); ok && err.Timeout() {
					timedOut = true
				}
			case net.Error:
				if err.Timeout() {
					timedOut = true
				}
			}

			if !timedOut {
				connectionError = true
			}
		} else {
			statusCode = response.StatusCode
			elapsedFirstByte = time.Since(start)
			if !isRedirect {
				_, err = response.Body.Read(buf)
				firstByteRead := true
				if err != nil {
					status = fmt.Sprintf("reading first byte failed: %s\n", err)
					firstByteRead = false
				}
				body, err := ioutil.ReadAll(response.Body)
				if firstByteRead {
					bytesRead = len(body) + 1
				}
				elapsedLastByte = time.Since(start)
				if err != nil {
					// todo: detect timeout here as well
					status = fmt.Sprintf("reading response body failed: %s\n", err)
					connectionError = true
				} else {
					status = "Success"
				}
			} else {
				status = "Redirect"
			}
			response.Body.Close()

			elapsed = time.Since(start)
		}
		//fmt.Printf("Request end: %d, elapsed: %d\n", time.Now().Sub(loadTestStartTime).Nanoseconds(), elapsed.Nanoseconds())
		result := RequestResult{
			start.Sub(loadTestStartTime).Nanoseconds(),
			req.URL.Host,
			req.Method,
			statusCode,
			elapsed.Nanoseconds(),
			elapsedFirstByte.Nanoseconds(),
			elapsedLastByte.Nanoseconds(),
			bytesRead,
			timedOut,
			connectionError,
			status,
		}
		ch <- result
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
)

func main() {
	var wg sync.WaitGroup
	// passed wg will be accounted at p.Wait() call
	p := mpb.New(mpb.WithWaitGroup(&wg))
	total := 100
	numBars := 3
	wg.Add(numBars)

	for i := 0; i < numBars; i++ {
		name := fmt.Sprintf("Bar#%d:", i)
		bar := p.AddBar(int64(total),
			mpb.BarID(i),
			mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0),
			mpb.PrependDecorators(
				decor.Name(name),
			),
			mpb.AppendDecorators(
				decor.Any(func(s decor.Statistics) string {
					return fmt.Sprintf("completed: %v", s.Completed)
				}, decor.WCSyncSpaceR),
				decor.Any(func(s decor.Statistics) string {
					return fmt.Sprintf("aborted: %v", s.Aborted)
				}, decor.WCSyncSpaceR),
				decor.OnComplete(decor.NewPercentage("%d", decor.WCSyncSpace), "done"),
				decor.OnAbort(decor.NewPercentage("%d", decor.WCSyncSpace), "ohno"),
			),
		)
		go func() {
			defer wg.Done()
			rng := rand.New(rand.NewSource(time.Now().UnixNano()))
			max := 100 * time.Millisecond
			for i := 0; bar.IsRunning(); i++ {
				if bar.ID() == 2 && i >= 42 {
					bar.Abort(false)
				}
				time.Sleep(time.Duration(rng.Intn(10)+1) * max / 10)
				bar.Increment()
			}
		}()
	}
	// wait for passed wg and for all bars to complete and flush
	p.Wait()
}
<commit_msg>abort bar asynchronously<commit_after>package main

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

	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
)

func main() {
	var wg sync.WaitGroup
	// passed wg will be accounted at p.Wait() call
	p := mpb.New(mpb.WithWaitGroup(&wg))
	total := 100
	numBars := 3
	wg.Add(numBars)

	for i := 0; i < numBars; i++ {
		name := fmt.Sprintf("Bar#%d:", i)
		bar := p.AddBar(int64(total),
			mpb.BarID(i),
			mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0),
			mpb.PrependDecorators(
				decor.Name(name),
			),
			mpb.AppendDecorators(
				decor.Any(func(s decor.Statistics) string {
					return fmt.Sprintf("completed: %v", s.Completed)
				}, decor.WCSyncSpaceR),
				decor.Any(func(s decor.Statistics) string {
					return fmt.Sprintf("aborted: %v", s.Aborted)
				}, decor.WCSyncSpaceR),
				decor.OnComplete(decor.NewPercentage("%d", decor.WCSyncSpace), "done"),
				decor.OnAbort(decor.NewPercentage("%d", decor.WCSyncSpace), "ohno"),
			),
		)
		go func() {
			defer wg.Done()
			rng := rand.New(rand.NewSource(time.Now().UnixNano()))
			max := 100 * time.Millisecond
			for i := 0; bar.IsRunning(); i++ {
				if bar.ID() == 2 && i >= 42 {
					go bar.Abort(false)
				}
				time.Sleep(time.Duration(rng.Intn(10)+1) * max / 10)
				bar.Increment()
			}
		}()
	}
	// wait for passed wg and for all bars to complete and flush
	p.Wait()
}
<|endoftext|>


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

package api

import (
	"fmt"
	"github.com/goharbor/harbor/src/common"
	"github.com/goharbor/harbor/src/common/dao"
	"github.com/goharbor/harbor/src/common/models"
	common_quota "github.com/goharbor/harbor/src/common/quota"
	"github.com/goharbor/harbor/src/common/utils/log"
	"github.com/goharbor/harbor/src/core/config"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/pkg/errors"
	"strconv"

	quota "github.com/goharbor/harbor/src/core/api/quota"
)

// InternalAPI handles request of harbor admin...
type InternalAPI struct {
	BaseController
}

// Prepare validates the URL and parms
func (ia *InternalAPI) Prepare() {
	ia.BaseController.Prepare()
	if !ia.SecurityCtx.IsAuthenticated() {
		ia.SendUnAuthorizedError(errors.New("UnAuthorized"))
		return
	}
	if !ia.SecurityCtx.IsSysAdmin() {
		ia.SendForbiddenError(errors.New(ia.SecurityCtx.GetUsername()))
		return
	}
}

// SyncRegistry ...
func (ia *InternalAPI) SyncRegistry() {
	err := SyncRegistry(ia.ProjectMgr)
	if err != nil {
		ia.SendInternalServerError(err)
		return
	}
}

// RenameAdmin we don't provide flexibility in this API, as this is a workaround.
func (ia *InternalAPI) RenameAdmin() {
	if !dao.IsSuperUser(ia.SecurityCtx.GetUsername()) {
		log.Errorf("User %s is not super user, not allow to rename admin.", ia.SecurityCtx.GetUsername())
		ia.SendForbiddenError(errors.New(ia.SecurityCtx.GetUsername()))
		return
	}
	newName := common.NewHarborAdminName
	if err := dao.ChangeUserProfile(models.User{
		UserID:   1,
		Username: newName,
	}, "username"); err != nil {
		log.Errorf("Failed to change admin's username, error: %v", err)
		ia.SendInternalServerError(errors.New("failed to rename admin user"))
		return
	}
	log.Debugf("The super user has been renamed to: %s", newName)
	ia.DestroySession()
}

// QuotaSwitcher ...
type QuotaSwitcher struct {
	Enabled bool
}

// SwitchQuota ...
func (ia *InternalAPI) SwitchQuota() {
	var req QuotaSwitcher
	if err := ia.DecodeJSONReq(&req); err != nil {
		ia.SendBadRequestError(err)
		return
	}
	// quota per project from disable to enable, it needs to update the quota usage bases on the DB records.
	if !config.QuotaPerProjectEnable() && req.Enabled {
		if err := ia.ensureQuota(); err != nil {
			ia.SendInternalServerError(err)
			return
		}
	}
	defer func() {
		config.GetCfgManager().Set(common.QuotaPerProjectEnable, req.Enabled)
		config.GetCfgManager().Save()
	}()
	return
}

func (ia *InternalAPI) ensureQuota() error {
	projects, err := dao.GetProjects(nil)
	if err != nil {
		return err
	}
	for _, project := range projects {
		pSize, err := dao.CountSizeOfProject(project.ProjectID)
		if err != nil {
			logger.Warningf("error happen on counting size of project:%d , error:%v, just skip it.", project.ProjectID, err)
			continue
		}
		afQuery := &models.ArtifactQuery{
			PID: project.ProjectID,
		}
		afs, err := dao.ListArtifacts(afQuery)
		if err != nil {
			logger.Warningf("error happen on counting number of project:%d , error:%v, just skip it.", project.ProjectID, err)
			continue
		}
		pCount := int64(len(afs))

		// it needs to append the chart count
		if config.WithChartMuseum() {
			count, err := chartController.GetCountOfCharts([]string{project.Name})
			if err != nil {
				err = errors.Wrap(err, fmt.Sprintf("get chart count of project %d failed", project.ProjectID))
				logger.Error(err)
				continue
			}
			pCount = pCount + int64(count)
		}

		quotaMgr, err := common_quota.NewManager("project", strconv.FormatInt(project.ProjectID, 10))
		if err != nil {
			logger.Errorf("Error occurred when to new quota manager %v, just skip it.", err)
			continue
		}
		used := common_quota.ResourceList{
			common_quota.ResourceStorage: pSize,
			common_quota.ResourceCount:   pCount,
		}
		if err := quotaMgr.EnsureQuota(used); err != nil {
			logger.Errorf("cannot ensure quota for the project: %d, err: %v, just skip it.", project.ProjectID, err)
			continue
		}
	}
	return nil
}

// SyncQuota ...
func (ia *InternalAPI) SyncQuota() {
	cur := config.ReadOnly()
	cfgMgr := config.GetCfgManager()
	if !cur {
		cfgMgr.Set(common.ReadOnly, true)
		cfgMgr.Save()
	}
	// For api call, to avoid the timeout, it should be asynchronous
	go func() {
		defer func() {
			cfgMgr.Set(common.ReadOnly, cur)
			cfgMgr.Save()
		}()
		log.Info("start to sync quota(API), the system will be set to ReadOnly and back it normal once it done.")
		err := quota.Sync(ia.ProjectMgr, false)
		if err != nil {
			log.Errorf("fail to sync quota(API), but with error: %v, please try to do it again.", err)
			return
		}
		log.Info("success to sync quota(API).")
	}()
	return
}
<commit_msg>add read only for quota switcher<commit_after>// Copyright 2018 Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"fmt"
	"github.com/goharbor/harbor/src/common"
	"github.com/goharbor/harbor/src/common/dao"
	"github.com/goharbor/harbor/src/common/models"
	common_quota "github.com/goharbor/harbor/src/common/quota"
	"github.com/goharbor/harbor/src/common/utils/log"
	"github.com/goharbor/harbor/src/core/config"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/pkg/errors"
	"strconv"

	quota "github.com/goharbor/harbor/src/core/api/quota"
)

// InternalAPI handles request of harbor admin...
type InternalAPI struct {
	BaseController
}

// Prepare validates the URL and parms
func (ia *InternalAPI) Prepare() {
	ia.BaseController.Prepare()
	if !ia.SecurityCtx.IsAuthenticated() {
		ia.SendUnAuthorizedError(errors.New("UnAuthorized"))
		return
	}
	if !ia.SecurityCtx.IsSysAdmin() {
		ia.SendForbiddenError(errors.New(ia.SecurityCtx.GetUsername()))
		return
	}
}

// SyncRegistry ...
func (ia *InternalAPI) SyncRegistry() {
	err := SyncRegistry(ia.ProjectMgr)
	if err != nil {
		ia.SendInternalServerError(err)
		return
	}
}

// RenameAdmin we don't provide flexibility in this API, as this is a workaround.
func (ia *InternalAPI) RenameAdmin() {
	if !dao.IsSuperUser(ia.SecurityCtx.GetUsername()) {
		log.Errorf("User %s is not super user, not allow to rename admin.", ia.SecurityCtx.GetUsername())
		ia.SendForbiddenError(errors.New(ia.SecurityCtx.GetUsername()))
		return
	}
	newName := common.NewHarborAdminName
	if err := dao.ChangeUserProfile(models.User{
		UserID:   1,
		Username: newName,
	}, "username"); err != nil {
		log.Errorf("Failed to change admin's username, error: %v", err)
		ia.SendInternalServerError(errors.New("failed to rename admin user"))
		return
	}
	log.Debugf("The super user has been renamed to: %s", newName)
	ia.DestroySession()
}

// QuotaSwitcher ...
type QuotaSwitcher struct {
	Enabled bool
}

// SwitchQuota ...
func (ia *InternalAPI) SwitchQuota() {
	var req QuotaSwitcher
	if err := ia.DecodeJSONReq(&req); err != nil {
		ia.SendBadRequestError(err)
		return
	}
	cur := config.ReadOnly()
	// quota per project from disable to enable, it needs to update the quota usage bases on the DB records.
	if !config.QuotaPerProjectEnable() && req.Enabled {
		if !cur {
			config.GetCfgManager().Set(common.ReadOnly, true)
			config.GetCfgManager().Save()
		}
		if err := ia.ensureQuota(); err != nil {
			ia.SendInternalServerError(err)
			return
		}
	}
	defer func() {
		config.GetCfgManager().Set(common.ReadOnly, cur)
		config.GetCfgManager().Set(common.QuotaPerProjectEnable, req.Enabled)
		config.GetCfgManager().Save()
	}()
	return
}

func (ia *InternalAPI) ensureQuota() error {
	projects, err := dao.GetProjects(nil)
	if err != nil {
		return err
	}
	for _, project := range projects {
		pSize, err := dao.CountSizeOfProject(project.ProjectID)
		if err != nil {
			logger.Warningf("error happen on counting size of project:%d , error:%v, just skip it.", project.ProjectID, err)
			continue
		}
		afQuery := &models.ArtifactQuery{
			PID: project.ProjectID,
		}
		afs, err := dao.ListArtifacts(afQuery)
		if err != nil {
			logger.Warningf("error happen on counting number of project:%d , error:%v, just skip it.", project.ProjectID, err)
			continue
		}
		pCount := int64(len(afs))

		// it needs to append the chart count
		if config.WithChartMuseum() {
			count, err := chartController.GetCountOfCharts([]string{project.Name})
			if err != nil {
				err = errors.Wrap(err, fmt.Sprintf("get chart count of project %d failed", project.ProjectID))
				logger.Error(err)
				continue
			}
			pCount = pCount + int64(count)
		}

		quotaMgr, err := common_quota.NewManager("project", strconv.FormatInt(project.ProjectID, 10))
		if err != nil {
			logger.Errorf("Error occurred when to new quota manager %v, just skip it.", err)
			continue
		}
		used := common_quota.ResourceList{
			common_quota.ResourceStorage: pSize,
			common_quota.ResourceCount:   pCount,
		}
		if err := quotaMgr.EnsureQuota(used); err != nil {
			logger.Errorf("cannot ensure quota for the project: %d, err: %v, just skip it.", project.ProjectID, err)
			continue
		}
	}
	return nil
}

// SyncQuota ...
func (ia *InternalAPI) SyncQuota() {
	cur := config.ReadOnly()
	cfgMgr := config.GetCfgManager()
	if !cur {
		cfgMgr.Set(common.ReadOnly, true)
		cfgMgr.Save()
	}
	// For api call, to avoid the timeout, it should be asynchronous
	go func() {
		defer func() {
			cfgMgr.Set(common.ReadOnly, cur)
			cfgMgr.Save()
		}()
		log.Info("start to sync quota(API), the system will be set to ReadOnly and back it normal once it done.")
		err := quota.Sync(ia.ProjectMgr, false)
		if err != nil {
			log.Errorf("fail to sync quota(API), but with error: %v, please try to do it again.", err)
			return
		}
		log.Info("success to sync quota(API).")
	}()
	return
}
<|endoftext|>


<commit_before>package main

import (
	"container/list"
	"flag"
	"log"
	"math/rand"
	"sort"
	"time"

	"github.com/afg984/thegame/server/go/src/thegame/pb"
)

var ticksPerSecond int

func init() {
	flag.IntVar(&ticksPerSecond, "tps", 30, "ticks per second")
}

type HeroControls struct {
	*Hero
	*pb.Controls
}

type joinRequest struct {
	name string
	ch   chan *list.Element
}

type Arena struct {
	polygons       [360]*Polygon
	heroCounter    int
	bulletCounter  int
	polygonCounter int
	heroes         *list.List
	bullets        []*Bullet
	controlChan    chan HeroControls
	joinChan       chan joinRequest
	quitChan       chan *list.Element
	commandChan    chan GameCommand
	viewRemotes    []chan *pb.GameState
	viewChan       chan chan *pb.GameState
}

func NewArena() *Arena {
	a := &Arena{
		heroes:      list.New(),
		controlChan: make(chan HeroControls),
		joinChan:    make(chan joinRequest),
		quitChan:    make(chan *list.Element),
		commandChan: make(chan GameCommand),
		viewChan:    make(chan chan *pb.GameState),
	}
	for i := 0; i < 30; i++ {
		a.polygons[i] = &Polygon{shape: Pentagon}
	}
	for i := 30; i < 90; i++ {
		a.polygons[i] = &Polygon{shape: Triangle}
	}
	for i := 90; i < 360; i++ {
		a.polygons[i] = &Polygon{shape: Square}
	}
	go a.Run()
	return a
}

func (a *Arena) Join(name string) *list.Element {
	c := make(chan *list.Element)
	a.joinChan <- joinRequest{name, c}
	return <-c
}

func (a *Arena) Quit(e *list.Element) {
	a.quitChan <- e
}

// gc cleans up unused objects
func (a *Arena) gc() {
	// remove disconnected and dead heroes
	var next *list.Element
	for e := a.heroes.Front(); e != nil; e = next {
		next = e.Next()

		h := e.Value.(*Hero)
		if h.health <= 0 && h.disconnected {
			a.heroes.Remove(e)
		}
	}
}

func (a *Arena) tick() {
	a.gc()

	var objects []Collidable
	for _, p := range a.polygons {
		TickPosition(p)
		if p.visible {
			objects = append(objects, p)
		}
	}
	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if !h.visible {
			h.TryRespawn()
		}
		if h.visible {
			TickPosition(h)
			objects = append(objects, h)
		}
	}
	for _, b := range a.bullets {
		TickPosition(b)
		if b.visible {
			objects = append(objects, b)
		}
	}

	DoCollision(objects)

	for _, b := range a.bullets {
		b.timeout--
	}
	a.bullets = filterBullets(a.bullets)

	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if h.visible {
			h.Action(a)
		}
	}

	// random respawn polygon
	for _, p := range a.polygons {
		if !p.visible {
			if rand.Float64() < 0.008 {
				a.polygonCounter++
				p.id = a.polygonCounter
				p.visible = true
				p.health = p.MaxHealth()
				p.position = RandomPosition()
			}
		}
	}
}

// remove dead and timed out bullets
func filterBullets(a []*Bullet) []*Bullet {
	// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
	b := a[:0]
	for _, x := range a {
		if x.health > 0 && x.timeout > 0 {
			b = append(b, x)
		}
	}
	return b
}

const fieldOfView = 800

func (a *Arena) broadcastView(gs *pb.GameState) {
	remotes := a.viewRemotes[:0]
	for _, remote := range a.viewRemotes {
		select {
		case remote <- gs:
			remotes = append(remotes, remote)
		default:
			close(remote)
			log.Println("view client not responding to updates")
		}
	}
	a.viewRemotes = remotes
}

func (a *Arena) broadcast() {
	var polygons []*pb.Polygon
	var bullets []*pb.Bullet
	var heroes []*pb.Hero
	var scores []*pb.ScoreEntry
	maxScore := 0
	var maxScoreHero *Hero
	for _, p := range a.polygons {
		if p.visible {
			polygons = append(polygons, p.ToProto())
		}
	}
	for _, b := range a.bullets {
		bullets = append(bullets, b.ToProto())
	}
	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if h.score > maxScore {
			maxScore = h.score
			maxScoreHero = h
		}
		if h.visible {
			heroes = append(heroes, h.ToProto())
		}
		scores = append(scores, h.ToScoreEntry())
	}
	sort.Sort(ByScore(scores))
	// send updates to clients
	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if h.disconnected {
			continue
		}
		canSee := func(w *pb.Entity) bool {
			x := real(h.position)
			y := imag(h.position)
			return (w.Position.X+w.Radius > x-fieldOfView &&
				w.Position.X-w.Radius < x+fieldOfView &&
				w.Position.Y-w.Radius > y-fieldOfView &&
				w.Position.Y-w.Radius < y+fieldOfView)
		}
		var spolygons []*pb.Polygon
		var sbullets []*pb.Bullet
		sheroes := []*pb.Hero{h.ToProto()}
		for _, p := range polygons {
			if canSee(p.Entity) {
				spolygons = append(spolygons, p)
			}
		}
		for _, b := range bullets {
			if canSee(b.Entity) {
				sbullets = append(sbullets, b)
			}
		}
		for _, oh := range heroes {
			if oh.Entity.Id == int32(h.id) {
				continue
			}
			if canSee(oh.Entity) {
				sheroes = append(sheroes, oh)
			}
		}
		state := &pb.GameState{
			Meta: &pb.GameState_Meta{
				HeroId: int32(h.id),
				Scores: scores,
			},
			Polygons: spolygons,
			Bullets:  sbullets,
			Heroes:   sheroes,
		}
		select {
		case h.UpdateChan <- state:
		default:
			log.Println(h, "not responding to updates")
		}
		if h == maxScoreHero {
			a.broadcastView(state)
		}
	}
}

func (a *Arena) Run() {
	tick := time.Tick(time.Second / time.Duration(ticksPerSecond))
	perfTick := time.Tick(time.Second)
	var tickCount int64
	var lastTick int64
	paused := false
	for {
		select {
		case <-tick:
			if !paused {
				tickCount++
				a.tick()
			}
			a.broadcast()
		case <-perfTick:
			log.Println("ticks per second:", tickCount-lastTick)
			lastTick = tickCount
		case hc := <-a.controlChan:
			hc.Hero.controls = hc.Controls
		case jr := <-a.joinChan:
			a.heroCounter++
			h := NewHero(a.heroCounter, jr.name)
			h.Spawn()
			log.Println(h, "joined the arena")
			jr.ch <- a.heroes.PushBack(h)
		case l := <-a.quitChan:
			h := l.Value.(*Hero)
			log.Println(h, "left the arena")
			h.disconnected = true
		case cgs := <-a.viewChan:
			a.viewRemotes = append(a.viewRemotes, cgs)
		case c := <-a.commandChan:
			switch c {
			case Pause:
				log.Println("Game Paused")
				paused = true
			case Resume:
				log.Println("Game Resumed")
				paused = false
			case c:
				log.Printf("No known action to %q", c)
			}
		}
	}
}
<commit_msg>server/go: fix: disconnected heroes should not be able to Action()<commit_after>package main

import (
	"container/list"
	"flag"
	"log"
	"math/rand"
	"sort"
	"time"

	"github.com/afg984/thegame/server/go/src/thegame/pb"
)

var ticksPerSecond int

func init() {
	flag.IntVar(&ticksPerSecond, "tps", 30, "ticks per second")
}

type HeroControls struct {
	*Hero
	*pb.Controls
}

type joinRequest struct {
	name string
	ch   chan *list.Element
}

type Arena struct {
	polygons       [360]*Polygon
	heroCounter    int
	bulletCounter  int
	polygonCounter int
	heroes         *list.List
	bullets        []*Bullet
	controlChan    chan HeroControls
	joinChan       chan joinRequest
	quitChan       chan *list.Element
	commandChan    chan GameCommand
	viewRemotes    []chan *pb.GameState
	viewChan       chan chan *pb.GameState
}

func NewArena() *Arena {
	a := &Arena{
		heroes:      list.New(),
		controlChan: make(chan HeroControls),
		joinChan:    make(chan joinRequest),
		quitChan:    make(chan *list.Element),
		commandChan: make(chan GameCommand),
		viewChan:    make(chan chan *pb.GameState),
	}
	for i := 0; i < 30; i++ {
		a.polygons[i] = &Polygon{shape: Pentagon}
	}
	for i := 30; i < 90; i++ {
		a.polygons[i] = &Polygon{shape: Triangle}
	}
	for i := 90; i < 360; i++ {
		a.polygons[i] = &Polygon{shape: Square}
	}
	go a.Run()
	return a
}

func (a *Arena) Join(name string) *list.Element {
	c := make(chan *list.Element)
	a.joinChan <- joinRequest{name, c}
	return <-c
}

func (a *Arena) Quit(e *list.Element) {
	a.quitChan <- e
}

// gc cleans up unused objects
func (a *Arena) gc() {
	// remove disconnected and dead heroes
	var next *list.Element
	for e := a.heroes.Front(); e != nil; e = next {
		next = e.Next()

		h := e.Value.(*Hero)
		if h.health <= 0 && h.disconnected {
			a.heroes.Remove(e)
		}
	}
}

func (a *Arena) tick() {
	a.gc()

	var objects []Collidable
	for _, p := range a.polygons {
		TickPosition(p)
		if p.visible {
			objects = append(objects, p)
		}
	}
	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if !h.visible {
			h.TryRespawn()
		}
		if h.visible {
			TickPosition(h)
			objects = append(objects, h)
		}
	}
	for _, b := range a.bullets {
		TickPosition(b)
		if b.visible {
			objects = append(objects, b)
		}
	}

	DoCollision(objects)

	for _, b := range a.bullets {
		b.timeout--
	}
	a.bullets = filterBullets(a.bullets)

	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if h.visible && !h.disconnected {
			h.Action(a)
		}
	}

	// random respawn polygon
	for _, p := range a.polygons {
		if !p.visible {
			if rand.Float64() < 0.008 {
				a.polygonCounter++
				p.id = a.polygonCounter
				p.visible = true
				p.health = p.MaxHealth()
				p.position = RandomPosition()
			}
		}
	}
}

// remove dead and timed out bullets
func filterBullets(a []*Bullet) []*Bullet {
	// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
	b := a[:0]
	for _, x := range a {
		if x.health > 0 && x.timeout > 0 {
			b = append(b, x)
		}
	}
	return b
}

const fieldOfView = 800

func (a *Arena) broadcastView(gs *pb.GameState) {
	remotes := a.viewRemotes[:0]
	for _, remote := range a.viewRemotes {
		select {
		case remote <- gs:
			remotes = append(remotes, remote)
		default:
			close(remote)
			log.Println("view client not responding to updates")
		}
	}
	a.viewRemotes = remotes
}

func (a *Arena) broadcast() {
	var polygons []*pb.Polygon
	var bullets []*pb.Bullet
	var heroes []*pb.Hero
	var scores []*pb.ScoreEntry
	maxScore := 0
	var maxScoreHero *Hero
	for _, p := range a.polygons {
		if p.visible {
			polygons = append(polygons, p.ToProto())
		}
	}
	for _, b := range a.bullets {
		bullets = append(bullets, b.ToProto())
	}
	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if h.score > maxScore {
			maxScore = h.score
			maxScoreHero = h
		}
		if h.visible {
			heroes = append(heroes, h.ToProto())
		}
		scores = append(scores, h.ToScoreEntry())
	}
	sort.Sort(ByScore(scores))
	// send updates to clients
	for e := a.heroes.Front(); e != nil; e = e.Next() {
		h := e.Value.(*Hero)
		if h.disconnected {
			continue
		}
		canSee := func(w *pb.Entity) bool {
			x := real(h.position)
			y := imag(h.position)
			return (w.Position.X+w.Radius > x-fieldOfView &&
				w.Position.X-w.Radius < x+fieldOfView &&
				w.Position.Y-w.Radius > y-fieldOfView &&
				w.Position.Y-w.Radius < y+fieldOfView)
		}
		var spolygons []*pb.Polygon
		var sbullets []*pb.Bullet
		sheroes := []*pb.Hero{h.ToProto()}
		for _, p := range polygons {
			if canSee(p.Entity) {
				spolygons = append(spolygons, p)
			}
		}
		for _, b := range bullets {
			if canSee(b.Entity) {
				sbullets = append(sbullets, b)
			}
		}
		for _, oh := range heroes {
			if oh.Entity.Id == int32(h.id) {
				continue
			}
			if canSee(oh.Entity) {
				sheroes = append(sheroes, oh)
			}
		}
		state := &pb.GameState{
			Meta: &pb.GameState_Meta{
				HeroId: int32(h.id),
				Scores: scores,
			},
			Polygons: spolygons,
			Bullets:  sbullets,
			Heroes:   sheroes,
		}
		select {
		case h.UpdateChan <- state:
		default:
			log.Println(h, "not responding to updates")
		}
		if h == maxScoreHero {
			a.broadcastView(state)
		}
	}
}

func (a *Arena) Run() {
	tick := time.Tick(time.Second / time.Duration(ticksPerSecond))
	perfTick := time.Tick(time.Second)
	var tickCount int64
	var lastTick int64
	paused := false
	for {
		select {
		case <-tick:
			if !paused {
				tickCount++
				a.tick()
			}
			a.broadcast()
		case <-perfTick:
			log.Println("ticks per second:", tickCount-lastTick)
			lastTick = tickCount
		case hc := <-a.controlChan:
			hc.Hero.controls = hc.Controls
		case jr := <-a.joinChan:
			a.heroCounter++
			h := NewHero(a.heroCounter, jr.name)
			h.Spawn()
			log.Println(h, "joined the arena")
			jr.ch <- a.heroes.PushBack(h)
		case l := <-a.quitChan:
			h := l.Value.(*Hero)
			log.Println(h, "left the arena")
			h.disconnected = true
		case cgs := <-a.viewChan:
			a.viewRemotes = append(a.viewRemotes, cgs)
		case c := <-a.commandChan:
			switch c {
			case Pause:
				log.Println("Game Paused")
				paused = true
			case Resume:
				log.Println("Game Resumed")
				paused = false
			case c:
				log.Printf("No known action to %q", c)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestGetSystemLoad(t *testing.T) {
	load, err := GetSystemLoad()

	assert.NotNil(t, load)
	assert.Nil(t, err)

	assert.True(t, (load.Load1 > 0))
	assert.True(t, (load.Load5 > 0))
	assert.True(t, (load.Load15 > 0))
}

func TestGetMemory(t *testing.T) {
	mem, err := GetMemory()

	assert.NotNil(t, mem)
	assert.Nil(t, err)

	assert.True(t, (mem.Total > 0))
	assert.True(t, (mem.Used > 0))
	assert.True(t, (mem.Free >= 0))
	assert.True(t, (mem.ActualFree >= 0))
	assert.True(t, (mem.ActualUsed > 0))
}

func TestGetSwap(t *testing.T) {
	swap, err := GetSwap()

	assert.NotNil(t, swap)
	assert.Nil(t, err)

	assert.True(t, (swap.Total >= 0))
	assert.True(t, (swap.Used >= 0))
	assert.True(t, (swap.Total >= 0))
}

func TestPids(t *testing.T) {
	pids, err := Pids()

	assert.NotNil(t, pids)
	assert.Nil(t, err)

	// Assuming at least 2 processes are running
	assert.True(t, (len(pids) > 1))
}

func TestGetProcess(t *testing.T) {
	pids, err := Pids()

	assert.Nil(t, err)

	process, err := GetProcess(pids[0])

	assert.NotNil(t, process)
	assert.Nil(t, err)

	assert.True(t, (process.Pid > 0))
	assert.True(t, (process.Ppid > 0))
	assert.True(t, (len(process.Name) > 0))
	assert.NotEqual(t, "unknown", process.State)

	// Memory Checks
	assert.True(t, (process.Mem.Size >= 0))
	assert.True(t, (process.Mem.Resident >= 0))
	assert.True(t, (process.Mem.Share >= 0))

	// CPU Checks
	assert.True(t, (len(process.Cpu.Start) > 0))
	assert.True(t, (process.Cpu.Total >= 0))
	assert.True(t, (process.Cpu.User >= 0))
	assert.True(t, (process.Cpu.System >= 0))

	assert.True(t, (process.lastCPUTime.Unix() <= time.Now().Unix()))
}
<commit_msg>Make Ppid = 0 also an option<commit_after>package main

import (
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestGetSystemLoad(t *testing.T) {
	load, err := GetSystemLoad()

	assert.NotNil(t, load)
	assert.Nil(t, err)

	assert.True(t, (load.Load1 > 0))
	assert.True(t, (load.Load5 > 0))
	assert.True(t, (load.Load15 > 0))
}

func TestGetMemory(t *testing.T) {
	mem, err := GetMemory()

	assert.NotNil(t, mem)
	assert.Nil(t, err)

	assert.True(t, (mem.Total > 0))
	assert.True(t, (mem.Used > 0))
	assert.True(t, (mem.Free >= 0))
	assert.True(t, (mem.ActualFree >= 0))
	assert.True(t, (mem.ActualUsed > 0))
}

func TestGetSwap(t *testing.T) {
	swap, err := GetSwap()

	assert.NotNil(t, swap)
	assert.Nil(t, err)

	assert.True(t, (swap.Total >= 0))
	assert.True(t, (swap.Used >= 0))
	assert.True(t, (swap.Total >= 0))
}

func TestPids(t *testing.T) {
	pids, err := Pids()

	assert.NotNil(t, pids)
	assert.Nil(t, err)

	// Assuming at least 2 processes are running
	assert.True(t, (len(pids) > 1))
}

func TestGetProcess(t *testing.T) {
	pids, err := Pids()

	assert.Nil(t, err)

	process, err := GetProcess(pids[0])

	assert.NotNil(t, process)
	assert.Nil(t, err)

	assert.True(t, (process.Pid > 0))
	assert.True(t, (process.Ppid >= 0))
	assert.True(t, (len(process.Name) > 0))
	assert.NotEqual(t, "unknown", process.State)

	// Memory Checks
	assert.True(t, (process.Mem.Size >= 0))
	assert.True(t, (process.Mem.Resident >= 0))
	assert.True(t, (process.Mem.Share >= 0))

	// CPU Checks
	assert.True(t, (len(process.Cpu.Start) > 0))
	assert.True(t, (process.Cpu.Total >= 0))
	assert.True(t, (process.Cpu.User >= 0))
	assert.True(t, (process.Cpu.System >= 0))

	assert.True(t, (process.lastCPUTime.Unix() <= time.Now().Unix()))
}
<|endoftext|>


<commit_before>package bot

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

	"github.com/nlopes/slack"
	"strings"
)

type commit struct {
	ID      string `json:"change_id"`
	Number  int    `json:"_number"`
	Subject string `json:"subject"`
	Branch  string `json:"branch"`
}

var errNonOK = errors.New("non-ok status")

func (b *Bot) MonitorGerrit() {
	tk := time.NewTicker(1 * time.Hour)
	defer tk.Stop()

	lastID := ""

	historyParams := slack.HistoryParameters{Count: 100}
	history, err := b.slackAPI.GetChannelHistory(b.channels["golang_cls"], historyParams)
	if err != nil {
		b.log(err)
	} else {
		for _, msg := range history.Messages {
			if msg.BotID != b.id {
				continue
			}
			parts := strings.Split(strings.ToLower(msg.Text), " #-# ")
			if len(parts) != 3 {
				continue
			}
			lastID = parts[1]
			break
		}
	}

	params := slack.PostMessageParameters{AsUser: true}
	for <-tk.C {
		req, err := http.NewRequest("GET", b.gerritLink, nil)
		req.Header.Add("User-Agent", "Gophers Slack bot")
		resp, err := b.client.Do(req)
		if err != nil {
			b.log("%s\n", err)
			continue
		}

		if resp.StatusCode != http.StatusOK {
			resp.Body.Close()
			continue
		}

		body, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			b.log("%s\n", err)
			continue
		}

		// Fix Gerrit adding a random prefix )]}'
		body = body[4:]
		cls := []commit{}
		err = json.Unmarshal(body, &cls)
		if err != nil {
			b.log("%s\n", err)
			continue
		}

		for _, cl := range cls {
			if cl.ID == lastID {
				break
			}
			clLine := fmt.Sprintf("%d #-# %s #-# %s", cl.Number, cl.ID, cl.Subject)
			_, _, err = b.slackAPI.PostMessage(b.channels["golang_cls"].slackID, clLine, params)
			if err != nil {
				b.log("%s\n", err)
				continue
			}
		}
	}
}
<commit_msg>Fix previous commit<commit_after>package bot

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

	"github.com/nlopes/slack"
	"strings"
)

type commit struct {
	ID      string `json:"change_id"`
	Number  int    `json:"_number"`
	Subject string `json:"subject"`
	Branch  string `json:"branch"`
}

var errNonOK = errors.New("non-ok status")

func (b *Bot) MonitorGerrit() {
	tk := time.NewTicker(1 * time.Hour)
	defer tk.Stop()

	lastID := ""

	historyParams := slack.HistoryParameters{Count: 100}
	history, err := b.slackAPI.GetChannelHistory(b.channels["golang_cls"].slackID, historyParams)
	if err != nil {
		b.log("error while fetching history: %v\n", err)
	} else {
		for _, msg := range history.Messages {
			if msg.BotID != b.id {
				continue
			}
			parts := strings.Split(strings.ToLower(msg.Text), " #-# ")
			if len(parts) != 3 {
				continue
			}
			lastID = parts[1]
			break
		}
	}

	params := slack.PostMessageParameters{AsUser: true}
	for range tk.C {
		req, err := http.NewRequest("GET", b.gerritLink, nil)
		req.Header.Add("User-Agent", "Gophers Slack bot")
		resp, err := b.client.Do(req)
		if err != nil {
			b.log("%s\n", err)
			continue
		}

		if resp.StatusCode != http.StatusOK {
			resp.Body.Close()
			continue
		}

		body, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			b.log("%s\n", err)
			continue
		}

		// Fix Gerrit adding a random prefix )]}'
		body = body[4:]
		cls := []commit{}
		err = json.Unmarshal(body, &cls)
		if err != nil {
			b.log("%s\n", err)
			continue
		}

		for _, cl := range cls {
			if cl.ID == lastID {
				break
			}
			clLine := fmt.Sprintf("%d #-# %s #-# %s", cl.Number, cl.ID, cl.Subject)
			_, _, err = b.slackAPI.PostMessage(b.channels["golang_cls"].slackID, clLine, params)
			if err != nil {
				b.log("%s\n", err)
				continue
			}
		}
	}
}
<|endoftext|>


<commit_before>package bitbucket

import (
	"context"
	"crypto/sha256"
	"fmt"
	"net/http"
	"path/filepath"
	"strings"
	"sync"

	"github.com/reviewdog/reviewdog"
	"github.com/reviewdog/reviewdog/proto/rdf"
	"github.com/reviewdog/reviewdog/service/bitbucket/openapi"
)

var _ reviewdog.CommentService = &ReportAnnotator{}

const (
	// avatar from https://github.com/apps/reviewdog
	logoURL  = "https://avatars1.githubusercontent.com/in/12131"
	reporter = "reviewdog"
	// max amount of annotations in one batch call
	annotationsBatchSize = 100
)

// ReportAnnotator is a comment service for Bitbucket Code Insights reports.
//
// API:
//  https://developer.atlassian.com/bitbucket/api/2/reference/resource/repositories/%7Bworkspace%7D/%7Brepo_slug%7D/commit/%7Bcommit%7D/reports/%7BreportId%7D/annotations#post
//  POST /2.0/repositories/{username}/{repo_slug}/commit/{commit}/reports/{reportId}/annotations
type ReportAnnotator struct {
	cli         *openapi.APIClient
	sha         string
	owner, repo string

	muAnnotations sync.Mutex
	// store annotations in map per tool name
	// so we can create report per tool
	annotations map[string][]openapi.ReportAnnotation
	severityMap map[rdf.Severity]string

	// wd is working directory relative to root of repository.
	wd         string
	duplicates map[string]struct{}
}

// NewReportAnnotator creates new Bitbucket Report Annotator
func NewReportAnnotator(cli *openapi.APIClient, owner, repo, sha string, runners []string) *ReportAnnotator {
	r := &ReportAnnotator{
		cli:         cli,
		sha:         sha,
		owner:       owner,
		repo:        repo,
		annotations: make(map[string][]openapi.ReportAnnotation, len(runners)),
		severityMap: map[rdf.Severity]string{
			rdf.Severity_INFO:    annotationSeverityLow,
			rdf.Severity_WARNING: annotationSeverityMedium,
			rdf.Severity_ERROR:   annotationSeverityHigh,
		},
		duplicates: map[string]struct{}{},
	}

	// pre populate map of annotations, so we still create passed (green) report
	// if no issues found from the specific tool
	for _, runner := range runners {
		if len(runner) == 0 {
			continue
		}
		r.annotations[runner] = []openapi.ReportAnnotation{}
		// create Pending report for each tool
		_ = r.createOrUpdateReport(context.Background(), reportID(runner, reporter), reportTitle(runner, reporter), reportResultPending)
	}

	return r
}

// Post accepts a comment and holds it. Flush method actually posts comments to
// Bitbucket in batch.
func (r *ReportAnnotator) Post(_ context.Context, c *reviewdog.Comment) error {
	c.Result.Diagnostic.GetLocation().Path = filepath.ToSlash(
		filepath.Join(r.wd, c.Result.Diagnostic.GetLocation().GetPath()))
	r.muAnnotations.Lock()
	defer r.muAnnotations.Unlock()

	anot := r.annotationFromReviewDogComment(*c)

	// deduplicate event, because some reporters might report
	// it twice, and bitbucket api will complain on duplicated
	// external id of annotation
	if _, ok := r.duplicates[*anot.ExternalId]; !ok {
		r.annotations[c.ToolName] = append(r.annotations[c.ToolName], anot)
	}

	r.duplicates[*anot.ExternalId] = struct{}{}
	return nil
}

// Flush posts comments which has not been posted yet.
func (r *ReportAnnotator) Flush(ctx context.Context) error {
	r.muAnnotations.Lock()
	defer r.muAnnotations.Unlock()

	// create/update/annotate report per tool
	for tool, annotations := range r.annotations {
		reportID := reportID(reporter, tool)
		title := reportTitle(tool, reporter)
		if len(annotations) == 0 {
			// if no annotation, create Passed report
			if err := r.createOrUpdateReport(ctx, reportID, title, reportResultPassed); err != nil {
				return err
			}
			// and move one
			continue
		}

		// create report or update report first, with the failed status
		if err := r.createOrUpdateReport(ctx, reportID, title, reportResultFailed); err != nil {
			return err
		}

		// send annotations in batches, because of the api max payload size limit
		for start, annCount := 0, len(annotations); start < annCount; start += annotationsBatchSize {
			end := start + annotationsBatchSize

			if end > annCount {
				end = annCount
			}

			// add annotations to the report
			_, resp, err := r.cli.ReportsApi.BulkCreateOrUpdateAnnotations(
				ctx, r.owner, r.repo, r.sha, reportID,
			).Body(annotations[start:end]).Execute()

			if err := checkAPIError(err, resp, http.StatusOK); err != nil {
				return fmt.Errorf("bitbucket.BulkCreateOrUpdateAnnotations: %s", err)
			}
		}
	}

	return nil
}

func (r *ReportAnnotator) annotationFromReviewDogComment(c reviewdog.Comment) openapi.ReportAnnotation {
	a := openapi.NewReportAnnotation()

	// TODO: allow providing different annotation types in future
	a.SetAnnotationType(annotationTypeCodeSmell)
	// hash the output of linter and use it as external id
	a.SetExternalId(hashString(c.Result.Diagnostic.OriginalOutput))
	a.SetSummary(c.Result.Diagnostic.GetMessage())
	a.SetDetails(fmt.Sprintf(`[%s] %s`, c.ToolName, c.Result.Diagnostic.GetMessage()))
	a.SetLine(c.Result.Diagnostic.GetLocation().GetRange().GetStart().GetLine())
	a.SetPath(c.Result.Diagnostic.GetLocation().GetPath())
	if v, ok := r.severityMap[c.Result.Diagnostic.GetSeverity()]; ok {
		a.SetSeverity(v)
	}
	if link := c.Result.Diagnostic.GetCode().GetUrl(); link != "" {
		a.SetLink(link)
	}

	return *a
}

func (r *ReportAnnotator) createOrUpdateReport(ctx context.Context, redportID, title, reportStatus string) error {
	var report = openapi.NewReport()
	report.SetTitle(title)
	// TODO: different report types?
	report.SetReportType(reportTypeBug)
	report.SetReporter(reporter)
	report.SetLogoUrl(logoURL)
	report.SetResult(reportStatus)
	if reportStatus == reportResultPassed {
		report.SetDetails("Great news! Reviewdog couldn't spot any issues!")
	} else {
		report.SetDetails("Woof-Woof! This report generated for you by reviewdog")
	}

	_, resp, err := r.cli.ReportsApi.CreateOrUpdateReport(
		ctx, r.owner, r.repo, r.sha, redportID,
	).Body(*report).Execute()

	if err := checkAPIError(err, resp, http.StatusOK); err != nil {
		return fmt.Errorf("bitbucket.CreateOrUpdateReport: %s", err)
	}

	return nil
}

func hashString(str string) string {
	h := sha256.New()
	_, _ = h.Write([]byte(str))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func reportID(ids ...string) string {
	return strings.ReplaceAll(strings.ToLower(strings.Join(ids, "-")), " ", "_")
}

func reportTitle(tool, reporter string) string {
	return fmt.Sprintf("[%s] %s report", tool, reporter)
}
<commit_msg>get external id from diagnostics message<commit_after>package bitbucket

import (
	"context"
	"crypto/sha256"
	"fmt"
	"net/http"
	"path/filepath"
	"strings"
	"sync"

	"github.com/reviewdog/reviewdog"
	"github.com/reviewdog/reviewdog/proto/rdf"
	"github.com/reviewdog/reviewdog/service/bitbucket/openapi"
	"google.golang.org/protobuf/proto"
)

var _ reviewdog.CommentService = &ReportAnnotator{}

const (
	// avatar from https://github.com/apps/reviewdog
	logoURL  = "https://avatars1.githubusercontent.com/in/12131"
	reporter = "reviewdog"
	// max amount of annotations in one batch call
	annotationsBatchSize = 100
)

// ReportAnnotator is a comment service for Bitbucket Code Insights reports.
//
// API:
//  https://developer.atlassian.com/bitbucket/api/2/reference/resource/repositories/%7Bworkspace%7D/%7Brepo_slug%7D/commit/%7Bcommit%7D/reports/%7BreportId%7D/annotations#post
//  POST /2.0/repositories/{username}/{repo_slug}/commit/{commit}/reports/{reportId}/annotations
type ReportAnnotator struct {
	cli         *openapi.APIClient
	sha         string
	owner, repo string

	muAnnotations sync.Mutex
	// store annotations in map per tool name
	// so we can create report per tool
	annotations map[string][]openapi.ReportAnnotation
	severityMap map[rdf.Severity]string

	// wd is working directory relative to root of repository.
	wd         string
	duplicates map[string]struct{}
}

// NewReportAnnotator creates new Bitbucket Report Annotator
func NewReportAnnotator(cli *openapi.APIClient, owner, repo, sha string, runners []string) *ReportAnnotator {
	r := &ReportAnnotator{
		cli:         cli,
		sha:         sha,
		owner:       owner,
		repo:        repo,
		annotations: make(map[string][]openapi.ReportAnnotation, len(runners)),
		severityMap: map[rdf.Severity]string{
			rdf.Severity_INFO:    annotationSeverityLow,
			rdf.Severity_WARNING: annotationSeverityMedium,
			rdf.Severity_ERROR:   annotationSeverityHigh,
		},
		duplicates: map[string]struct{}{},
	}

	// pre populate map of annotations, so we still create passed (green) report
	// if no issues found from the specific tool
	for _, runner := range runners {
		if len(runner) == 0 {
			continue
		}
		r.annotations[runner] = []openapi.ReportAnnotation{}
		// create Pending report for each tool
		_ = r.createOrUpdateReport(context.Background(), reportID(runner, reporter), reportTitle(runner, reporter), reportResultPending)
	}

	return r
}

// Post accepts a comment and holds it. Flush method actually posts comments to
// Bitbucket in batch.
func (r *ReportAnnotator) Post(_ context.Context, c *reviewdog.Comment) error {
	c.Result.Diagnostic.GetLocation().Path = filepath.ToSlash(
		filepath.Join(r.wd, c.Result.Diagnostic.GetLocation().GetPath()))
	r.muAnnotations.Lock()
	defer r.muAnnotations.Unlock()

	anot := r.annotationFromReviewDogComment(*c)

	// deduplicate event, because some reporters might report
	// it twice, and bitbucket api will complain on duplicated
	// external id of annotation
	if _, ok := r.duplicates[*anot.ExternalId]; !ok {
		r.annotations[c.ToolName] = append(r.annotations[c.ToolName], anot)
	}

	r.duplicates[*anot.ExternalId] = struct{}{}
	return nil
}

// Flush posts comments which has not been posted yet.
func (r *ReportAnnotator) Flush(ctx context.Context) error {
	r.muAnnotations.Lock()
	defer r.muAnnotations.Unlock()

	// create/update/annotate report per tool
	for tool, annotations := range r.annotations {
		reportID := reportID(reporter, tool)
		title := reportTitle(tool, reporter)
		if len(annotations) == 0 {
			// if no annotation, create Passed report
			if err := r.createOrUpdateReport(ctx, reportID, title, reportResultPassed); err != nil {
				return err
			}
			// and move one
			continue
		}

		// create report or update report first, with the failed status
		if err := r.createOrUpdateReport(ctx, reportID, title, reportResultFailed); err != nil {
			return err
		}

		// send annotations in batches, because of the api max payload size limit
		for start, annCount := 0, len(annotations); start < annCount; start += annotationsBatchSize {
			end := start + annotationsBatchSize

			if end > annCount {
				end = annCount
			}

			// add annotations to the report
			_, resp, err := r.cli.ReportsApi.BulkCreateOrUpdateAnnotations(
				ctx, r.owner, r.repo, r.sha, reportID,
			).Body(annotations[start:end]).Execute()

			if err := checkAPIError(err, resp, http.StatusOK); err != nil {
				return fmt.Errorf("bitbucket.BulkCreateOrUpdateAnnotations: %s", err)
			}
		}
	}

	return nil
}

func (r *ReportAnnotator) annotationFromReviewDogComment(c reviewdog.Comment) openapi.ReportAnnotation {
	a := openapi.NewReportAnnotation()

	// TODO: allow providing different annotation types in future
	a.SetAnnotationType(annotationTypeCodeSmell)
	// hash the output of linter and use it as external id
	a.SetExternalId(externalIDFromDiagnostic(c.Result.Diagnostic))
	a.SetSummary(c.Result.Diagnostic.GetMessage())
	a.SetDetails(fmt.Sprintf(`[%s] %s`, c.ToolName, c.Result.Diagnostic.GetMessage()))
	a.SetLine(c.Result.Diagnostic.GetLocation().GetRange().GetStart().GetLine())
	a.SetPath(c.Result.Diagnostic.GetLocation().GetPath())
	if v, ok := r.severityMap[c.Result.Diagnostic.GetSeverity()]; ok {
		a.SetSeverity(v)
	}
	if link := c.Result.Diagnostic.GetCode().GetUrl(); link != "" {
		a.SetLink(link)
	}

	return *a
}

func (r *ReportAnnotator) createOrUpdateReport(ctx context.Context, redportID, title, reportStatus string) error {
	var report = openapi.NewReport()
	report.SetTitle(title)
	// TODO: different report types?
	report.SetReportType(reportTypeBug)
	report.SetReporter(reporter)
	report.SetLogoUrl(logoURL)
	report.SetResult(reportStatus)
	if reportStatus == reportResultPassed {
		report.SetDetails("Great news! Reviewdog couldn't spot any issues!")
	} else {
		report.SetDetails("Woof-Woof! This report generated for you by reviewdog")
	}

	_, resp, err := r.cli.ReportsApi.CreateOrUpdateReport(
		ctx, r.owner, r.repo, r.sha, redportID,
	).Body(*report).Execute()

	if err := checkAPIError(err, resp, http.StatusOK); err != nil {
		return fmt.Errorf("bitbucket.CreateOrUpdateReport: %s", err)
	}

	return nil
}

func hash(b []byte) string {
	h := sha256.New()
	_, _ = h.Write(b)
	return fmt.Sprintf("%x", h.Sum(nil))
}

// Note that it might be good to create external ID from the diagnostic
// content along with *original line* (by using git blame for example) to
// generate unique ID, but it hashes the Diagnostic message for simplicity.
func externalIDFromDiagnostic(d *rdf.Diagnostic) string {
	b, err := proto.Marshal(d)
	if err != nil {
		b = []byte(d.OriginalOutput)
	}
	return hash(b)
}

func reportID(ids ...string) string {
	return strings.ReplaceAll(strings.ToLower(strings.Join(ids, "-")), " ", "_")
}

func reportTitle(tool, reporter string) string {
	return fmt.Sprintf("[%s] %s report", tool, reporter)
}
<|endoftext|>