<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
)

const (
	VERSION = "0.0.1"
)

// sort interface-typed arrays by first-class functions
type ByFn struct{
	elems []interface{}
	comp func(a, b interface{}) bool
}
func (c ByFn) Len() int { return len(c.elems)  }
func (c ByFn) Less(i, j int) bool { return c.comp(c.elems[i], c.elems[j]) }
func (c ByFn) Swap(i, j int) { c.elems[i], c.elems[j] = c.elems[j], c.elems[i] }

// generic api requests
func apiReq(meth string, url string) interface{} {
	client := &http.Client{}
	req, err := http.NewRequest(meth, url, nil)
	req.SetBasicAuth("x", os.Getenv("HEROKU_API_KEY"))
	req.Header.Add("User-Agent", fmt.Sprintf("hk/%s", VERSION))
	req.Header.Add("Accept", "application/json")
	res, err := client.Do(req)
		if err != nil {
		panic(err)
	}
	if (res.StatusCode == 401) {
		error("Unauthorized")
	}
	if (res.StatusCode == 403) {
		error("Unauthorized")
	}
	if (res.StatusCode != 200) {
		fmt.Printf("%v\n", res)
		error("Unexpected error")
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}
	var data interface{}
	err = json.Unmarshal(body, &data)
	if err != nil {
		panic(err)
	}
	return data
}

// error formatting
func error(msg string) {
	fmt.Fprintf(os.Stderr, "Error: %s.\n", msg)
	os.Exit(1)
}

func unrecArg(arg string) {
	error(fmt.Sprintf("Unrecognized argument '%s'", arg))
}

func unrecCmd(cmd string) {
	error(fmt.Sprintf("'%s' is not an hk command. See 'hk help'", cmd))
}

// info formatting
func cmdHelp(usage string, desc string) {
	fmt.Printf("Usage: %s\n\n", usage)
	fmt.Printf("%s.\n", desc)
	os.Exit(0)
}

// commands
func envHelp() {
	cmdHelp("hk env -a <app>", "Show all config vars")
}

func env() {
	if (len(os.Args) != 4) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help env'")
	}
	appName := os.Args[3]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s/config_vars", appName))
	config := data.(map[string]interface{})
	for k, v := range config {
		fmt.Printf("%s=%v\n", k, v)
	}
	os.Exit(0)
}

func getHelp() {
	cmdHelp("hk get -a <app> <key>", "Get the value of a config var")
}

func get() {
	if (len(os.Args) != 5) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help get'")
	}
	appName := os.Args[3]
	key := os.Args[4]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s/config_vars", appName))
	config := data.(map[string]interface{})
	value, found := config[key].(string)
	if !found {
		error(fmt.Sprintf("No such key as '%s'", key))
	}
	fmt.Println(value)
	os.Exit(0)
}

func infoHelp() {
	cmdHelp("hk info -a <app>", "Show app info")
}

func info() {
	if (len(os.Args) != 4) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help info'")
	}
	appName := os.Args[3]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s", appName))
	info := data.(map[string]interface{})
	fmt.Printf("Name:     %s\n", info["name"])
	fmt.Printf("Owner:    %s\n", info["owner_email"])
	fmt.Printf("Stack:    %s\n", info["stack"])
	fmt.Printf("Git URL:  %s\n", info["git_url"])
	fmt.Printf("Web URL:  %s\n", info["web_url"])
	os.Exit(0)
}

func listHelp() {
	cmdHelp("hk list", "List accessible apps")
}

func list() {
	if len(os.Args) != 2 {
		unrecArg(os.Args[2])
	}
	data := apiReq("GET", "https://api.heroku.com/apps")
	apps := data.([]interface{})
	for i := range apps {
		app := apps[i].(map[string]interface{})
		fmt.Printf("%s\n", app["name"])
	}
	os.Exit(0)
}

func psHelp() {
	cmdHelp("hk ps -a <app>", "List app processes")
}

func ps() {
	if (len(os.Args) != 4) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help ps'")
	}
	appName := os.Args[3]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s/ps", appName))
	processes := data.([]interface{})
	sort.Sort(ByFn{
		processes,
		func(a, b interface{}) bool {
			p1 := a.(map[string]interface{})["process"].(string)
			p2 := b.(map[string]interface{})["process"].(string)
		  return p1 < p2
	  }})
	fmt.Printf("Process           State       Command\n")
	fmt.Printf("----------------  ----------  ------------------------\n")
	for i := range processes {
		process := processes[i].(map[string]interface{})
		fmt.Printf("%-16s  %-10s  %s\n", process["process"], process["state"], process["command"])
	}
	os.Exit(0)
}

func versionHelp() {
	cmdHelp("hk version", "Show hk client version")
}

func version() {
	if len(os.Args) != 2 {
	  unrecArg(os.Args[2])
  }
	fmt.Printf("%s\n", VERSION)
	os.Exit(0)
}

func help() {
	if len(os.Args) <= 2 {
		usage()
	} else {
		cmd := os.Args[2]
		switch cmd {
	  case "env":
		  envHelp()
	  case "get":
			getHelp()
		case "info":
			infoHelp()
		case "list":
		  listHelp()
	  case "ps":
		  psHelp()
		case "version":
			versionHelp()
		}
		unrecCmd(cmd)
	}
}

// top-level usage
func usage() {
	fmt.Printf("Usage: hk <command> [-a <app>] [command-specific-options]\n\n")
	fmt.Printf("Supported hk commands are:\n")
	fmt.Printf("  addons          List add-ons\n")
	fmt.Printf("  addons-add      Add an add-on\n")
	fmt.Printf("  addons-open     Open an add-on page\n")
	fmt.Printf("  addons-remove   Remove an add-on \n")
	fmt.Printf("  create          Create an app\n")
	fmt.Printf("  destroy         Destroy an app\n")
	fmt.Printf("  env             List config vars\n")
	fmt.Printf("  get             Get config var\n")
	fmt.Printf("  help            Show this help\n")
	fmt.Printf("  info            Show app info\n")
	fmt.Printf("  list            List apps\n")
	fmt.Printf("  login           Log in\n")
	fmt.Printf("  logout          Log out\n")
	fmt.Printf("  logs            Show logs\n")
	fmt.Printf("  open            Open app\n")
	fmt.Printf("  pg              List databases\n")
	fmt.Printf("  pg-info         Show database info\n")
	fmt.Printf("  pg-promote      Promote a database\n")
	fmt.Printf("  ps-psql         Open a psql database shell\n")
	fmt.Printf("  pg-wait         Await a database\n")
	fmt.Printf("  ps              List processes\n")
	fmt.Printf("  release         Show release info\n")
	fmt.Printf("  releases        List releases\n")
	fmt.Printf("  rename          Rename an app\n")
	fmt.Printf("  restart         Restart processes\n")
	fmt.Printf("  rollback        Rollback to a previous release\n")
	fmt.Printf("  run             Run a process\n")
	fmt.Printf("  set             Set config var\n")
	fmt.Printf("  scale           Scale processes\n")
	fmt.Printf("  run             Run a process\n")
	fmt.Printf("  stop            Stop a process\n")
	fmt.Printf("  token           Show auth token\n")
	fmt.Printf("  unset           Unset config vars\n")
	fmt.Printf("  version         Display version\n\n")
	fmt.Printf("See 'hk help <command>' for more information on a specific command.\n")
	os.Exit(0)
}

// entry point
func main() {
	if len(os.Args) <= 1 {
		usage()
	} else {
		cmd := os.Args[1]
		switch cmd {
		case "env":
			env()
		case "get":
			get()
		case "help":
			help()
		case "info":
			info()
		case "list":
			list()
		case "ps":
			ps()
		case "version":
			version()
		}
		unrecCmd(cmd)
	}
}
<commit_msg>first pass at token<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
)

const (
	VERSION = "0.0.1"
)

// sort interface-typed arrays by first-class functions
type ByFn struct{
	elems []interface{}
	comp func(a, b interface{}) bool
}
func (c ByFn) Len() int { return len(c.elems)  }
func (c ByFn) Less(i, j int) bool { return c.comp(c.elems[i], c.elems[j]) }
func (c ByFn) Swap(i, j int) { c.elems[i], c.elems[j] = c.elems[j], c.elems[i] }

// generic api requests
func apiReq(meth string, url string) interface{} {
	client := &http.Client{}
	req, err := http.NewRequest(meth, url, nil)
	req.SetBasicAuth("x", os.Getenv("HEROKU_API_KEY"))
	req.Header.Add("User-Agent", fmt.Sprintf("hk/%s", VERSION))
	req.Header.Add("Accept", "application/json")
	res, err := client.Do(req)
		if err != nil {
		panic(err)
	}
	if (res.StatusCode == 401) {
		error("Unauthorized")
	}
	if (res.StatusCode == 403) {
		error("Unauthorized")
	}
	if (res.StatusCode != 200) {
		fmt.Printf("%v\n", res)
		error("Unexpected error")
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}
	var data interface{}
	err = json.Unmarshal(body, &data)
	if err != nil {
		panic(err)
	}
	return data
}

// error formatting
func error(msg string) {
	fmt.Fprintf(os.Stderr, "Error: %s.\n", msg)
	os.Exit(1)
}

func unrecArg(arg string) {
	error(fmt.Sprintf("Unrecognized argument '%s'", arg))
}

func unrecCmd(cmd string) {
	error(fmt.Sprintf("'%s' is not an hk command. See 'hk help'", cmd))
}

// info formatting
func cmdHelp(usage string, desc string) {
	fmt.Printf("Usage: %s\n\n", usage)
	fmt.Printf("%s.\n", desc)
	os.Exit(0)
}

// commands
func envHelp() {
	cmdHelp("hk env -a <app>", "Show all config vars")
}

func env() {
	if (len(os.Args) != 4) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help env'")
	}
	appName := os.Args[3]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s/config_vars", appName))
	config := data.(map[string]interface{})
	for k, v := range config {
		fmt.Printf("%s=%v\n", k, v)
	}
	os.Exit(0)
}

func getHelp() {
	cmdHelp("hk get -a <app> <key>", "Get the value of a config var")
}

func get() {
	if (len(os.Args) != 5) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help get'")
	}
	appName := os.Args[3]
	key := os.Args[4]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s/config_vars", appName))
	config := data.(map[string]interface{})
	value, found := config[key].(string)
	if !found {
		error(fmt.Sprintf("No such key as '%s'", key))
	}
	fmt.Println(value)
	os.Exit(0)
}

func infoHelp() {
	cmdHelp("hk info -a <app>", "Show app info")
}

func info() {
	if (len(os.Args) != 4) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help info'")
	}
	appName := os.Args[3]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s", appName))
	info := data.(map[string]interface{})
	fmt.Printf("Name:     %s\n", info["name"])
	fmt.Printf("Owner:    %s\n", info["owner_email"])
	fmt.Printf("Stack:    %s\n", info["stack"])
	fmt.Printf("Git URL:  %s\n", info["git_url"])
	fmt.Printf("Web URL:  %s\n", info["web_url"])
	os.Exit(0)
}

func tokenHelp() {
	cmdHelp("hk token", "Show API token")
}

func token() {
	fmt.Println(os.Getenv("HEROKU_API_KEY"))
	os.Exit(0)
}

func listHelp() {
	cmdHelp("hk list", "List accessible apps")
}

func list() {
	if len(os.Args) != 2 {
		unrecArg(os.Args[2])
	}
	data := apiReq("GET", "https://api.heroku.com/apps")
	apps := data.([]interface{})
	for i := range apps {
		app := apps[i].(map[string]interface{})
		fmt.Printf("%s\n", app["name"])
	}
	os.Exit(0)
}

func psHelp() {
	cmdHelp("hk ps -a <app>", "List app processes")
}

func ps() {
	if (len(os.Args) != 4) || (os.Args[2] != "-a") {
		error("Invalid usage. See 'hk help ps'")
	}
	appName := os.Args[3]
	data := apiReq("GET", fmt.Sprintf("https://api.heroku.com/apps/%s/ps", appName))
	processes := data.([]interface{})
	sort.Sort(ByFn{
		processes,
		func(a, b interface{}) bool {
			p1 := a.(map[string]interface{})["process"].(string)
			p2 := b.(map[string]interface{})["process"].(string)
		  return p1 < p2
	  }})
	fmt.Printf("Process           State       Command\n")
	fmt.Printf("----------------  ----------  ------------------------\n")
	for i := range processes {
		process := processes[i].(map[string]interface{})
		fmt.Printf("%-16s  %-10s  %s\n", process["process"], process["state"], process["command"])
	}
	os.Exit(0)
}

func versionHelp() {
	cmdHelp("hk version", "Show hk client version")
}

func version() {
	if len(os.Args) != 2 {
	  unrecArg(os.Args[2])
  }
	fmt.Printf("%s\n", VERSION)
	os.Exit(0)
}

func help() {
	if len(os.Args) <= 2 {
		usage()
	} else {
		cmd := os.Args[2]
		switch cmd {
	  case "env":
		  envHelp()
	  case "get":
			getHelp()
		case "info":
			infoHelp()
		case "token":
			tokenHelp()
		case "list":
		  listHelp()
	  case "ps":
		  psHelp()
		case "version":
			versionHelp()
		}
		unrecCmd(cmd)
	}
}

// top-level usage
func usage() {
	fmt.Printf("Usage: hk <command> [-a <app>] [command-specific-options]\n\n")
	fmt.Printf("Supported hk commands are:\n")
	fmt.Printf("  addons          List add-ons\n")
	fmt.Printf("  addons-add      Add an add-on\n")
	fmt.Printf("  addons-open     Open an add-on page\n")
	fmt.Printf("  addons-remove   Remove an add-on \n")
	fmt.Printf("  create          Create an app\n")
	fmt.Printf("  destroy         Destroy an app\n")
	fmt.Printf("  env             List config vars\n")
	fmt.Printf("  get             Get config var\n")
	fmt.Printf("  help            Show this help\n")
	fmt.Printf("  info            Show app info\n")
	fmt.Printf("  list            List apps\n")
	fmt.Printf("  login           Log in\n")
	fmt.Printf("  logout          Log out\n")
	fmt.Printf("  logs            Show logs\n")
	fmt.Printf("  open            Open app\n")
	fmt.Printf("  pg              List databases\n")
	fmt.Printf("  pg-info         Show database info\n")
	fmt.Printf("  pg-promote      Promote a database\n")
	fmt.Printf("  ps-psql         Open a psql database shell\n")
	fmt.Printf("  pg-wait         Await a database\n")
	fmt.Printf("  ps              List processes\n")
	fmt.Printf("  release         Show release info\n")
	fmt.Printf("  releases        List releases\n")
	fmt.Printf("  rename          Rename an app\n")
	fmt.Printf("  restart         Restart processes\n")
	fmt.Printf("  rollback        Rollback to a previous release\n")
	fmt.Printf("  run             Run a process\n")
	fmt.Printf("  set             Set config var\n")
	fmt.Printf("  scale           Scale processes\n")
	fmt.Printf("  run             Run a process\n")
	fmt.Printf("  stop            Stop a process\n")
	fmt.Printf("  token           Show auth token\n")
	fmt.Printf("  unset           Unset config vars\n")
	fmt.Printf("  version         Display version\n\n")
	fmt.Printf("See 'hk help <command>' for more information on a specific command.\n")
	os.Exit(0)
}

// entry point
func main() {
	if len(os.Args) <= 1 {
		usage()
	} else {
		cmd := os.Args[1]
		switch cmd {
		case "env":
			env()
		case "get":
			get()
		case "help":
			help()
		case "info":
			info()
		case "token":
			token()
		case "list":
			list()
		case "ps":
			ps()
		case "version":
			version()
		}
		unrecCmd(cmd)
	}
}
<|endoftext|>


<commit_before>// Package event implements an event multiplexer.
package event

import (
	"errors"
	"reflect"
	"sync"
)

// Subscription is implemented by event subscriptions.
type Subscription interface {
	// Chan returns a channel that carries events.
	// Implementations should return the same channel
	// for any subsequent calls to Chan.
	Chan() <-chan interface{}

	// Unsubscribe stops delivery of events to a subscription.
	// The event channel is closed.
	// Unsubscribe can be called more than once.
	Unsubscribe()
}

// A TypeMux dispatches events to registered receivers. Receivers can be
// registered to handle events of certain type. Any operation
// called after mux is stopped will return ErrMuxClosed.
type TypeMux struct {
	mutex   sync.RWMutex
	subm    map[reflect.Type][]*muxsub
	stopped bool
}

// ErrMuxClosed is returned when Posting on a closed TypeMux.
var ErrMuxClosed = errors.New("event: mux closed")

// NewTypeMux creates a running mux.
func NewTypeMux() *TypeMux {
	return &TypeMux{subm: make(map[reflect.Type][]*muxsub)}
}

// Subscribe creates a subscription for events of the given types. The
// subscription's channel is closed when it is unsubscribed
// or the mux is closed.
func (mux *TypeMux) Subscribe(types ...interface{}) Subscription {
	sub := newsub(mux)
	mux.mutex.Lock()
	if mux.stopped {
		mux.mutex.Unlock()
		close(sub.postC)
	} else {
		for _, t := range types {
			rtyp := reflect.TypeOf(t)
			oldsubs := mux.subm[rtyp]
			subs := make([]*muxsub, len(oldsubs)+1)
			copy(subs, oldsubs)
			subs[len(oldsubs)] = sub
			mux.subm[rtyp] = subs
		}
		mux.mutex.Unlock()
	}
	return sub
}

// Post sends an event to all receivers registered for the given type.
// It returns ErrMuxClosed if the mux has been stopped.
func (mux *TypeMux) Post(ev interface{}) error {
	rtyp := reflect.TypeOf(ev)
	mux.mutex.RLock()
	if mux.stopped {
		mux.mutex.RUnlock()
		return ErrMuxClosed
	}
	subs := mux.subm[rtyp]
	mux.mutex.RUnlock()
	for _, sub := range subs {
		sub.deliver(ev)
	}
	return nil
}

// Stop closes a mux. The mux can no longer be used.
// Future Post calls will fail with ErrMuxClosed.
// Stop blocks until all current deliveries have finished.
func (mux *TypeMux) Stop() {
	mux.mutex.Lock()
	for _, subs := range mux.subm {
		for _, sub := range subs {
			sub.closewait()
		}
	}
	mux.subm = nil
	mux.stopped = true
	mux.mutex.Unlock()
}

func (mux *TypeMux) del(s *muxsub) {
	mux.mutex.Lock()
	for typ, subs := range mux.subm {
		if pos := find(subs, s); pos >= 0 {
			if len(subs) == 1 {
				delete(mux.subm, typ)
			} else {
				mux.subm[typ] = posdelete(subs, pos)
			}
		}
	}
	s.mux.mutex.Unlock()
}

func find(slice []*muxsub, item *muxsub) int {
	for i, v := range slice {
		if v == item {
			return i
		}
	}
	return -1
}

func posdelete(slice []*muxsub, pos int) []*muxsub {
	news := make([]*muxsub, len(slice)-1)
	copy(news[:pos], slice[:pos])
	copy(news[pos:], slice[pos+1:])
	return news
}

type muxsub struct {
	mux     *TypeMux
	mutex   sync.RWMutex
	closing chan struct{}

	// these two are the same channel. they are stored separately so
	// postC can be set to nil without affecting the return value of
	// Chan.
	readC <-chan interface{}
	postC chan<- interface{}
}

func newsub(mux *TypeMux) *muxsub {
	c := make(chan interface{})
	return &muxsub{
		mux:     mux,
		readC:   c,
		postC:   c,
		closing: make(chan struct{}),
	}
}

func (s *muxsub) Chan() <-chan interface{} {
	return s.readC
}

func (s *muxsub) Unsubscribe() {
	s.mux.del(s)
	s.closewait()
}

func (s *muxsub) closewait() {
	close(s.closing)
	s.mutex.Lock()
	close(s.postC)
	s.postC = nil
	s.mutex.Unlock()
}

func (s *muxsub) deliver(ev interface{}) {
	s.mutex.RLock()
	select {
	case s.postC <- ev:
	case <-s.closing:
	}
	s.mutex.RUnlock()
}
<commit_msg>event: make Unsubscribe idempotent<commit_after>// Package event implements an event multiplexer.
package event

import (
	"errors"
	"reflect"
	"sync"
)

// Subscription is implemented by event subscriptions.
type Subscription interface {
	// Chan returns a channel that carries events.
	// Implementations should return the same channel
	// for any subsequent calls to Chan.
	Chan() <-chan interface{}

	// Unsubscribe stops delivery of events to a subscription.
	// The event channel is closed.
	// Unsubscribe can be called more than once.
	Unsubscribe()
}

// A TypeMux dispatches events to registered receivers. Receivers can be
// registered to handle events of certain type. Any operation
// called after mux is stopped will return ErrMuxClosed.
type TypeMux struct {
	mutex   sync.RWMutex
	subm    map[reflect.Type][]*muxsub
	stopped bool
}

// ErrMuxClosed is returned when Posting on a closed TypeMux.
var ErrMuxClosed = errors.New("event: mux closed")

// NewTypeMux creates a running mux.
func NewTypeMux() *TypeMux {
	return &TypeMux{subm: make(map[reflect.Type][]*muxsub)}
}

// Subscribe creates a subscription for events of the given types. The
// subscription's channel is closed when it is unsubscribed
// or the mux is closed.
func (mux *TypeMux) Subscribe(types ...interface{}) Subscription {
	sub := newsub(mux)
	mux.mutex.Lock()
	if mux.stopped {
		mux.mutex.Unlock()
		close(sub.postC)
	} else {
		for _, t := range types {
			rtyp := reflect.TypeOf(t)
			oldsubs := mux.subm[rtyp]
			subs := make([]*muxsub, len(oldsubs)+1)
			copy(subs, oldsubs)
			subs[len(oldsubs)] = sub
			mux.subm[rtyp] = subs
		}
		mux.mutex.Unlock()
	}
	return sub
}

// Post sends an event to all receivers registered for the given type.
// It returns ErrMuxClosed if the mux has been stopped.
func (mux *TypeMux) Post(ev interface{}) error {
	rtyp := reflect.TypeOf(ev)
	mux.mutex.RLock()
	if mux.stopped {
		mux.mutex.RUnlock()
		return ErrMuxClosed
	}
	subs := mux.subm[rtyp]
	mux.mutex.RUnlock()
	for _, sub := range subs {
		sub.deliver(ev)
	}
	return nil
}

// Stop closes a mux. The mux can no longer be used.
// Future Post calls will fail with ErrMuxClosed.
// Stop blocks until all current deliveries have finished.
func (mux *TypeMux) Stop() {
	mux.mutex.Lock()
	for _, subs := range mux.subm {
		for _, sub := range subs {
			sub.closewait()
		}
	}
	mux.subm = nil
	mux.stopped = true
	mux.mutex.Unlock()
}

func (mux *TypeMux) del(s *muxsub) {
	mux.mutex.Lock()
	for typ, subs := range mux.subm {
		if pos := find(subs, s); pos >= 0 {
			if len(subs) == 1 {
				delete(mux.subm, typ)
			} else {
				mux.subm[typ] = posdelete(subs, pos)
			}
		}
	}
	s.mux.mutex.Unlock()
}

func find(slice []*muxsub, item *muxsub) int {
	for i, v := range slice {
		if v == item {
			return i
		}
	}
	return -1
}

func posdelete(slice []*muxsub, pos int) []*muxsub {
	news := make([]*muxsub, len(slice)-1)
	copy(news[:pos], slice[:pos])
	copy(news[pos:], slice[pos+1:])
	return news
}

type muxsub struct {
	mux     *TypeMux
	closeMu sync.Mutex
	closing chan struct{}
	closed  bool

	// these two are the same channel. they are stored separately so
	// postC can be set to nil without affecting the return value of
	// Chan.
	postMu sync.RWMutex
	readC  <-chan interface{}
	postC  chan<- interface{}
}

func newsub(mux *TypeMux) *muxsub {
	c := make(chan interface{})
	return &muxsub{
		mux:     mux,
		readC:   c,
		postC:   c,
		closing: make(chan struct{}),
	}
}

func (s *muxsub) Chan() <-chan interface{} {
	return s.readC
}

func (s *muxsub) Unsubscribe() {
	s.mux.del(s)
	s.closewait()
}

func (s *muxsub) closewait() {
	s.closeMu.Lock()
	defer s.closeMu.Unlock()
	if s.closed {
		return
	}
	close(s.closing)
	s.closed = true

	s.postMu.Lock()
	close(s.postC)
	s.postC = nil
	s.postMu.Unlock()
}

func (s *muxsub) deliver(ev interface{}) {
	s.postMu.RLock()
	select {
	case s.postC <- ev:
	case <-s.closing:
	}
	s.postMu.RUnlock()
}
<|endoftext|>


<commit_before>package routes

import (
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"golang.org/x/net/context"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestValidateToken(t *testing.T) {
	req, err := http.NewRequest("GET", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("Authorization", "Bearer abc123")

	rec := httptest.NewRecorder()

	ctx := req.Context()
	storage := new(MockStorage)

	storage.On("RetrieveUsernameFromToken", "abc123").Return("foo", nil)

	ctx = context.WithValue(ctx, "storage", storage)

	ValidateToken(rec, req.WithContext(ctx), func(w http.ResponseWriter, req *http.Request) {
		assert.Equal(t, rec.Code, http.StatusOK)
		assert.Equal(t, req.Context().Value("user").(string), "foo")
		assert.Equal(t, req.Context().Value("token").(string), "abc123")
	})
}

type MockStorage struct {
	mock.Mock
}

func (s *MockStorage) RetrieveUsernameFromToken(token string) (string, error) {
	args := s.Called(token)
	return args.String(0), args.Error(1)
}
<commit_msg>Adds negative test case for ValidateToken middleware<commit_after>package routes

import (
	"errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"golang.org/x/net/context"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestValidateToken(t *testing.T) {
	req, err := http.NewRequest("GET", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("Authorization", "Bearer abc123")

	rec := httptest.NewRecorder()

	ctx := req.Context()
	storage := new(MockStorage)

	storage.On("RetrieveUsernameFromToken", "abc123").Return("foo", nil)

	ctx = context.WithValue(ctx, "storage", storage)

	ValidateToken(rec, req.WithContext(ctx), func(w http.ResponseWriter, req *http.Request) {
		assert.Equal(t, rec.Code, http.StatusOK)
		assert.Equal(t, req.Context().Value("user").(string), "foo")
		assert.Equal(t, req.Context().Value("token").(string), "abc123")
	})
}

func TestValidateTokenWithInvalidToken(t *testing.T) {
	req, err := http.NewRequest("GET", "", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("Authorization", "Bearer abc123")

	rec := httptest.NewRecorder()

	ctx := req.Context()
	storage := new(MockStorage)

	storage.
		On("RetrieveUsernameFromToken", "abc123").
		Return("", errors.New("oh no!"))

	ctx = context.WithValue(ctx, "storage", storage)

	ValidateToken(rec, req.WithContext(ctx), func(w http.ResponseWriter, req *http.Request) {
		t.Fail()
	})

	assert.Equal(t, rec.Code, http.StatusUnauthorized)
}

type MockStorage struct {
	mock.Mock
}

func (s *MockStorage) RetrieveUsernameFromToken(token string) (string, error) {
	args := s.Called(token)
	return args.String(0), args.Error(1)
}
<|endoftext|>


<commit_before>package gkvlite

import (
	"sync"
	"sync/atomic"
	"unsafe"
)

var freeNodeLock sync.Mutex
var freeNodes *node

var freeNodeLocLock sync.Mutex
var freeNodeLocs *nodeLoc

var freeRootNodeLocLock sync.Mutex
var freeRootNodeLocs *rootNodeLoc

var allocStats AllocStats

type AllocStats struct {
	MkNodes      int64
	FreeNodes    int64 // Number of invocations of the freeNode() API.
	AllocNodes   int64
	CurFreeNodes int64 // Current length of freeNodes list.

	MkNodeLocs      int64
	FreeNodeLocs    int64 // Number of invocations of the freeNodeLoc() API.
	AllocNodeLocs   int64
	CurFreeNodeLocs int64 // Current length of freeNodeLocs list.

	MkRootNodeLocs      int64
	FreeRootNodeLocs    int64 // Number of invocations of the freeRootNodeLoc() API.
	AllocRootNodeLocs   int64
	CurFreeRootNodeLocs int64 // Current length of freeRootNodeLocs list.
}

func withAllocLocks(cb func()) {
	freeNodeLock.Lock()
	freeNodeLocLock.Lock()
	freeRootNodeLocLock.Lock()
	defer freeNodeLock.Unlock()
	defer freeNodeLocLock.Unlock()
	defer freeRootNodeLocLock.Unlock()
	cb()
}

func (t *Collection) markReclaimable(n *node, reclaimMark *node) {
	t.rootLock.Lock()
	defer t.rootLock.Unlock()
	if n == nil || n.next != nil || n == reclaimMark {
		return
	}
	n.next = reclaimMark
}

func (t *Collection) reclaimNodes_unlocked(n *node,
	reclaimLater *[2]*node, reclaimMark *node) int64 {
	if n == nil {
		return 0
	}
	if reclaimLater != nil {
		for i := 0; i < len(reclaimLater); i++ {
			if reclaimLater[i] == n {
				reclaimLater[i] = nil
			}
		}
	}
	if n.next != reclaimMark {
		return 0
	}
	var left *node
	var right *node
	if !n.left.isEmpty() {
		left = n.left.Node()
	}
	if !n.right.isEmpty() {
		right = n.right.Node()
	}
	t.freeNode_unlocked(n, reclaimMark)
	numLeft := t.reclaimNodes_unlocked(left, reclaimLater, reclaimMark)
	numRight := t.reclaimNodes_unlocked(right, reclaimLater, reclaimMark)
	return 1 + numLeft + numRight
}

// Assumes that the caller serializes invocations.
func (t *Collection) mkNode(itemIn *itemLoc, leftIn *nodeLoc, rightIn *nodeLoc,
	numNodesIn uint64, numBytesIn uint64) *node {
	freeNodeLock.Lock()
	allocStats.MkNodes++
	t.allocStats.MkNodes++
	n := freeNodes
	if n == nil {
		allocStats.AllocNodes++
		t.allocStats.AllocNodes++
		freeNodeLock.Unlock()
		atomic.AddUint64(&t.store.nodeAllocs, 1)
		n = &node{}
	} else {
		freeNodes = n.next
		allocStats.CurFreeNodes--
		freeNodeLock.Unlock()
	}
	n.item.Copy(itemIn)
	n.left.Copy(leftIn)
	n.right.Copy(rightIn)
	n.numNodes = numNodesIn
	n.numBytes = numBytesIn
	n.next = nil
	return n
}

func (t *Collection) freeNode_unlocked(n *node, reclaimMark *node) {
	if n == nil || n == reclaimMark {
		return
	}
	if n.next != nil && n.next != reclaimMark {
		panic("double free node")
	}
	n.item = *empty_itemLoc
	n.left = *empty_nodeLoc
	n.right = *empty_nodeLoc
	n.numNodes = 0
	n.numBytes = 0

	n.next = freeNodes
	freeNodes = n
	allocStats.CurFreeNodes++
	allocStats.FreeNodes++
	t.allocStats.FreeNodes++
}

// Assumes that the caller serializes invocations.
func (t *Collection) mkNodeLoc(n *node) *nodeLoc {
	freeNodeLocLock.Lock()
	allocStats.MkNodeLocs++
	t.allocStats.MkNodeLocs++
	nloc := freeNodeLocs
	if nloc == nil {
		allocStats.AllocNodeLocs++
		t.allocStats.AllocNodeLocs++
		freeNodeLocLock.Unlock()
		nloc = &nodeLoc{}
	} else {
		freeNodeLocs = nloc.next
		allocStats.CurFreeNodeLocs--
		freeNodeLocLock.Unlock()
	}
	nloc.loc = unsafe.Pointer(nil)
	nloc.node = unsafe.Pointer(n)
	nloc.next = nil
	return nloc
}

// Assumes that the caller serializes invocations.
func (t *Collection) freeNodeLoc(nloc *nodeLoc) {
	if nloc == nil || nloc == empty_nodeLoc {
		return
	}
	if nloc.next != nil {
		panic("double free nodeLoc")
	}
	nloc.loc = unsafe.Pointer(nil)
	nloc.node = unsafe.Pointer(nil)

	freeNodeLocLock.Lock()
	nloc.next = freeNodeLocs
	freeNodeLocs = nloc
	allocStats.CurFreeNodeLocs++
	allocStats.FreeNodeLocs++
	t.allocStats.FreeNodeLocs++
	freeNodeLocLock.Unlock()
}

func (t *Collection) mkRootNodeLoc(root *nodeLoc) *rootNodeLoc {
	freeRootNodeLocLock.Lock()
	allocStats.MkRootNodeLocs++
	t.allocStats.MkRootNodeLocs++
	rnl := freeRootNodeLocs
	if rnl == nil {
		allocStats.AllocRootNodeLocs++
		t.allocStats.AllocRootNodeLocs++
		freeRootNodeLocLock.Unlock()
		rnl = &rootNodeLoc{}
	} else {
		freeRootNodeLocs = rnl.next
		allocStats.CurFreeRootNodeLocs--
		freeRootNodeLocLock.Unlock()
	}
	rnl.refs = 1
	rnl.root = root
	rnl.next = nil
	rnl.chainedCollection = nil
	rnl.chainedRootNodeLoc = nil
	for i := 0; i < len(rnl.reclaimLater); i++ {
		rnl.reclaimLater[i] = nil
	}
	return rnl
}

func (t *Collection) freeRootNodeLoc(rnl *rootNodeLoc) {
	if rnl == nil {
		return
	}
	if rnl.next != nil {
		panic("double free rootNodeLoc")
	}
	rnl.refs = 0
	rnl.root = nil
	rnl.chainedCollection = nil
	rnl.chainedRootNodeLoc = nil
	for i := 0; i < len(rnl.reclaimLater); i++ {
		rnl.reclaimLater[i] = nil
	}
	freeRootNodeLocLock.Lock()
	rnl.next = freeRootNodeLocs
	freeRootNodeLocs = rnl
	allocStats.CurFreeRootNodeLocs++
	allocStats.FreeRootNodeLocs++
	t.allocStats.FreeRootNodeLocs++
	freeRootNodeLocLock.Unlock()
}
<commit_msg>mkNode() does ItemValAddRef().<commit_after>package gkvlite

import (
	"sync"
	"sync/atomic"
	"unsafe"
)

var freeNodeLock sync.Mutex
var freeNodes *node

var freeNodeLocLock sync.Mutex
var freeNodeLocs *nodeLoc

var freeRootNodeLocLock sync.Mutex
var freeRootNodeLocs *rootNodeLoc

var allocStats AllocStats

type AllocStats struct {
	MkNodes      int64
	FreeNodes    int64 // Number of invocations of the freeNode() API.
	AllocNodes   int64
	CurFreeNodes int64 // Current length of freeNodes list.

	MkNodeLocs      int64
	FreeNodeLocs    int64 // Number of invocations of the freeNodeLoc() API.
	AllocNodeLocs   int64
	CurFreeNodeLocs int64 // Current length of freeNodeLocs list.

	MkRootNodeLocs      int64
	FreeRootNodeLocs    int64 // Number of invocations of the freeRootNodeLoc() API.
	AllocRootNodeLocs   int64
	CurFreeRootNodeLocs int64 // Current length of freeRootNodeLocs list.
}

func withAllocLocks(cb func()) {
	freeNodeLock.Lock()
	freeNodeLocLock.Lock()
	freeRootNodeLocLock.Lock()
	defer freeNodeLock.Unlock()
	defer freeNodeLocLock.Unlock()
	defer freeRootNodeLocLock.Unlock()
	cb()
}

func (t *Collection) markReclaimable(n *node, reclaimMark *node) {
	t.rootLock.Lock()
	defer t.rootLock.Unlock()
	if n == nil || n.next != nil || n == reclaimMark {
		return
	}
	n.next = reclaimMark
}

func (t *Collection) reclaimNodes_unlocked(n *node,
	reclaimLater *[2]*node, reclaimMark *node) int64 {
	if n == nil {
		return 0
	}
	if reclaimLater != nil {
		for i := 0; i < len(reclaimLater); i++ {
			if reclaimLater[i] == n {
				reclaimLater[i] = nil
			}
		}
	}
	if n.next != reclaimMark {
		return 0
	}
	var left *node
	var right *node
	if !n.left.isEmpty() {
		left = n.left.Node()
	}
	if !n.right.isEmpty() {
		right = n.right.Node()
	}
	t.freeNode_unlocked(n, reclaimMark)
	numLeft := t.reclaimNodes_unlocked(left, reclaimLater, reclaimMark)
	numRight := t.reclaimNodes_unlocked(right, reclaimLater, reclaimMark)
	return 1 + numLeft + numRight
}

// Assumes that the caller serializes invocations.
func (t *Collection) mkNode(itemIn *itemLoc, leftIn *nodeLoc, rightIn *nodeLoc,
	numNodesIn uint64, numBytesIn uint64) *node {
	freeNodeLock.Lock()
	allocStats.MkNodes++
	t.allocStats.MkNodes++
	n := freeNodes
	if n == nil {
		allocStats.AllocNodes++
		t.allocStats.AllocNodes++
		freeNodeLock.Unlock()
		atomic.AddUint64(&t.store.nodeAllocs, 1)
		n = &node{}
	} else {
		freeNodes = n.next
		allocStats.CurFreeNodes--
		freeNodeLock.Unlock()
	}
	if itemIn != nil {
		i := itemIn.Item()
		if i != nil {
			t.store.ItemValAddRef(t, i)
		}
	}
	n.item.Copy(itemIn)
	n.left.Copy(leftIn)
	n.right.Copy(rightIn)
	n.numNodes = numNodesIn
	n.numBytes = numBytesIn
	n.next = nil
	return n
}

func (t *Collection) freeNode_unlocked(n *node, reclaimMark *node) {
	if n == nil || n == reclaimMark {
		return
	}
	if n.next != nil && n.next != reclaimMark {
		panic("double free node")
	}
	n.item = *empty_itemLoc
	n.left = *empty_nodeLoc
	n.right = *empty_nodeLoc
	n.numNodes = 0
	n.numBytes = 0

	n.next = freeNodes
	freeNodes = n
	allocStats.CurFreeNodes++
	allocStats.FreeNodes++
	t.allocStats.FreeNodes++
}

// Assumes that the caller serializes invocations.
func (t *Collection) mkNodeLoc(n *node) *nodeLoc {
	freeNodeLocLock.Lock()
	allocStats.MkNodeLocs++
	t.allocStats.MkNodeLocs++
	nloc := freeNodeLocs
	if nloc == nil {
		allocStats.AllocNodeLocs++
		t.allocStats.AllocNodeLocs++
		freeNodeLocLock.Unlock()
		nloc = &nodeLoc{}
	} else {
		freeNodeLocs = nloc.next
		allocStats.CurFreeNodeLocs--
		freeNodeLocLock.Unlock()
	}
	nloc.loc = unsafe.Pointer(nil)
	nloc.node = unsafe.Pointer(n)
	nloc.next = nil
	return nloc
}

// Assumes that the caller serializes invocations.
func (t *Collection) freeNodeLoc(nloc *nodeLoc) {
	if nloc == nil || nloc == empty_nodeLoc {
		return
	}
	if nloc.next != nil {
		panic("double free nodeLoc")
	}
	nloc.loc = unsafe.Pointer(nil)
	nloc.node = unsafe.Pointer(nil)

	freeNodeLocLock.Lock()
	nloc.next = freeNodeLocs
	freeNodeLocs = nloc
	allocStats.CurFreeNodeLocs++
	allocStats.FreeNodeLocs++
	t.allocStats.FreeNodeLocs++
	freeNodeLocLock.Unlock()
}

func (t *Collection) mkRootNodeLoc(root *nodeLoc) *rootNodeLoc {
	freeRootNodeLocLock.Lock()
	allocStats.MkRootNodeLocs++
	t.allocStats.MkRootNodeLocs++
	rnl := freeRootNodeLocs
	if rnl == nil {
		allocStats.AllocRootNodeLocs++
		t.allocStats.AllocRootNodeLocs++
		freeRootNodeLocLock.Unlock()
		rnl = &rootNodeLoc{}
	} else {
		freeRootNodeLocs = rnl.next
		allocStats.CurFreeRootNodeLocs--
		freeRootNodeLocLock.Unlock()
	}
	rnl.refs = 1
	rnl.root = root
	rnl.next = nil
	rnl.chainedCollection = nil
	rnl.chainedRootNodeLoc = nil
	for i := 0; i < len(rnl.reclaimLater); i++ {
		rnl.reclaimLater[i] = nil
	}
	return rnl
}

func (t *Collection) freeRootNodeLoc(rnl *rootNodeLoc) {
	if rnl == nil {
		return
	}
	if rnl.next != nil {
		panic("double free rootNodeLoc")
	}
	rnl.refs = 0
	rnl.root = nil
	rnl.chainedCollection = nil
	rnl.chainedRootNodeLoc = nil
	for i := 0; i < len(rnl.reclaimLater); i++ {
		rnl.reclaimLater[i] = nil
	}
	freeRootNodeLocLock.Lock()
	rnl.next = freeRootNodeLocs
	freeRootNodeLocs = rnl
	allocStats.CurFreeRootNodeLocs++
	allocStats.FreeRootNodeLocs++
	t.allocStats.FreeRootNodeLocs++
	freeRootNodeLocLock.Unlock()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/bwmarrin/discordgo"
	"strconv"
	"strings"
)

func (b *Bot) onDiscordReady(s *discordgo.Session, event *discordgo.Ready) {
	_ = s.UpdateStatus(0, "Starbase Online")
}

func (b *Bot) onDiscordGuildCreate(s *discordgo.Session, event *discordgo.GuildCreate) {
	if event.Guild.Unavailable {
		log.WithField("guildID", event.Guild.ID).Debug("Received GuildCreate event for unavaiable guild")
		return
	}

	targetFound := false
	if strings.EqualFold(event.Guild.ID, b.config.Discord.GuildID) {
		for _, channel := range event.Guild.Channels {
			if strings.EqualFold(channel.ID, b.config.Discord.ChannelID) {
				targetFound = true
				if b.config.Discord.Debug {
					s.ChannelMessageSend(channel.ID, ":robot: POSbot online and ready to serve :rocket:")
				}
			}
		}
	}

	if !targetFound {
		log.WithFields(logrus.Fields{
			"guildID":   b.config.Discord.GuildID,
			"channelID": b.config.Discord.ChannelID,
		}).Fatal("Target Discord guild or channel not found")
	}
}

func (b *Bot) onDiscordMessageCreate(s *discordgo.Session, message *discordgo.MessageCreate) {
	if !b.isRelevantDiscordChannel(s, message) {
		return
	}

	if strings.HasPrefix(strings.ToLower(message.Content), "!fuel") {
		b.handleDiscordFuelCommand(message)
	} else if strings.HasPrefix(strings.ToLower(message.Content), "!pos") {
		b.handleDiscordPOSCommand(message)
	}
}

func (b *Bot) getChannelFromMessage(s *discordgo.Session, message *discordgo.MessageCreate) (*discordgo.Channel, error) {
	return s.Channel(message.ChannelID)
}

func (b *Bot) isRelevantDiscordChannel(s *discordgo.Session, message *discordgo.MessageCreate) bool {
	channel, err := b.getChannelFromMessage(s, message)
	if err != nil {
		log.WithFields(logrus.Fields{
			"channelID": message.ChannelID,
			"author":    message.Author.Username,
		}).Debug("Failed to get channel from message")
		return false
	}

	return strings.EqualFold(channel.GuildID, b.config.Discord.GuildID) && strings.EqualFold(channel.ID, b.config.Discord.ChannelID)
}

func (b *Bot) handleDiscordFuelCommand(message *discordgo.MessageCreate) {
	b.discord.ChannelMessageSend(message.ChannelID, "How the fuck would I know? ¯\\_(ツ)_/¯")
}

func (b *Bot) handleDiscordPOSCommand(message *discordgo.MessageCreate) {
	messageParts := strings.Split(message.Content, " ")

	if len(messageParts) == 1 || (len(messageParts) == 2 && strings.EqualFold(messageParts[1], "list")) {
		log.WithField("author", message.Author.Username).Debug("Retrieving POS list for Discord command")

		starbases, err := b.retrieveStarbaseList()
		if err != nil {
			log.WithError(err).Warn("Failed to retrieve POS list")
			b.discord.ChannelMessageSend(message.ChannelID, ":poop: Looks like there was an error retrieving the POS list :frowning:")
			return
		}

		b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("Would you kindly fuck off, please? :innocent: BTW, there's %d POSes", len(starbases.Starbases)))
		log.WithField("author", message.Author.Username).Debug("Retrieved POS list for Discord command")
	} else if len(messageParts) >= 2 {
		starbaseID, err := strconv.ParseInt(messageParts[1], 10, 64)
		if err != nil || starbaseID <= 0 {
			log.WithField("starbaseID", starbaseID).WithError(err).Debug("Failed to parse starbaseID for Discord POS command")
			b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("<@%s>: Seems like you've provided an invalid POS ID :poop:", message.Author.ID))
			return
		}

		log.WithFields(logrus.Fields{
			"author":     message.Author.Username,
			"starbaseID": starbaseID,
		}).Debug("Retrieving POS details for Discord command")

		starbase, err := b.retrieveStarbaseDetails(int(starbaseID))
		if err != nil {
			log.WithError(err).Warn("Failed to retrieve POS details")
			b.discord.ChannelMessageSend(message.ChannelID, ":poop: Looks like there was an error retrieving the POS details :frowning:")
			return
		}

		b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("Amazing. State: %s", starbase.State))
		log.WithFields(logrus.Fields{
			"author":     message.Author.Username,
			"starbaseID": starbaseID,
		}).Debug("Retrieved POS details for Discord command")
	}
}
<commit_msg>Added "typing" indicator at the beginning of POS commands to indicate loading times Slight restructure of POS subcommands<commit_after>package main

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/bwmarrin/discordgo"
	"strconv"
	"strings"
)

func (b *Bot) onDiscordReady(s *discordgo.Session, event *discordgo.Ready) {
	_ = s.UpdateStatus(0, "Starbase Online")
}

func (b *Bot) onDiscordGuildCreate(s *discordgo.Session, event *discordgo.GuildCreate) {
	if event.Guild.Unavailable {
		log.WithField("guildID", event.Guild.ID).Debug("Received GuildCreate event for unavaiable guild")
		return
	}

	targetFound := false
	if strings.EqualFold(event.Guild.ID, b.config.Discord.GuildID) {
		for _, channel := range event.Guild.Channels {
			if strings.EqualFold(channel.ID, b.config.Discord.ChannelID) {
				targetFound = true
				if b.config.Discord.Debug {
					s.ChannelMessageSend(channel.ID, ":robot: POSbot online and ready to serve :rocket:")
				}
			}
		}
	}

	if !targetFound {
		log.WithFields(logrus.Fields{
			"guildID":   b.config.Discord.GuildID,
			"channelID": b.config.Discord.ChannelID,
		}).Fatal("Target Discord guild or channel not found")
	}
}

func (b *Bot) onDiscordMessageCreate(s *discordgo.Session, message *discordgo.MessageCreate) {
	if !b.isRelevantDiscordChannel(s, message) {
		return
	}

	if strings.HasPrefix(strings.ToLower(message.Content), "!fuel") {
		b.handleDiscordFuelCommand(message)
	} else if strings.HasPrefix(strings.ToLower(message.Content), "!pos") {
		b.handleDiscordPOSCommand(message)
	}
}

func (b *Bot) getChannelFromMessage(s *discordgo.Session, message *discordgo.MessageCreate) (*discordgo.Channel, error) {
	return s.Channel(message.ChannelID)
}

func (b *Bot) isRelevantDiscordChannel(s *discordgo.Session, message *discordgo.MessageCreate) bool {
	channel, err := b.getChannelFromMessage(s, message)
	if err != nil {
		log.WithFields(logrus.Fields{
			"channelID": message.ChannelID,
			"author":    message.Author.Username,
		}).Debug("Failed to get channel from message")
		return false
	}

	return strings.EqualFold(channel.GuildID, b.config.Discord.GuildID) && strings.EqualFold(channel.ID, b.config.Discord.ChannelID)
}

func (b *Bot) handleDiscordFuelCommand(message *discordgo.MessageCreate) {
	b.discord.ChannelMessageSend(message.ChannelID, "How the fuck would I know? ¯\\_(ツ)_/¯")
}

func (b *Bot) handleDiscordPOSCommand(message *discordgo.MessageCreate) {
	b.discord.ChannelTyping(message.ChannelID)

	messageParts := strings.Split(message.Content, " ")

	if len(messageParts) == 1 || (len(messageParts) == 2 && strings.EqualFold(messageParts[1], "list")) {
		log.WithField("author", message.Author.Username).Debug("Retrieving POS list for Discord command")

		starbases, err := b.retrieveStarbaseList()
		if err != nil {
			log.WithError(err).Warn("Failed to retrieve POS list")
			b.discord.ChannelMessageSend(message.ChannelID, ":poop: Looks like there was an error retrieving the POS list :frowning:")
			return
		}

		b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("Would you kindly fuck off, please? :innocent: BTW, there's %d POSes", len(starbases.Starbases)))
		log.WithField("author", message.Author.Username).Debug("Retrieved POS list for Discord command")
		return
	} else if len(messageParts) >= 2 {
		if strings.EqualFold(messageParts[1], "details") {
			if len(messageParts) < 3 {
				b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("<@%s>: You'll have to tell me which POS you want to know more about...", message.Author.ID))
				return
			}

			starbaseID, err := strconv.ParseInt(messageParts[2], 10, 64)
			if err != nil || starbaseID <= 0 {
				log.WithField("starbaseID", starbaseID).WithError(err).Debug("Failed to parse starbaseID for Discord POS command")
				b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("<@%s>: Seems like you've provided an invalid POS ID :poop:", message.Author.ID))
				return
			}

			log.WithFields(logrus.Fields{
				"author":     message.Author.Username,
				"starbaseID": starbaseID,
			}).Debug("Retrieving POS details for Discord command")

			starbase, err := b.retrieveStarbaseDetails(int(starbaseID))
			if err != nil {
				log.WithError(err).Warn("Failed to retrieve POS details")
				b.discord.ChannelMessageSend(message.ChannelID, ":poop: Looks like there was an error retrieving the POS details :frowning:")
				return
			}

			b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("Amazing. State: %s", starbase.State))
			log.WithFields(logrus.Fields{
				"author":     message.Author.Username,
				"starbaseID": starbaseID,
			}).Debug("Retrieved POS details for Discord command")
			return
		}
	}

	b.discord.ChannelMessageSend(message.ChannelID, fmt.Sprintf("<@%s> seems to be drunk, there's no command like this :thinking:", message.Author.ID))
}
<|endoftext|>


<commit_before>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains high level helper functions and easy entry points for the
// entire discordgo package.  These functions are beling developed and are very
// experimental at this point.  They will most likley change so please use the
// low level functions if that's a problem.

// Package discordgo provides Discord binding for Go
package discordgo

import (
	"fmt"
	"reflect"
)

// VERSION of Discordgo, follows Symantic Versioning. (http://semver.org/)
const VERSION = "0.13.0-alpha"

// New creates a new Discord session and will automate some startup
// tasks if given enough information to do so.  Currently you can pass zero
// arguments and it will return an empty Discord session.
// There are 3 ways to call New:
//     With a single auth token - All requests will use the token blindly,
//         no verification of the token will be done and requests may fail.
//     With an email and password - Discord will sign in with the provided
//         credentials.
//     With an email, password and auth token - Discord will verify the auth
//         token, if it is invalid it will sign in with the provided
//         credentials. This is the Discord recommended way to sign in.
func New(args ...interface{}) (s *Session, err error) {

	// Create an empty Session interface.
	s = &Session{
		State:                  NewState(),
		StateEnabled:           true,
		Compress:               true,
		ShouldReconnectOnError: true,
	}

	// If no arguments are passed return the empty Session interface.
	if args == nil {
		return
	}

	// Variables used below when parsing func arguments
	var auth, pass string

	// Parse passed arguments
	for _, arg := range args {

		switch v := arg.(type) {

		case []string:
			if len(v) > 3 {
				err = fmt.Errorf("Too many string parameters provided.")
				return
			}

			// First string is either token or username
			if len(v) > 0 {
				auth = v[0]
			}

			// If second string exists, it must be a password.
			if len(v) > 1 {
				pass = v[1]
			}

			// If third string exists, it must be an auth token.
			if len(v) > 2 {
				s.Token = v[2]
			}

		case string:
			// First string must be either auth token or username.
			// Second string must be a password.
			// Only 2 input strings are supported.

			if auth == "" {
				auth = v
			} else if pass == "" {
				pass = v
			} else if s.Token == "" {
				s.Token = v
			} else {
				err = fmt.Errorf("Too many string parameters provided.")
				return
			}

			//		case Config:
			// TODO: Parse configuration struct

		default:
			err = fmt.Errorf("Unsupported parameter type provided.")
			return
		}
	}

	// If only one string was provided, assume it is an auth token.
	// Otherwise get auth token from Discord, if a token was specified
	// Discord will verify it for free, or log the user in if it is
	// invalid.
	if pass == "" {
		s.Token = auth
	} else {
		err = s.Login(auth, pass)
		if err != nil || s.Token == "" {
			err = fmt.Errorf("Unable to fetch discord authentication token. %v", err)
			return
		}
	}

	// The Session is now able to have RestAPI methods called on it.
	// It is recommended that you now call Open() so that events will trigger.

	return
}

// validateHandler takes an event handler func, and returns the type of event.
// eg.
//     Session.validateHandler(func (s *discordgo.Session, m *discordgo.MessageCreate))
//     will return the reflect.Type of *discordgo.MessageCreate
func (s *Session) validateHandler(handler interface{}) reflect.Type {

	handlerType := reflect.TypeOf(handler)

	if handlerType.NumIn() != 2 {
		panic("Unable to add event handler, handler must be of the type func(*discordgo.Session, *discordgo.EventType).")
	}

	if handlerType.In(0) != reflect.TypeOf(s) {
		panic("Unable to add event handler, first argument must be of type *discordgo.Session.")
	}

	eventType := handlerType.In(1)

	// Support handlers of type interface{}, this is a special handler, which is triggered on every event.
	if eventType.Kind() == reflect.Interface {
		eventType = nil
	}

	return eventType
}

// AddHandler allows you to add an event handler that will be fired anytime
// the Discord WSAPI event that matches the interface fires.
// eventToInterface in events.go has a list of all the Discord WSAPI events
// and their respective interface.
// eg:
//     Session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
//     })
//
// or:
//     Session.AddHandler(func(s *discordgo.Session, m *discordgo.PresenceUpdate) {
//     })
// The return value of this method is a function, that when called will remove the
// event handler.
func (s *Session) AddHandler(handler interface{}) func() {

	s.initialize()

	eventType := s.validateHandler(handler)

	s.handlersMu.Lock()
	defer s.handlersMu.Unlock()

	h := reflect.ValueOf(handler)

	s.handlers[eventType] = append(s.handlers[eventType], h)

	// This must be done as we need a consistent reference to the
	// reflected value, otherwise a RemoveHandler method would have
	// been nice.
	return func() {
		s.handlersMu.Lock()
		defer s.handlersMu.Unlock()

		handlers := s.handlers[eventType]
		for i, v := range handlers {
			if h == v {
				s.handlers[eventType] = append(handlers[:i], handlers[i+1:]...)
				return
			}
		}
	}
}

// handle calls any handlers that match the event type and any handlers of
// interface{}.
func (s *Session) handle(event interface{}) {

	s.handlersMu.RLock()
	defer s.handlersMu.RUnlock()

	if s.handlers == nil {
		return
	}

	handlerParameters := []reflect.Value{reflect.ValueOf(s), reflect.ValueOf(event)}

	if handlers, ok := s.handlers[nil]; ok {
		for _, handler := range handlers {
			handler.Call(handlerParameters)
		}
	}

	if handlers, ok := s.handlers[reflect.TypeOf(event)]; ok {
		for _, handler := range handlers {
			handler.Call(handlerParameters)
		}
	}
}

// initialize adds all internal handlers and state tracking handlers.
func (s *Session) initialize() {

	s.handlersMu.Lock()
	if s.handlers != nil {
		s.handlersMu.Unlock()
		return
	}

	s.handlers = map[interface{}][]reflect.Value{}
	s.handlersMu.Unlock()

	s.AddHandler(s.onReady)
	s.AddHandler(s.onResumed)
	s.AddHandler(s.onVoiceServerUpdate)
	s.AddHandler(s.onVoiceStateUpdate)
	s.AddHandler(s.State.onInterface)
}

// onReady handles the ready event.
func (s *Session) onReady(se *Session, r *Ready) {

	// Store the SessionID within the Session struct.
	s.sessionID = r.SessionID

	// Start the heartbeat to keep the connection alive.
	go s.heartbeat(s.wsConn, s.listening, r.HeartbeatInterval)
}

// onResumed handles the resumed event.
func (s *Session) onResumed(se *Session, r *Resumed) {

	// Start the heartbeat to keep the connection alive.
	go s.heartbeat(s.wsConn, s.listening, r.HeartbeatInterval)
}
<commit_msg>Fix possible mutex deadlock in handle()<commit_after>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains high level helper functions and easy entry points for the
// entire discordgo package.  These functions are beling developed and are very
// experimental at this point.  They will most likley change so please use the
// low level functions if that's a problem.

// Package discordgo provides Discord binding for Go
package discordgo

import (
	"fmt"
	"reflect"
)

// VERSION of Discordgo, follows Symantic Versioning. (http://semver.org/)
const VERSION = "0.13.0-alpha"

// New creates a new Discord session and will automate some startup
// tasks if given enough information to do so.  Currently you can pass zero
// arguments and it will return an empty Discord session.
// There are 3 ways to call New:
//     With a single auth token - All requests will use the token blindly,
//         no verification of the token will be done and requests may fail.
//     With an email and password - Discord will sign in with the provided
//         credentials.
//     With an email, password and auth token - Discord will verify the auth
//         token, if it is invalid it will sign in with the provided
//         credentials. This is the Discord recommended way to sign in.
func New(args ...interface{}) (s *Session, err error) {

	// Create an empty Session interface.
	s = &Session{
		State:                  NewState(),
		StateEnabled:           true,
		Compress:               true,
		ShouldReconnectOnError: true,
	}

	// If no arguments are passed return the empty Session interface.
	if args == nil {
		return
	}

	// Variables used below when parsing func arguments
	var auth, pass string

	// Parse passed arguments
	for _, arg := range args {

		switch v := arg.(type) {

		case []string:
			if len(v) > 3 {
				err = fmt.Errorf("Too many string parameters provided.")
				return
			}

			// First string is either token or username
			if len(v) > 0 {
				auth = v[0]
			}

			// If second string exists, it must be a password.
			if len(v) > 1 {
				pass = v[1]
			}

			// If third string exists, it must be an auth token.
			if len(v) > 2 {
				s.Token = v[2]
			}

		case string:
			// First string must be either auth token or username.
			// Second string must be a password.
			// Only 2 input strings are supported.

			if auth == "" {
				auth = v
			} else if pass == "" {
				pass = v
			} else if s.Token == "" {
				s.Token = v
			} else {
				err = fmt.Errorf("Too many string parameters provided.")
				return
			}

			//		case Config:
			// TODO: Parse configuration struct

		default:
			err = fmt.Errorf("Unsupported parameter type provided.")
			return
		}
	}

	// If only one string was provided, assume it is an auth token.
	// Otherwise get auth token from Discord, if a token was specified
	// Discord will verify it for free, or log the user in if it is
	// invalid.
	if pass == "" {
		s.Token = auth
	} else {
		err = s.Login(auth, pass)
		if err != nil || s.Token == "" {
			err = fmt.Errorf("Unable to fetch discord authentication token. %v", err)
			return
		}
	}

	// The Session is now able to have RestAPI methods called on it.
	// It is recommended that you now call Open() so that events will trigger.

	return
}

// validateHandler takes an event handler func, and returns the type of event.
// eg.
//     Session.validateHandler(func (s *discordgo.Session, m *discordgo.MessageCreate))
//     will return the reflect.Type of *discordgo.MessageCreate
func (s *Session) validateHandler(handler interface{}) reflect.Type {

	handlerType := reflect.TypeOf(handler)

	if handlerType.NumIn() != 2 {
		panic("Unable to add event handler, handler must be of the type func(*discordgo.Session, *discordgo.EventType).")
	}

	if handlerType.In(0) != reflect.TypeOf(s) {
		panic("Unable to add event handler, first argument must be of type *discordgo.Session.")
	}

	eventType := handlerType.In(1)

	// Support handlers of type interface{}, this is a special handler, which is triggered on every event.
	if eventType.Kind() == reflect.Interface {
		eventType = nil
	}

	return eventType
}

// AddHandler allows you to add an event handler that will be fired anytime
// the Discord WSAPI event that matches the interface fires.
// eventToInterface in events.go has a list of all the Discord WSAPI events
// and their respective interface.
// eg:
//     Session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) {
//     })
//
// or:
//     Session.AddHandler(func(s *discordgo.Session, m *discordgo.PresenceUpdate) {
//     })
// The return value of this method is a function, that when called will remove the
// event handler.
func (s *Session) AddHandler(handler interface{}) func() {

	s.initialize()

	eventType := s.validateHandler(handler)

	s.handlersMu.Lock()
	defer s.handlersMu.Unlock()

	h := reflect.ValueOf(handler)

	s.handlers[eventType] = append(s.handlers[eventType], h)

	// This must be done as we need a consistent reference to the
	// reflected value, otherwise a RemoveHandler method would have
	// been nice.
	return func() {
		s.handlersMu.Lock()
		defer s.handlersMu.Unlock()

		handlers := s.handlers[eventType]
		for i, v := range handlers {
			if h == v {
				s.handlers[eventType] = append(handlers[:i], handlers[i+1:]...)
				return
			}
		}
	}
}

// handle calls any handlers that match the event type and any handlers of
// interface{}.
func (s *Session) handle(event interface{}) {

	s.handlersMu.RLock()
	defer s.handlersMu.RUnlock()

	if s.handlers == nil {
		return
	}

	handlerParameters := []reflect.Value{reflect.ValueOf(s), reflect.ValueOf(event)}

	if handlers, ok := s.handlers[nil]; ok {
		for _, handler := range handlers {
			go handler.Call(handlerParameters)
		}
	}

	if handlers, ok := s.handlers[reflect.TypeOf(event)]; ok {
		for _, handler := range handlers {
			go handler.Call(handlerParameters)
		}
	}
}

// initialize adds all internal handlers and state tracking handlers.
func (s *Session) initialize() {

	s.log(LogInformational, "called")

	s.handlersMu.Lock()
	if s.handlers != nil {
		s.handlersMu.Unlock()
		return
	}

	s.handlers = map[interface{}][]reflect.Value{}
	s.handlersMu.Unlock()

	s.AddHandler(s.onReady)
	s.AddHandler(s.onResumed)
	s.AddHandler(s.onVoiceServerUpdate)
	s.AddHandler(s.onVoiceStateUpdate)
	s.AddHandler(s.State.onInterface)
}

// onReady handles the ready event.
func (s *Session) onReady(se *Session, r *Ready) {

	// Store the SessionID within the Session struct.
	s.sessionID = r.SessionID

	// Start the heartbeat to keep the connection alive.
	go s.heartbeat(s.wsConn, s.listening, r.HeartbeatInterval)
}

// onResumed handles the resumed event.
func (s *Session) onResumed(se *Session, r *Resumed) {

	// Start the heartbeat to keep the connection alive.
	go s.heartbeat(s.wsConn, s.listening, r.HeartbeatInterval)
}
<|endoftext|>


<commit_before>package git

/*
#include <git2.h>

extern int _go_git_odb_foreach(git_odb *db, void *payload);
extern void _go_git_odb_backend_free(git_odb_backend *backend);
*/
import "C"
import (
	"reflect"
	"runtime"
	"unsafe"
)

type Odb struct {
	ptr *C.git_odb
}

type OdbBackend struct {
	ptr *C.git_odb_backend
}

func NewOdb() (odb *Odb, err error) {
	odb = new(Odb)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_new(&odb.ptr)
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(odb, (*Odb).Free)
	return odb, nil
}

func NewOdbBackendFromC(ptr *C.git_odb_backend) (backend *OdbBackend) {
	backend = &OdbBackend{ptr}
	return backend
}

func (v *Odb) AddBackend(backend *OdbBackend, priority int) (err error) {

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_add_backend(v.ptr, backend.ptr, C.int(priority))
	if ret < 0 {
		backend.Free()
		return MakeGitError(ret)
	}
	return nil
}

func (v *Odb) ReadHeader(oid *Oid) (uint64, ObjectType, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	var sz C.size_t
	var cotype C.git_otype

	ret := C.git_odb_read_header(&sz, &cotype, v.ptr, oid.toC())
	if ret < 0 {
		return 0, C.GIT_OBJ_BAD, MakeGitError(ret)
	}

	return uint64(sz), ObjectType(cotype), nil
}

func (v *Odb) Exists(oid *Oid) bool {
	ret := C.git_odb_exists(v.ptr, oid.toC())
	return ret != 0
}

func (v *Odb) Write(data []byte, otype ObjectType) (oid *Oid, err error) {
	oid = new(Oid)
	var cptr unsafe.Pointer
	if len(data) > 0 {
		cptr = unsafe.Pointer(&data[0])
	}

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_write(oid.toC(), v.ptr, cptr, C.size_t(len(data)), C.git_otype(otype))

	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	return oid, nil
}

func (v *Odb) Read(oid *Oid) (obj *OdbObject, err error) {
	obj = new(OdbObject)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_read(&obj.ptr, v.ptr, oid.toC())
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(obj, (*OdbObject).Free)
	return obj, nil
}

type OdbForEachCallback func(id *Oid) error

type foreachData struct {
	callback OdbForEachCallback
	err      error
}

//export odbForEachCb
func odbForEachCb(id *C.git_oid, handle unsafe.Pointer) int {
	data, ok := pointerHandles.Get(handle).(*foreachData)

	if !ok {
		panic("could not retrieve handle")
	}

	err := data.callback(newOidFromC(id))
	if err != nil {
		data.err = err
		return C.GIT_EUSER
	}

	return 0
}

func (v *Odb) ForEach(callback OdbForEachCallback) error {
	data := foreachData{
		callback: callback,
		err:      nil,
	}

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	handle := pointerHandles.Track(&data)
	defer pointerHandles.Untrack(handle)

	ret := C._go_git_odb_foreach(v.ptr, handle)
	if ret == C.GIT_EUSER {
		return data.err
	} else if ret < 0 {
		return MakeGitError(ret)
	}

	return nil
}

// Hash determines the object-ID (sha1) of a data buffer.
func (v *Odb) Hash(data []byte, otype ObjectType) (oid *Oid, err error) {
	oid = new(Oid)
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	ptr := unsafe.Pointer(header.Data)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_hash(oid.toC(), ptr, C.size_t(header.Len), C.git_otype(otype))
	if ret < 0 {
		return nil, MakeGitError(ret)
	}
	return oid, nil
}

// NewReadStream opens a read stream from the ODB. Reading from it will give you the
// contents of the object.
func (v *Odb) NewReadStream(id *Oid) (*OdbReadStream, error) {
	stream := new(OdbReadStream)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_open_rstream(&stream.ptr, v.ptr, id.toC())
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(stream, (*OdbReadStream).Free)
	return stream, nil
}

// NewWriteStream opens a write stream to the ODB, which allows you to
// create a new object in the database. The size and type must be
// known in advance
func (v *Odb) NewWriteStream(size int64, otype ObjectType) (*OdbWriteStream, error) {
	stream := new(OdbWriteStream)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_open_wstream(&stream.ptr, v.ptr, C.git_off_t(size), C.git_otype(otype))
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(stream, (*OdbWriteStream).Free)
	return stream, nil
}

func (v *OdbBackend) Free() {
	C._go_git_odb_backend_free(v.ptr)
}

type OdbObject struct {
	ptr *C.git_odb_object
}

func (v *OdbObject) Free() {
	runtime.SetFinalizer(v, nil)
	C.git_odb_object_free(v.ptr)
}

func (object *OdbObject) Id() (oid *Oid) {
	return newOidFromC(C.git_odb_object_id(object.ptr))
}

func (object *OdbObject) Len() (len uint64) {
	return uint64(C.git_odb_object_size(object.ptr))
}

func (object *OdbObject) Data() (data []byte) {
	var c_blob unsafe.Pointer = C.git_odb_object_data(object.ptr)
	var blob []byte

	len := int(C.git_odb_object_size(object.ptr))

	sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&blob)))
	sliceHeader.Cap = len
	sliceHeader.Len = len
	sliceHeader.Data = uintptr(c_blob)

	return blob
}

type OdbReadStream struct {
	ptr *C.git_odb_stream
}

// Read reads from the stream
func (stream *OdbReadStream) Read(data []byte) (int, error) {
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	ptr := (*C.char)(unsafe.Pointer(header.Data))
	size := C.size_t(header.Cap)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_stream_read(stream.ptr, ptr, size)
	if ret < 0 {
		return 0, MakeGitError(ret)
	}

	header.Len = int(ret)

	return len(data), nil
}

// Close is a dummy function in order to implement the Closer and
// ReadCloser interfaces
func (stream *OdbReadStream) Close() error {
	return nil
}

func (stream *OdbReadStream) Free() {
	runtime.SetFinalizer(stream, nil)
	C.git_odb_stream_free(stream.ptr)
}

type OdbWriteStream struct {
	ptr *C.git_odb_stream
	Id  Oid
}

// Write writes to the stream
func (stream *OdbWriteStream) Write(data []byte) (int, error) {
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	ptr := (*C.char)(unsafe.Pointer(header.Data))
	size := C.size_t(header.Len)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_stream_write(stream.ptr, ptr, size)
	if ret < 0 {
		return 0, MakeGitError(ret)
	}

	return len(data), nil
}

// Close signals that all the data has been written and stores the
// resulting object id in the stream's Id field.
func (stream *OdbWriteStream) Close() error {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_stream_finalize_write(stream.Id.toC(), stream.ptr)
	if ret < 0 {
		return MakeGitError(ret)
	}

	return nil
}

func (stream *OdbWriteStream) Free() {
	runtime.SetFinalizer(stream, nil)
	C.git_odb_stream_free(stream.ptr)
}
<commit_msg>odb: Expose git_odb_object_type() as OdbObject.Type()<commit_after>package git

/*
#include <git2.h>

extern int _go_git_odb_foreach(git_odb *db, void *payload);
extern void _go_git_odb_backend_free(git_odb_backend *backend);
*/
import "C"
import (
	"reflect"
	"runtime"
	"unsafe"
)

type Odb struct {
	ptr *C.git_odb
}

type OdbBackend struct {
	ptr *C.git_odb_backend
}

func NewOdb() (odb *Odb, err error) {
	odb = new(Odb)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_new(&odb.ptr)
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(odb, (*Odb).Free)
	return odb, nil
}

func NewOdbBackendFromC(ptr *C.git_odb_backend) (backend *OdbBackend) {
	backend = &OdbBackend{ptr}
	return backend
}

func (v *Odb) AddBackend(backend *OdbBackend, priority int) (err error) {

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_add_backend(v.ptr, backend.ptr, C.int(priority))
	if ret < 0 {
		backend.Free()
		return MakeGitError(ret)
	}
	return nil
}

func (v *Odb) ReadHeader(oid *Oid) (uint64, ObjectType, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	var sz C.size_t
	var cotype C.git_otype

	ret := C.git_odb_read_header(&sz, &cotype, v.ptr, oid.toC())
	if ret < 0 {
		return 0, C.GIT_OBJ_BAD, MakeGitError(ret)
	}

	return uint64(sz), ObjectType(cotype), nil
}

func (v *Odb) Exists(oid *Oid) bool {
	ret := C.git_odb_exists(v.ptr, oid.toC())
	return ret != 0
}

func (v *Odb) Write(data []byte, otype ObjectType) (oid *Oid, err error) {
	oid = new(Oid)
	var cptr unsafe.Pointer
	if len(data) > 0 {
		cptr = unsafe.Pointer(&data[0])
	}

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_write(oid.toC(), v.ptr, cptr, C.size_t(len(data)), C.git_otype(otype))

	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	return oid, nil
}

func (v *Odb) Read(oid *Oid) (obj *OdbObject, err error) {
	obj = new(OdbObject)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_read(&obj.ptr, v.ptr, oid.toC())
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(obj, (*OdbObject).Free)
	return obj, nil
}

type OdbForEachCallback func(id *Oid) error

type foreachData struct {
	callback OdbForEachCallback
	err      error
}

//export odbForEachCb
func odbForEachCb(id *C.git_oid, handle unsafe.Pointer) int {
	data, ok := pointerHandles.Get(handle).(*foreachData)

	if !ok {
		panic("could not retrieve handle")
	}

	err := data.callback(newOidFromC(id))
	if err != nil {
		data.err = err
		return C.GIT_EUSER
	}

	return 0
}

func (v *Odb) ForEach(callback OdbForEachCallback) error {
	data := foreachData{
		callback: callback,
		err:      nil,
	}

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	handle := pointerHandles.Track(&data)
	defer pointerHandles.Untrack(handle)

	ret := C._go_git_odb_foreach(v.ptr, handle)
	if ret == C.GIT_EUSER {
		return data.err
	} else if ret < 0 {
		return MakeGitError(ret)
	}

	return nil
}

// Hash determines the object-ID (sha1) of a data buffer.
func (v *Odb) Hash(data []byte, otype ObjectType) (oid *Oid, err error) {
	oid = new(Oid)
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	ptr := unsafe.Pointer(header.Data)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_hash(oid.toC(), ptr, C.size_t(header.Len), C.git_otype(otype))
	if ret < 0 {
		return nil, MakeGitError(ret)
	}
	return oid, nil
}

// NewReadStream opens a read stream from the ODB. Reading from it will give you the
// contents of the object.
func (v *Odb) NewReadStream(id *Oid) (*OdbReadStream, error) {
	stream := new(OdbReadStream)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_open_rstream(&stream.ptr, v.ptr, id.toC())
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(stream, (*OdbReadStream).Free)
	return stream, nil
}

// NewWriteStream opens a write stream to the ODB, which allows you to
// create a new object in the database. The size and type must be
// known in advance
func (v *Odb) NewWriteStream(size int64, otype ObjectType) (*OdbWriteStream, error) {
	stream := new(OdbWriteStream)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_open_wstream(&stream.ptr, v.ptr, C.git_off_t(size), C.git_otype(otype))
	if ret < 0 {
		return nil, MakeGitError(ret)
	}

	runtime.SetFinalizer(stream, (*OdbWriteStream).Free)
	return stream, nil
}

func (v *OdbBackend) Free() {
	C._go_git_odb_backend_free(v.ptr)
}

type OdbObject struct {
	ptr *C.git_odb_object
}

func (v *OdbObject) Free() {
	runtime.SetFinalizer(v, nil)
	C.git_odb_object_free(v.ptr)
}

func (object *OdbObject) Id() (oid *Oid) {
	return newOidFromC(C.git_odb_object_id(object.ptr))
}

func (object *OdbObject) Len() (len uint64) {
	return uint64(C.git_odb_object_size(object.ptr))
}

func (object *OdbObject) Type() ObjectType {
	return ObjectType(C.git_odb_object_type(object.ptr))
}

func (object *OdbObject) Data() (data []byte) {
	var c_blob unsafe.Pointer = C.git_odb_object_data(object.ptr)
	var blob []byte

	len := int(C.git_odb_object_size(object.ptr))

	sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&blob)))
	sliceHeader.Cap = len
	sliceHeader.Len = len
	sliceHeader.Data = uintptr(c_blob)

	return blob
}

type OdbReadStream struct {
	ptr *C.git_odb_stream
}

// Read reads from the stream
func (stream *OdbReadStream) Read(data []byte) (int, error) {
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	ptr := (*C.char)(unsafe.Pointer(header.Data))
	size := C.size_t(header.Cap)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_stream_read(stream.ptr, ptr, size)
	if ret < 0 {
		return 0, MakeGitError(ret)
	}

	header.Len = int(ret)

	return len(data), nil
}

// Close is a dummy function in order to implement the Closer and
// ReadCloser interfaces
func (stream *OdbReadStream) Close() error {
	return nil
}

func (stream *OdbReadStream) Free() {
	runtime.SetFinalizer(stream, nil)
	C.git_odb_stream_free(stream.ptr)
}

type OdbWriteStream struct {
	ptr *C.git_odb_stream
	Id  Oid
}

// Write writes to the stream
func (stream *OdbWriteStream) Write(data []byte) (int, error) {
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	ptr := (*C.char)(unsafe.Pointer(header.Data))
	size := C.size_t(header.Len)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_stream_write(stream.ptr, ptr, size)
	if ret < 0 {
		return 0, MakeGitError(ret)
	}

	return len(data), nil
}

// Close signals that all the data has been written and stores the
// resulting object id in the stream's Id field.
func (stream *OdbWriteStream) Close() error {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	ret := C.git_odb_stream_finalize_write(stream.Id.toC(), stream.ptr)
	if ret < 0 {
		return MakeGitError(ret)
	}

	return nil
}

func (stream *OdbWriteStream) Free() {
	runtime.SetFinalizer(stream, nil)
	C.git_odb_stream_free(stream.ptr)
}
<|endoftext|>


<commit_before>// delayd2 is an available setTimeout() service for scheduling message sends
package delayd2

import (
	"context"
	"log"
	"runtime"
	"sync"
	"time"

	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/pmylund/go-cache"
)

// QueueMessage is a message queued in the database.
type QueueMessage struct {
	QueueID   string
	WorkerID  string
	ReleaseAt time.Time
	RelayTo   string
	Payload   string
}

// WorkerConfig is configuration for Worker.
type WorkerConfig struct {
	ID string

	LeaveMessagesOrphanedAtShutdown bool

	NumConsumerFactor int
	NumRelayFactor    int
}

// Worker is the delayd2 worker.
type Worker struct {
	config *WorkerConfig

	driver   Driver
	consumer *Consumer
	relay    *Relay

	succededIDsCache *cache.Cache

	relayCh chan releaseJob

	// signaled from external
	shutdownCh chan struct{}

	// wait for all worker goroutine to be finished
	stoppped sync.WaitGroup
}

// NewWorker creates a new worker.
func NewWorker(c *WorkerConfig, driver Driver, consumer *Consumer, relay *Relay) *Worker {
	return &Worker{
		driver:   driver,
		consumer: consumer,
		relay:    relay,

		// Cache succededIDs to prepare for transient errors in the database (e.g. failover/network problems)
		succededIDsCache: cache.New(15*time.Minute, 1*time.Minute),

		config: c,

		relayCh:    make(chan releaseJob, runtime.NumCPU()*c.NumRelayFactor),
		shutdownCh: make(chan struct{}),
	}
}

func (w *Worker) runWorker(f func()) {
	w.stoppped.Add(1)
	go func() {
		f()
		w.stoppped.Done()
	}()
}

// Run starts the worker process.
func (w *Worker) Run() error {
	log.Print("worker: starting delayd2 worker process")

	log.Print("worker: registering delayd2 worker process")
	if err := w.driver.RegisterSession(); err != nil {
		return err
	}

	defer func() {
		log.Print("worker: resetting aborted messages remaining in active queue")
		if _, err := w.driver.ResetActive(); err != nil {
			log.Print("worker: unable to reset messages in active queue:", err)
		}

		if w.config.LeaveMessagesOrphanedAtShutdown {
			log.Print("worker: leaving messages as orphaned")

			if err := w.driver.MarkOrphaned(); err != nil {
				log.Print("worker: unable to mark messages as orphaned:", err)
			}
		}

		log.Print("worker: deregistering delayd2 worker process")
		if err := w.driver.DeregisterSession(); err != nil {
			log.Print("worker: unable to deregister session:", err)
		}
	}()

	baseCtx := context.Background()
	ctx, cancel := context.WithCancel(baseCtx)

	nCPU := runtime.NumCPU()
	for i := 0; i < nCPU*w.config.NumConsumerFactor; i++ {
		log.Print("worker: launching consumer process")
		w.runWorker(func() { w.consumeWorker(ctx) })

	}

	for i := 0; i < nCPU*w.config.NumRelayFactor; i++ {
		log.Print("worker: launching relay worker process")
		w.runWorker(func() { w.relayWorker(ctx) })
	}

	w.runWorker(func() { w.adoptOrphansWorker(ctx) })
	w.runWorker(func() { w.keepAliveWorker(ctx) })
	w.runWorker(func() { w.markActiveWorker(ctx) })
	w.runWorker(func() { w.releaseWorker(ctx) })

	log.Print("worker: starting delayd2 process")

	<-w.shutdownCh
	log.Print("worker: receiving shutdown signal. waiting for the process finished.")
	cancel()

	w.stoppped.Wait()
	log.Print("worker: the process finished.")

	return nil
}

func (w *Worker) Stop() {
	close(w.shutdownCh)
}

func (w *Worker) consume() (int64, error) {
	return w.consumer.ConsumeMessages()
}

func (w *Worker) keepAliveWorker(ctx context.Context) {
	for range time.Tick(1 * time.Second) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down keepalive worker")
			return
		default:
		}

		if err := w.driver.KeepAliveSession(); err != nil {
			log.Printf("worker: unable to keep alived: %s", err)
		}
	}
}

// consumeWorker consumes messages in SQS.
func (w *Worker) consumeWorker(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down consuming worker")
			return
		default:
		}

		begin := time.Now()
		n, err := w.consume()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to consume messages: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages consumed in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) adoptOrphansWorker(ctx context.Context) {
	for range time.Tick(10 * time.Second) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down adopting worker")
			return
		default:
		}

		begin := time.Now()
		n, err := w.driver.AdoptOrphans()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to adopt orphans: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d orphans adopted in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) markActive(begin time.Time) (int64, error) {
	return w.driver.MarkActive(begin)
}

// markActiveWorker marks coming messages in queue.
func (w *Worker) markActiveWorker(ctx context.Context) {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down marking worker")
			return
		default:
		}

		begin := time.Now().Truncate(time.Second)
		n, err := w.markActive(begin)
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to mark messages active: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages marked as active in %s", n, end.Sub(begin))
		}
		time.Sleep(1 * time.Second)
	}
}

// releaseWorker releases messages that is ready to fire.
func (w *Worker) releaseWorker(ctx context.Context) {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down releasing worker")
			return
		default:
		}

		if err := w.release(); err != nil {
			log.Printf("worker: unable to release messages: %s", err)
			continue
		}
	}
}

type releaseJob struct {
	*sync.WaitGroup

	RelayTo  string
	Messages []releaseMessage
}

type releaseMessage struct {
	QueueID string
	Payload string
}

// release releases active messages. It returns the number of released message n.
func (w *Worker) release() error {
	messages, err := w.driver.GetActiveMessages()

	if err != nil {
		return err
	}

	// we need QueueID to delete from queue
	batchMap := BuildBatchMap(messages)

	var wg sync.WaitGroup
	for relayTo, batchMsgs := range batchMap {
		for _, rms := range batchMsgs {
			wg.Add(1)
			w.relayCh <- releaseJob{
				WaitGroup: &wg,

				RelayTo:  relayTo,
				Messages: rms,
			}
		}
	}
	wg.Wait()
	return nil
}

func (w *Worker) relayWorker(ctx context.Context) {
	for rj := range w.relayCh {
		begin := time.Now()
		n := w.releaseBatch(rj)
		end := time.Now()

		if n > 0 {
			log.Printf("worker: %d messages relayed in %s", n, end.Sub(begin))
		}
		rj.Done()

		select {
		case <-ctx.Done():
			return
		default:
		}
	}
	log.Print("relayworker: closed")
}

func (w *Worker) releaseBatch(rj releaseJob) int64 {
	payloads := make([]string, 0, len(rj.Messages))
	for _, m := range rj.Messages {
		if _, found := w.succededIDsCache.Get(m.QueueID); found {
			// skip since we already relayed
			log.Printf("worker: seeing queue_id in the cache so continueing...")
			continue
		}
		payloads = append(payloads, m.Payload)
	}

	var n int64
	failedIndex := make(map[int]struct{})
	if len(payloads) > 0 {
		if err := w.relay.Relay(rj.RelayTo, payloads); err != nil {
			// only print log here in batch operation
			// TODO: a dead letter queue support
			berrs, batchOk := queue.IsBatchError(err)
			if !batchOk {
				log.Printf("worker: unable to send message due to non batch error. skipping: %s", err)
				return 0
			}

			for _, berr := range berrs {
				if berr.SenderFault {
					log.Printf("worker: unable to send message due to sender's fault: skipping: %s", berr.Message)
				}
				failedIndex[berr.Index] = struct{}{}
			}
		}
	}

	succededIDs := make([]string, 0, len(rj.Messages))
	for i, m := range rj.Messages {
		if _, failed := failedIndex[i]; failed {
			continue
		}
		succededIDs = append(succededIDs, m.QueueID)

		// remember succededIDs for a while to prevent us from relaying message in transient error
		w.succededIDsCache.Set(m.QueueID, struct{}{}, cache.DefaultExpiration)
	}

	if err := w.driver.RemoveMessages(succededIDs); err != nil {
		log.Printf("worker: unable to remove messages from queue: %s", err)
	} else {
		// reset succededIDs here since we succeeded to remove messages from the database
		// When we succeeded to remove messages, the messages will not appear again so it's safe.
		for _, id := range succededIDs {
			w.succededIDsCache.Delete(id)
		}
		n += int64(len(succededIDs))
	}

	return n
}

func BuildBatchMap(messages []*QueueMessage) map[string][][]releaseMessage {
	batchMap := make(map[string][][]releaseMessage)
	for _, m := range messages {
		rm := releaseMessage{
			QueueID: m.QueueID,
			Payload: m.Payload,
		}

		var i int
		i = len(batchMap[m.RelayTo])
		if i == 0 || len(batchMap[m.RelayTo][i-1]) >= 10 {
			batchMap[m.RelayTo] = append(batchMap[m.RelayTo], []releaseMessage{})
		}

		i = len(batchMap[m.RelayTo])

		batchMap[m.RelayTo][i-1] = append(batchMap[m.RelayTo][i-1], rm)
	}
	return batchMap
}
<commit_msg>worker: Close relayCh while shutting down<commit_after>// delayd2 is an available setTimeout() service for scheduling message sends
package delayd2

import (
	"context"
	"log"
	"runtime"
	"sync"
	"time"

	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/pmylund/go-cache"
)

// QueueMessage is a message queued in the database.
type QueueMessage struct {
	QueueID   string
	WorkerID  string
	ReleaseAt time.Time
	RelayTo   string
	Payload   string
}

// WorkerConfig is configuration for Worker.
type WorkerConfig struct {
	ID string

	LeaveMessagesOrphanedAtShutdown bool

	NumConsumerFactor int
	NumRelayFactor    int
}

// Worker is the delayd2 worker.
type Worker struct {
	config *WorkerConfig

	driver   Driver
	consumer *Consumer
	relay    *Relay

	succededIDsCache *cache.Cache

	relayCh chan releaseJob

	// signaled from external
	shutdownCh chan struct{}

	// wait for all worker goroutine to be finished
	stoppped sync.WaitGroup
}

// NewWorker creates a new worker.
func NewWorker(c *WorkerConfig, driver Driver, consumer *Consumer, relay *Relay) *Worker {
	return &Worker{
		driver:   driver,
		consumer: consumer,
		relay:    relay,

		// Cache succededIDs to prepare for transient errors in the database (e.g. failover/network problems)
		succededIDsCache: cache.New(15*time.Minute, 1*time.Minute),

		config: c,

		relayCh:    make(chan releaseJob, runtime.NumCPU()*c.NumRelayFactor),
		shutdownCh: make(chan struct{}),
	}
}

func (w *Worker) runWorker(f func()) {
	w.stoppped.Add(1)
	go func() {
		f()
		w.stoppped.Done()
	}()
}

// Run starts the worker process.
func (w *Worker) Run() error {
	log.Print("worker: starting delayd2 worker process")

	log.Print("worker: registering delayd2 worker process")
	if err := w.driver.RegisterSession(); err != nil {
		return err
	}

	defer func() {
		log.Print("worker: resetting aborted messages remaining in active queue")
		if _, err := w.driver.ResetActive(); err != nil {
			log.Print("worker: unable to reset messages in active queue:", err)
		}

		if w.config.LeaveMessagesOrphanedAtShutdown {
			log.Print("worker: leaving messages as orphaned")

			if err := w.driver.MarkOrphaned(); err != nil {
				log.Print("worker: unable to mark messages as orphaned:", err)
			}
		}

		log.Print("worker: deregistering delayd2 worker process")
		if err := w.driver.DeregisterSession(); err != nil {
			log.Print("worker: unable to deregister session:", err)
		}
	}()

	baseCtx := context.Background()
	ctx, cancel := context.WithCancel(baseCtx)

	nCPU := runtime.NumCPU()
	for i := 0; i < nCPU*w.config.NumConsumerFactor; i++ {
		log.Print("worker: launching consumer process")
		w.runWorker(func() { w.consumeWorker(ctx) })

	}

	for i := 0; i < nCPU*w.config.NumRelayFactor; i++ {
		log.Print("worker: launching relay worker process")
		w.runWorker(func() { w.relayWorker(ctx) })
	}

	w.runWorker(func() { w.adoptOrphansWorker(ctx) })
	w.runWorker(func() { w.keepAliveWorker(ctx) })
	w.runWorker(func() { w.markActiveWorker(ctx) })
	w.runWorker(func() { w.releaseWorker(ctx) })

	log.Print("worker: started delayd2 process")

	<-w.shutdownCh
	log.Print("worker: receiving shutdown signal. waiting for the process finished.")
	cancel()
	close(w.relayCh)

	w.stoppped.Wait()
	log.Print("worker: the process finished.")

	return nil
}

func (w *Worker) Stop() {
	close(w.shutdownCh)
}

func (w *Worker) consume() (int64, error) {
	return w.consumer.ConsumeMessages()
}

func (w *Worker) keepAliveWorker(ctx context.Context) {
	for range time.Tick(1 * time.Second) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down keepalive worker")
			return
		default:
		}

		if err := w.driver.KeepAliveSession(); err != nil {
			log.Printf("worker: unable to keep alived: %s", err)
		}
	}
}

// consumeWorker consumes messages in SQS.
func (w *Worker) consumeWorker(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down consuming worker")
			return
		default:
		}

		begin := time.Now()
		n, err := w.consume()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to consume messages: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages consumed in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) adoptOrphansWorker(ctx context.Context) {
	for range time.Tick(10 * time.Second) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down adopting worker")
			return
		default:
		}

		begin := time.Now()
		n, err := w.driver.AdoptOrphans()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to adopt orphans: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d orphans adopted in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) markActive(begin time.Time) (int64, error) {
	return w.driver.MarkActive(begin)
}

// markActiveWorker marks coming messages in queue.
func (w *Worker) markActiveWorker(ctx context.Context) {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down marking worker")
			return
		default:
		}

		begin := time.Now().Truncate(time.Second)
		n, err := w.markActive(begin)
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to mark messages active: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages marked as active in %s", n, end.Sub(begin))
		}
		time.Sleep(1 * time.Second)
	}
}

// releaseWorker releases messages that is ready to fire.
func (w *Worker) releaseWorker(ctx context.Context) {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-ctx.Done():
			log.Print("worker: shutting down releasing worker")
			return
		default:
		}

		if err := w.release(); err != nil {
			log.Printf("worker: unable to release messages: %s", err)
			continue
		}
	}
}

type releaseJob struct {
	*sync.WaitGroup

	RelayTo  string
	Messages []releaseMessage
}

type releaseMessage struct {
	QueueID string
	Payload string
}

// release releases active messages. It returns the number of released message n.
func (w *Worker) release() error {
	messages, err := w.driver.GetActiveMessages()

	if err != nil {
		return err
	}

	// we need QueueID to delete from queue
	batchMap := BuildBatchMap(messages)

	var wg sync.WaitGroup
	for relayTo, batchMsgs := range batchMap {
		for _, rms := range batchMsgs {
			wg.Add(1)
			w.relayCh <- releaseJob{
				WaitGroup: &wg,

				RelayTo:  relayTo,
				Messages: rms,
			}
		}
	}
	wg.Wait()
	return nil
}

func (w *Worker) relayWorker(ctx context.Context) {
	for rj := range w.relayCh {
		begin := time.Now()
		n := w.releaseBatch(rj)
		end := time.Now()

		if n > 0 {
			log.Printf("worker: %d messages relayed in %s", n, end.Sub(begin))
		}
		rj.Done()

		select {
		case <-ctx.Done():
			log.Print("worker: shutting relay worker")
			return
		default:
		}
	}
}

func (w *Worker) releaseBatch(rj releaseJob) int64 {
	payloads := make([]string, 0, len(rj.Messages))
	for _, m := range rj.Messages {
		if _, found := w.succededIDsCache.Get(m.QueueID); found {
			// skip since we already relayed
			log.Printf("worker: seeing queue_id in the cache so continueing...")
			continue
		}
		payloads = append(payloads, m.Payload)
	}

	var n int64
	failedIndex := make(map[int]struct{})
	if len(payloads) > 0 {
		if err := w.relay.Relay(rj.RelayTo, payloads); err != nil {
			// only print log here in batch operation
			// TODO: a dead letter queue support
			berrs, batchOk := queue.IsBatchError(err)
			if !batchOk {
				log.Printf("worker: unable to send message due to non batch error. skipping: %s", err)
				return 0
			}

			for _, berr := range berrs {
				if berr.SenderFault {
					log.Printf("worker: unable to send message due to sender's fault: skipping: %s", berr.Message)
				}
				failedIndex[berr.Index] = struct{}{}
			}
		}
	}

	succededIDs := make([]string, 0, len(rj.Messages))
	for i, m := range rj.Messages {
		if _, failed := failedIndex[i]; failed {
			continue
		}
		succededIDs = append(succededIDs, m.QueueID)

		// remember succededIDs for a while to prevent us from relaying message in transient error
		w.succededIDsCache.Set(m.QueueID, struct{}{}, cache.DefaultExpiration)
	}

	if err := w.driver.RemoveMessages(succededIDs); err != nil {
		log.Printf("worker: unable to remove messages from queue: %s", err)
	} else {
		// reset succededIDs here since we succeeded to remove messages from the database
		// When we succeeded to remove messages, the messages will not appear again so it's safe.
		for _, id := range succededIDs {
			w.succededIDsCache.Delete(id)
		}
		n += int64(len(succededIDs))
	}

	return n
}

func BuildBatchMap(messages []*QueueMessage) map[string][][]releaseMessage {
	batchMap := make(map[string][][]releaseMessage)
	for _, m := range messages {
		rm := releaseMessage{
			QueueID: m.QueueID,
			Payload: m.Payload,
		}

		var i int
		i = len(batchMap[m.RelayTo])
		if i == 0 || len(batchMap[m.RelayTo][i-1]) >= 10 {
			batchMap[m.RelayTo] = append(batchMap[m.RelayTo], []releaseMessage{})
		}

		i = len(batchMap[m.RelayTo])

		batchMap[m.RelayTo][i-1] = append(batchMap[m.RelayTo][i-1], rm)
	}
	return batchMap
}
<|endoftext|>