<commit_before>package blobstore_test

import (
	"reflect"
	"sort"
	"testing"

	"github.com/nyaxt/otaru/blobstore"
	"github.com/nyaxt/otaru/flags"
	tu "github.com/nyaxt/otaru/testutils"
)

func TestCachedBlobStore(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	if err := tu.WriteVersionedBlob(backendbs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}
	// assert cache not yet filled
	if err := tu.AssertBlobVersion(cachebs, "backendonly", 0); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersionRA(bs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}

	// assert cache fill
	if err := tu.AssertBlobVersion(cachebs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := tu.WriteVersionedBlobRA(bs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}

	if err := tu.AssertBlobVersionRA(bs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(cachebs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
}

func TestCachedBlobStore_Invalidate(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	if err := tu.WriteVersionedBlob(cachebs, "backendnewer", 2); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlob(backendbs, "backendnewer", 3); err != nil {
		t.Errorf("%v", err)
		return
	}

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}
	if err := tu.AssertBlobVersionRA(bs, "backendnewer", 3); err != nil {
		t.Errorf("%v", err)
		return
	}

	// assert cache fill
	if err := tu.AssertBlobVersion(cachebs, "backendnewer", 3); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := tu.WriteVersionedBlobRA(bs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := tu.AssertBlobVersionRA(bs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.AssertBlobVersion(cachebs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}
}

func TestCachedBlobStore_NewEntry(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	if err := tu.WriteVersionedBlobRA(bs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersionRA(bs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.AssertBlobVersion(cachebs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
}

func TestCachedBlobStore_AutoExpandLen(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	bh, err := bs.Open("hoge", flags.O_RDWRCREATE)
	if err != nil {
		t.Errorf("Failed to open blobhandle")
	}
	defer bh.Close()

	if size := bh.Size(); size != 0 {
		t.Errorf("New bh size non-zero: %d", size)
	}

	if err := bh.PWrite(0, []byte("Hello")); err != nil {
		t.Errorf("PWrite failed: %v", err)
	}

	if size := bh.Size(); size != 5 {
		t.Errorf("bh size not auto expanded! size  %d", size)
	}
}

func TestCachedBlobStore_ListBlobs(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	if err := tu.WriteVersionedBlob(backendbs, "backendonly", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlob(cachebs, "cacheonly", 2); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "synced", 3); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "unsynced", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	bpaths, err := bs.ListBlobs()
	if err != nil {
		t.Errorf("ListBlobs failed: %v", err)
		return
	}
	sort.Strings(bpaths)
	if !reflect.DeepEqual([]string{"backendonly", "synced", "unsynced"}, bpaths) {
		t.Errorf("ListBlobs returned unexpected result: %v", bpaths)
	}
}

func TestCachedBlobStore_RemoveBlob(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	if err := tu.WriteVersionedBlob(backendbs, "backendonly", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlob(cachebs, "cacheonly", 2); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "synced", 3); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "unsynced", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := bs.RemoveBlob("backendonly"); err != nil {
		t.Errorf("RemoveBlob failed: %v", err)
		return
	}
	if err := bs.RemoveBlob("synced"); err != nil {
		t.Errorf("RemoveBlob failed: %v", err)
		return
	}
	if err := bs.RemoveBlob("unsynced"); err != nil {
		t.Errorf("RemoveBlob failed: %v", err)
		return
	}

	bpaths, err := bs.ListBlobs()
	if err != nil {
		t.Errorf("ListBlobs failed: %v", err)
		return
	}
	if len(bpaths) > 0 {
		t.Errorf("Left over blobs: %v", bpaths)
	}
}
<commit_msg>Add test exposing RemoveBlob bug<commit_after>package blobstore_test

import (
	"reflect"
	"sort"
	"testing"

	"github.com/nyaxt/otaru/blobstore"
	"github.com/nyaxt/otaru/flags"
	tu "github.com/nyaxt/otaru/testutils"
)

func TestCachedBlobStore(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	if err := tu.WriteVersionedBlob(backendbs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}
	// assert cache not yet filled
	if err := tu.AssertBlobVersion(cachebs, "backendonly", 0); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersionRA(bs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}

	// assert cache fill
	if err := tu.AssertBlobVersion(cachebs, "backendonly", 5); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := tu.WriteVersionedBlobRA(bs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}

	if err := tu.AssertBlobVersionRA(bs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(cachebs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "backendonly", 10); err != nil {
		t.Errorf("%v", err)
		return
	}
}

func TestCachedBlobStore_Invalidate(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	if err := tu.WriteVersionedBlob(cachebs, "backendnewer", 2); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlob(backendbs, "backendnewer", 3); err != nil {
		t.Errorf("%v", err)
		return
	}

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}
	if err := tu.AssertBlobVersionRA(bs, "backendnewer", 3); err != nil {
		t.Errorf("%v", err)
		return
	}

	// assert cache fill
	if err := tu.AssertBlobVersion(cachebs, "backendnewer", 3); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := tu.WriteVersionedBlobRA(bs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := tu.AssertBlobVersionRA(bs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.AssertBlobVersion(cachebs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "backendnewer", 4); err != nil {
		t.Errorf("%v", err)
		return
	}
}

func TestCachedBlobStore_NewEntry(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	if err := tu.WriteVersionedBlobRA(bs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersionRA(bs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.AssertBlobVersion(cachebs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.AssertBlobVersion(backendbs, "newentry", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
}

func TestCachedBlobStore_AutoExpandLen(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	bh, err := bs.Open("hoge", flags.O_RDWRCREATE)
	if err != nil {
		t.Errorf("Failed to open blobhandle")
	}
	defer bh.Close()

	if size := bh.Size(); size != 0 {
		t.Errorf("New bh size non-zero: %d", size)
	}

	if err := bh.PWrite(0, []byte("Hello")); err != nil {
		t.Errorf("PWrite failed: %v", err)
	}

	if size := bh.Size(); size != 5 {
		t.Errorf("bh size not auto expanded! size  %d", size)
	}
}

func TestCachedBlobStore_ListBlobs(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	if err := tu.WriteVersionedBlob(backendbs, "backendonly", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlob(cachebs, "cacheonly", 2); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "synced", 3); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "unsynced", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	bpaths, err := bs.ListBlobs()
	if err != nil {
		t.Errorf("ListBlobs failed: %v", err)
		return
	}
	sort.Strings(bpaths)
	if !reflect.DeepEqual([]string{"backendonly", "synced", "unsynced"}, bpaths) {
		t.Errorf("ListBlobs returned unexpected result: %v", bpaths)
	}
}

func TestCachedBlobStore_RemoveBlob(t *testing.T) {
	backendbs := tu.TestFileBlobStoreOfName("backend")
	cachebs := tu.TestFileBlobStoreOfName("cache")

	bs, err := blobstore.NewCachedBlobStore(backendbs, cachebs, flags.O_RDWRCREATE, tu.TestQueryVersion)
	if err != nil {
		t.Errorf("Failed to create CachedBlobStore: %v", err)
		return
	}

	if err := tu.WriteVersionedBlob(backendbs, "backendonly", 1); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlob(cachebs, "cacheonly", 2); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "synced", 3); err != nil {
		t.Errorf("%v", err)
		return
	}
	if err := bs.Sync(); err != nil {
		t.Errorf("Sync failed: %v", err)
		return
	}
	if err := tu.WriteVersionedBlobRA(bs, "unsynced", 4); err != nil {
		t.Errorf("%v", err)
		return
	}

	if err := bs.RemoveBlob("backendonly"); err != nil {
		t.Errorf("RemoveBlob failed: %v", err)
		return
	}
	if err := bs.RemoveBlob("synced"); err != nil {
		t.Errorf("RemoveBlob failed: %v", err)
		return
	}
	if err := bs.RemoveBlob("unsynced"); err != nil {
		t.Errorf("RemoveBlob failed: %v", err)
		return
	}

	bpaths, err := bs.ListBlobs()
	if err != nil {
		t.Errorf("ListBlobs failed: %v", err)
		return
	}
	if len(bpaths) > 0 {
		t.Errorf("Left over blobs: %v", bpaths)
	}

	for _, bp := range []string{"backendonly", "synced", "unsynced"} {
		if err := tu.AssertBlobVersionRA(bs, bp, 0); err != nil {
			t.Errorf("left over blob in bs: %s", bp)
		}
		if err := tu.AssertBlobVersion(cachebs, bp, 0); err != nil {
			t.Errorf("left over blob in cachebs: %s", bp)
		}
		if err := tu.AssertBlobVersion(backendbs, bp, 0); err != nil {
			t.Errorf("left over blob in backendbs: %s", bp)
		}
	}
}
<|endoftext|>


<commit_before>package mixpanel

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

var (
	ErrTrackFailed = errors.New("Mixpanel did not return 1 when tracking")

	IgnoreTime *time.Time = &time.Time{}
)

// The Mixapanel struct store the mixpanel endpoint and the project token
type Mixpanel interface {
	// Create a mixpanel event
	Track(distinctId, eventName string, e *Event) error

	// Set properties for a mixpanel user.
	Update(distinctId string, u *Update) error

	Alias(distinctId, newId string) error
}

// The Mixapanel struct store the mixpanel endpoint and the project token
type mixpanel struct {
	Client *http.Client
	Token  string
	ApiURL string
}

// A mixpanel event
type Event struct {
	// IP-address of the user. Leave empty to use autodetect, or set to "0" to
	// not specify an ip-address.
	IP string

	// Timestamp. Set to nil to use the current time.
	Timestamp *time.Time

	// Custom properties. At least one must be specified.
	Properties map[string]interface{}
}

// An update of a user in mixpanel
type Update struct {
	// IP-address of the user. Leave empty to use autodetect, or set to "0" to
	// not specify an ip-address at all.
	IP string

	// Timestamp. Set to nil to use the current time, or IgnoreTime to not use a
	// timestamp.
	Timestamp *time.Time

	// Update operation such as "$set", "$update" etc.
	Operation string

	// Custom properties. At least one must be specified.
	Properties map[string]interface{}
}

// Track create a events to current distinct id
func (m *mixpanel) Alias(distinctId, newId string) error {
	props := map[string]interface{}{
		"token":       m.Token,
		"distinct_id": distinctId,
		"alias":       newId,
	}

	params := map[string]interface{}{
		"event":      "$create_alias",
		"properties": props,
	}

	return m.send("track", params, false)
}

// Track create a events to current distinct id
func (m *mixpanel) Track(distinctId, eventName string, e *Event) error {
	props := map[string]interface{}{
		"token":       m.Token,
		"distinct_id": distinctId,
	}
	if e.IP != "" {
		props["ip"] = e.IP
	}
	if e.Timestamp != nil {
		props["time"] = e.Timestamp.Unix()
	}

	for key, value := range e.Properties {
		props[key] = value
	}

	params := map[string]interface{}{
		"event":      eventName,
		"properties": props,
	}

	autoGeolocate := e.IP == ""

	return m.send("track", params, autoGeolocate)
}

// Updates a user in mixpanel. See
// https://mixpanel.com/help/reference/http#people-analytics-updates
func (m *mixpanel) Update(distinctId string, u *Update) error {
	params := map[string]interface{}{
		"$token":       m.Token,
		"$distinct_id": distinctId,
	}

	if u.IP != "" {
		params["$ip"] = u.IP
	}
	if u.Timestamp == IgnoreTime {
		params["$ignore_time"] = true
	} else if u.Timestamp != nil {
		params["$time"] = u.Timestamp.Unix()
	}

	params[u.Operation] = u.Properties

	autoGeolocate := u.IP == ""

	return m.send("engage", params, autoGeolocate)
}

func (m *mixpanel) to64(data string) string {
	bytes := []byte(data)
	return base64.StdEncoding.EncodeToString(bytes)
}

func (m *mixpanel) send(eventType string, params interface{}, autoGeolocate bool) error {
	dataJSON, _ := json.Marshal(params)
	data := string(dataJSON)

	url := m.ApiURL + "/" + eventType + "?data=" + m.to64(data)

	if autoGeolocate {
		url += "&ip=1"
	}

	if resp, err := m.Client.Get(url); err != nil {
		return fmt.Errorf("mixpanel: %s", err.Error())
	} else {
		defer resp.Body.Close()
		body, bodyErr := ioutil.ReadAll(resp.Body)
		if bodyErr != nil {
			return fmt.Errorf("mixpanel: %s", bodyErr.Error())
		}
		if string(body) != "1" && string(body) != "1\n" {
			return ErrTrackFailed
		}
	}

	return nil
}

// New returns the client instance. If apiURL is blank, the default will be used
// ("https://api.mixpanel.com").
func New(token, apiURL string) Mixpanel {
	return NewFromClient(http.DefaultClient, token, apiURL)
}

// Creates a client instance using the specified client instance. This is useful
// when using a proxy.
func NewFromClient(c *http.Client, token, apiURL string) Mixpanel {
	if apiURL == "" {
		apiURL = "https://api.mixpanel.com"
	}

	return &mixpanel{
		Client: c,
		Token:  token,
		ApiURL: apiURL,
	}
}
<commit_msg>Make ErrTrackFailed a type to carry usefull information<commit_after>package mixpanel

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

var IgnoreTime *time.Time = &time.Time{}

type ErrTrackFailed struct {
	Body string
	Resp *http.Response
}

func (err *ErrTrackFailed) Error() string {
	return fmt.Sprintf("Mixpanel did not return 1 when tracking: %s", err.Body)
}

// The Mixapanel struct store the mixpanel endpoint and the project token
type Mixpanel interface {
	// Create a mixpanel event
	Track(distinctId, eventName string, e *Event) error

	// Set properties for a mixpanel user.
	Update(distinctId string, u *Update) error

	Alias(distinctId, newId string) error
}

// The Mixapanel struct store the mixpanel endpoint and the project token
type mixpanel struct {
	Client *http.Client
	Token  string
	ApiURL string
}

// A mixpanel event
type Event struct {
	// IP-address of the user. Leave empty to use autodetect, or set to "0" to
	// not specify an ip-address.
	IP string

	// Timestamp. Set to nil to use the current time.
	Timestamp *time.Time

	// Custom properties. At least one must be specified.
	Properties map[string]interface{}
}

// An update of a user in mixpanel
type Update struct {
	// IP-address of the user. Leave empty to use autodetect, or set to "0" to
	// not specify an ip-address at all.
	IP string

	// Timestamp. Set to nil to use the current time, or IgnoreTime to not use a
	// timestamp.
	Timestamp *time.Time

	// Update operation such as "$set", "$update" etc.
	Operation string

	// Custom properties. At least one must be specified.
	Properties map[string]interface{}
}

// Track create a events to current distinct id
func (m *mixpanel) Alias(distinctId, newId string) error {
	props := map[string]interface{}{
		"token":       m.Token,
		"distinct_id": distinctId,
		"alias":       newId,
	}

	params := map[string]interface{}{
		"event":      "$create_alias",
		"properties": props,
	}

	return m.send("track", params, false)
}

// Track create a events to current distinct id
func (m *mixpanel) Track(distinctId, eventName string, e *Event) error {
	props := map[string]interface{}{
		"token":       m.Token,
		"distinct_id": distinctId,
	}
	if e.IP != "" {
		props["ip"] = e.IP
	}
	if e.Timestamp != nil {
		props["time"] = e.Timestamp.Unix()
	}

	for key, value := range e.Properties {
		props[key] = value
	}

	params := map[string]interface{}{
		"event":      eventName,
		"properties": props,
	}

	autoGeolocate := e.IP == ""

	return m.send("track", params, autoGeolocate)
}

// Updates a user in mixpanel. See
// https://mixpanel.com/help/reference/http#people-analytics-updates
func (m *mixpanel) Update(distinctId string, u *Update) error {
	params := map[string]interface{}{
		"$token":       m.Token,
		"$distinct_id": distinctId,
	}

	if u.IP != "" {
		params["$ip"] = u.IP
	}
	if u.Timestamp == IgnoreTime {
		params["$ignore_time"] = true
	} else if u.Timestamp != nil {
		params["$time"] = u.Timestamp.Unix()
	}

	params[u.Operation] = u.Properties

	autoGeolocate := u.IP == ""

	return m.send("engage", params, autoGeolocate)
}

func (m *mixpanel) to64(data string) string {
	bytes := []byte(data)
	return base64.StdEncoding.EncodeToString(bytes)
}

func (m *mixpanel) send(eventType string, params interface{}, autoGeolocate bool) error {
	dataJSON, _ := json.Marshal(params)
	data := string(dataJSON)

	url := m.ApiURL + "/" + eventType + "?data=" + m.to64(data)

	if autoGeolocate {
		url += "&ip=1"
	}

	resp, err := m.Client.Get(url)

	if err != nil {
		return fmt.Errorf("mixpanel: %s", err.Error())
	}

	defer resp.Body.Close()

	body, bodyErr := ioutil.ReadAll(resp.Body)

	if bodyErr != nil {
		return fmt.Errorf("mixpanel: %s", bodyErr.Error())
	}

	if strBody := string(body); strBody != "1" && strBody != "1\n" {
		return &ErrTrackFailed{Body: strBody, Resp: resp}
	}

	return nil
}

// New returns the client instance. If apiURL is blank, the default will be used
// ("https://api.mixpanel.com").
func New(token, apiURL string) Mixpanel {
	return NewFromClient(http.DefaultClient, token, apiURL)
}

// Creates a client instance using the specified client instance. This is useful
// when using a proxy.
func NewFromClient(c *http.Client, token, apiURL string) Mixpanel {
	if apiURL == "" {
		apiURL = "https://api.mixpanel.com"
	}

	return &mixpanel{
		Client: c,
		Token:  token,
		ApiURL: apiURL,
	}
}
<|endoftext|>


<commit_before>package main

/*

bulk command
force bulk insert mydata.csv

The load process involves these steps
	1. Create a job
	https://instance_name—api.salesforce.com/services/async/APIversion/job
		payload:
			<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 				<operation>insert</operation>
 				<object>Account</object>
 				<contentType>CSV</contentType>
			</jobInfo>
	2. Add batches to the created job
	https://instance_name—api.salesforce.com/services/async/APIversion/job/jobid/batch
		payload:
			<sObjects xmlns="http://www.force.com/2009/06/asyncapi/dataload">
  				<sObject>
    				<description>Created from Bulk API on Tue Apr 14 11:15:59 PDT 2009</description>
    				<name>[Bulk API] Account 0 (batch 0)</name>
  				</sObject>
  				<sObject>
    				<description>Created from Bulk API on Tue Apr 14 11:15:59 PDT 2009</description>
    				<name>[Bulk API] Account 1 (batch 0)</name>
  				</sObject>
			</sObjects>
	3. Close job (I assume this submits the job???)
	https://instance_name—api.salesforce.com/services/async/APIversion/job/jobId
		payload:
			<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 				<state>Closed</state>
			</jobInfo>

Jobs and batches can be monitored.

bulk command
force bulk job <jobId>

bulk command
force bulk batches <jobId>

bulk command
force bulk batch <batchId>

*/

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"strings"
)

var cmdBulk = &Command{
	Run:   runBulk,
	Usage: "bulk insert Account [csv file]",
	Short: "Load csv file use Bulk API",
	Long: `
Load csv file use Bulk API

Examples:

  force bulk insert Account [csv file]

  force bulk update Account [csv file]

  force bulk job [job id]

  force bulk batches [job id]

  force Bulk batch [job id] [batch id]

  force bulk batch retrieve [job id] [batch id]

  force bulk query Account [SOQL]

  force bulk query retrieve [job id] [batch id]
`,
}

func runBulk(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.printUsage()
	} else if len(args) == 1 {
		ErrorAndExit("Invalid command")
	} else if len(args) == 2 {
		if args[0] == "insert" {
			ErrorAndExit("Missing argument for insert")
		} else if args[0] == "update" {
			ErrorAndExit("Missing argument for update")
		} else if args[0] == "job" {
			showJobDetails(args[1])
		} else if args[0] == "batches" {
			listBatches(args[1])
		} else {
			ErrorAndExit("Invalid command")
		}
	} else if len(args) == 3 {
		if args[0] == "insert" {
			createBulkInsertJob(args[2], args[1], "CSV")
		} else if args[0] == "update" {
			createBulkUpdateJob(args[2], args[1], "CSV")
		} else if args[0] == "batch" {
			showBatchDetails(args[1], args[2])
		} else if args[0] == "query" {
			if args[1] == "retrieve" {
				ErrorAndExit("Query retrieve requires a job id and a batch id")
			} else {
				doBulkQuery(args[1], args[2], "CSV")
			}
		}
	} else if len(args) == 4 {
		if args[0] == "insert" {
			createBulkInsertJob(args[2], args[1], args[3])
		} else if args[0] == "update" {
			createBulkUpdateJob(args[2], args[1], args[3])
		} else if args[0] == "batch" {
			getBatchResults(args[2], args[3])
		} else if args[0] == "query" {
			if args[1] == "retrieve" {
				fmt.Println(string(getBulkQueryResults(args[2], args[3])))
			} else if args[1] == "status" {
				DisplayBatchInfo(getBatchDetails(args[2], args[3]))
			} else {
				doBulkQuery(args[1], args[2], args[3])
			}
		}
	}
}

func doBulkQuery(objectType string, soql string, contenttype string) {
	jobInfo, err := createBulkJob(objectType, "query", contenttype)
	force, _ := ActiveForce()

	result, err := force.BulkQuery(soql, jobInfo.Id, contenttype)
	if err != nil {
		closeBulkJob(jobInfo.Id)
		ErrorAndExit(err.Error())
	}
	fmt.Println("Query Submitted")
	fmt.Printf("To retrieve query status use\nforce bulk query status %s %s\n\n", jobInfo.Id, result.Id)
	fmt.Printf("To retrieve query data use\nforce bulk query retrieve %s %s\n\n", jobInfo.Id, result.Id)
	closeBulkJob(jobInfo.Id)
}

func getBulkQueryResults(jobId string, batchId string) (data []byte) {
	resultIds := retrieveBulkQuery(jobId, batchId)
	hasMultipleResultFiles := len(resultIds) > 1

	for _, resultId := range resultIds {
		//since this is going to stdOut, simply add header to separate "files"
		//if it's all in the same file, don't print this separator.
		if hasMultipleResultFiles {
			resultHeader := fmt.Sprint("ResultId: ", resultId, "\n")
			data = append(data[:], []byte(resultHeader)...)
		}
		//get next file, and append
		var newData []byte = retrieveBulkQueryResults(jobId, batchId, resultId)
		data = append(data[:], newData...)
	}

	return
}

func retrieveBulkQuery(jobId string, batchId string) (resultIds []string) {
	force, _ := ActiveForce()

	jobInfo, err := force.RetrieveBulkQuery(jobId, batchId)
	if err != nil {
		ErrorAndExit(err.Error())
	}

	var resultList struct {
		results []string `xml:"result"`
	}

	xml.Unmarshal(jobInfo, &resultList)
	resultIds = resultList.results
	return
}

func retrieveBulkQueryResults(jobId string, batchId string, resultId string) (data []byte) {
	force, _ := ActiveForce()

	data, err := force.RetrieveBulkQueryResults(jobId, batchId, resultId)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func showJobDetails(jobId string) {
	jobInfo := getJobDetails(jobId)
	DisplayJobInfo(jobInfo)
}

func listBatches(jobId string) {
	batchInfos := getBatches(jobId)
	DisplayBatchList(batchInfos)
}

func showBatchDetails(jobId string, batchId string) {
	batchInfo := getBatchDetails(jobId, batchId)
	DisplayBatchInfo(batchInfo)
}

func getBatchResults(jobId string, batchId string) {
	force, _ := ActiveForce()

	data, err := force.RetrieveBulkBatchResults(jobId, batchId)
	fmt.Println(data)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func getJobDetails(jobId string) (jobInfo JobInfo) {
	force, _ := ActiveForce()

	jobInfo, err := force.GetJobInfo(jobId)

	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func getBatches(jobId string) (batchInfos []BatchInfo) {
	force, _ := ActiveForce()

	batchInfos, err := force.GetBatches(jobId)

	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func getBatchDetails(jobId string, batchId string) (batchInfo BatchInfo) {
	force, _ := ActiveForce()

	batchInfo, err := force.GetBatchInfo(jobId, batchId)

	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func createBulkInsertJob(csvFilePath string, objectType string, format string) {
	jobInfo, err := createBulkJob(objectType, "insert", format)
	if err != nil {
		ErrorAndExit(err.Error())
	} else {
		batchInfo, err := addBatchToJob(csvFilePath, jobInfo.Id)
		if err != nil {
			closeBulkJob(jobInfo.Id)
			ErrorAndExit(err.Error())
		} else {
			closeBulkJob(jobInfo.Id)
			fmt.Printf("Job created ( %s ) - for job status use\n force bulk batch %s %s\n", jobInfo.Id, jobInfo.Id, batchInfo.Id)
		}
	}
}

func createBulkUpdateJob(csvFilePath string, objectType string, format string) {
	jobInfo, err := createBulkJob(objectType, "update", format)
	if err != nil {
		ErrorAndExit(err.Error())
	} else {
		batchInfo, err := addBatchToJob(csvFilePath, jobInfo.Id)
		if err != nil {
			closeBulkJob(jobInfo.Id)
			ErrorAndExit(err.Error())
		} else {
			closeBulkJob(jobInfo.Id)
			fmt.Printf("Job created ( %s ) - for job status use\n force bulk batch %s %s\n", jobInfo.Id, jobInfo.Id, batchInfo.Id)
		}
	}
}

func addBatchToJob(csvFilePath string, jobId string) (result BatchInfo, err error) {

	force, _ := ActiveForce()

	filedata, err := ioutil.ReadFile(csvFilePath)
	batches := splitFileIntoBatches(filedata)
	for b := range batches {
		result, err = force.AddBatchToJob(batches[b], jobId)
		if err != nil {
			break
		} else {
			fmt.Printf("Batch %d of %d added with Id %s \n", b+1, len(batches), result.Id)
		}
	}
	return
}

func splitFileIntoBatches(filedata []byte) (batches []string) {
	batchsize := 10000
	rows := strings.Split(string(filedata), "\n")
	headerRow, rows := rows[0], rows[1:]
	for len(rows) > 1 {
		if len(rows) < batchsize {
			batchsize = len(rows)
		}
		batch := []string{headerRow}
		batch = append(batch, rows[0:batchsize]...)
		batches = append(batches, strings.Join(batch, "\n"))
		rows = rows[batchsize:]
	}
	return
}

func getBatchInfo(jobId string, batchId string) (batchInfo BatchInfo, err error) {
	force, _ := ActiveForce()
	batchInfo, err = force.GetBatchInfo(jobId, batchId)
	return
}

func createBulkJob(objectType string, operation string, fileFormat string) (jobInfo JobInfo, err error) {
	force, _ := ActiveForce()

	xml := `
	<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 		<operation>%s</operation>
 		<object>%s</object>
 		<contentType>%s</contentType>
	</jobInfo>
	`
	data := fmt.Sprintf(xml, operation, objectType, fileFormat)
	jobInfo, err = force.CreateBulkJob(data)
	return
}

func closeBulkJob(jobId string) (jobInfo JobInfo, err error) {
	force, _ := ActiveForce()

	xml := `
	<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 		<state>Closed</state>
	</jobInfo>
	`
	jobInfo, err = force.CloseBulkJob(jobId, xml)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}
<commit_msg>Bugfix for `force bulk query retrieve` to parse resultIds<commit_after>package main

/*

bulk command
force bulk insert mydata.csv

The load process involves these steps
	1. Create a job
	https://instance_name—api.salesforce.com/services/async/APIversion/job
		payload:
			<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 				<operation>insert</operation>
 				<object>Account</object>
 				<contentType>CSV</contentType>
			</jobInfo>
	2. Add batches to the created job
	https://instance_name—api.salesforce.com/services/async/APIversion/job/jobid/batch
		payload:
			<sObjects xmlns="http://www.force.com/2009/06/asyncapi/dataload">
  				<sObject>
    				<description>Created from Bulk API on Tue Apr 14 11:15:59 PDT 2009</description>
    				<name>[Bulk API] Account 0 (batch 0)</name>
  				</sObject>
  				<sObject>
    				<description>Created from Bulk API on Tue Apr 14 11:15:59 PDT 2009</description>
    				<name>[Bulk API] Account 1 (batch 0)</name>
  				</sObject>
			</sObjects>
	3. Close job (I assume this submits the job???)
	https://instance_name—api.salesforce.com/services/async/APIversion/job/jobId
		payload:
			<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 				<state>Closed</state>
			</jobInfo>

Jobs and batches can be monitored.

bulk command
force bulk job <jobId>

bulk command
force bulk batches <jobId>

bulk command
force bulk batch <batchId>

*/

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"strings"
)

var cmdBulk = &Command{
	Run:   runBulk,
	Usage: "bulk insert Account [csv file]",
	Short: "Load csv file use Bulk API",
	Long: `
Load csv file use Bulk API

Examples:

  force bulk insert Account [csv file]

  force bulk update Account [csv file]

  force bulk job [job id]

  force bulk batches [job id]

  force Bulk batch [job id] [batch id]

  force bulk batch retrieve [job id] [batch id]

  force bulk query Account [SOQL]

  force bulk query retrieve [job id] [batch id]
`,
}

func runBulk(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.printUsage()
	} else if len(args) == 1 {
		ErrorAndExit("Invalid command")
	} else if len(args) == 2 {
		if args[0] == "insert" {
			ErrorAndExit("Missing argument for insert")
		} else if args[0] == "update" {
			ErrorAndExit("Missing argument for update")
		} else if args[0] == "job" {
			showJobDetails(args[1])
		} else if args[0] == "batches" {
			listBatches(args[1])
		} else {
			ErrorAndExit("Invalid command")
		}
	} else if len(args) == 3 {
		if args[0] == "insert" {
			createBulkInsertJob(args[2], args[1], "CSV")
		} else if args[0] == "update" {
			createBulkUpdateJob(args[2], args[1], "CSV")
		} else if args[0] == "batch" {
			showBatchDetails(args[1], args[2])
		} else if args[0] == "query" {
			if args[1] == "retrieve" {
				ErrorAndExit("Query retrieve requires a job id and a batch id")
			} else {
				doBulkQuery(args[1], args[2], "CSV")
			}
		}
	} else if len(args) == 4 {
		if args[0] == "insert" {
			createBulkInsertJob(args[2], args[1], args[3])
		} else if args[0] == "update" {
			createBulkUpdateJob(args[2], args[1], args[3])
		} else if args[0] == "batch" {
			getBatchResults(args[2], args[3])
		} else if args[0] == "query" {
			if args[1] == "retrieve" {
				fmt.Println(string(getBulkQueryResults(args[2], args[3])))
			} else if args[1] == "status" {
				DisplayBatchInfo(getBatchDetails(args[2], args[3]))
			} else {
				doBulkQuery(args[1], args[2], args[3])
			}
		}
	}
}

func doBulkQuery(objectType string, soql string, contenttype string) {
	jobInfo, err := createBulkJob(objectType, "query", contenttype)
	force, _ := ActiveForce()

	result, err := force.BulkQuery(soql, jobInfo.Id, contenttype)
	if err != nil {
		closeBulkJob(jobInfo.Id)
		ErrorAndExit(err.Error())
	}
	fmt.Println("Query Submitted")
	fmt.Printf("To retrieve query status use\nforce bulk query status %s %s\n\n", jobInfo.Id, result.Id)
	fmt.Printf("To retrieve query data use\nforce bulk query retrieve %s %s\n\n", jobInfo.Id, result.Id)
	closeBulkJob(jobInfo.Id)
}

func getBulkQueryResults(jobId string, batchId string) (data []byte) {
	resultIds := retrieveBulkQuery(jobId, batchId)
	hasMultipleResultFiles := len(resultIds) > 1

	for _, resultId := range resultIds {
		//since this is going to stdOut, simply add header to separate "files"
		//if it's all in the same file, don't print this separator.
		if hasMultipleResultFiles {
			resultHeader := fmt.Sprint("ResultId: ", resultId, "\n")
			data = append(data[:], []byte(resultHeader)...)
		}
		//get next file, and append
		var newData []byte = retrieveBulkQueryResults(jobId, batchId, resultId)
		data = append(data[:], newData...)
	}

	return
}

func retrieveBulkQuery(jobId string, batchId string) (resultIds []string) {
	force, _ := ActiveForce()

	jobInfo, err := force.RetrieveBulkQuery(jobId, batchId)
	if err != nil {
		ErrorAndExit(err.Error())
	}

	var resultList struct {
		Results []string `xml:"result"`
	}

	xml.Unmarshal(jobInfo, &resultList)
	resultIds = resultList.Results
	return
}

func retrieveBulkQueryResults(jobId string, batchId string, resultId string) (data []byte) {
	force, _ := ActiveForce()

	data, err := force.RetrieveBulkQueryResults(jobId, batchId, resultId)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func showJobDetails(jobId string) {
	jobInfo := getJobDetails(jobId)
	DisplayJobInfo(jobInfo)
}

func listBatches(jobId string) {
	batchInfos := getBatches(jobId)
	DisplayBatchList(batchInfos)
}

func showBatchDetails(jobId string, batchId string) {
	batchInfo := getBatchDetails(jobId, batchId)
	DisplayBatchInfo(batchInfo)
}

func getBatchResults(jobId string, batchId string) {
	force, _ := ActiveForce()

	data, err := force.RetrieveBulkBatchResults(jobId, batchId)
	fmt.Println(data)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func getJobDetails(jobId string) (jobInfo JobInfo) {
	force, _ := ActiveForce()

	jobInfo, err := force.GetJobInfo(jobId)

	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func getBatches(jobId string) (batchInfos []BatchInfo) {
	force, _ := ActiveForce()

	batchInfos, err := force.GetBatches(jobId)

	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func getBatchDetails(jobId string, batchId string) (batchInfo BatchInfo) {
	force, _ := ActiveForce()

	batchInfo, err := force.GetBatchInfo(jobId, batchId)

	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}

func createBulkInsertJob(csvFilePath string, objectType string, format string) {
	jobInfo, err := createBulkJob(objectType, "insert", format)
	if err != nil {
		ErrorAndExit(err.Error())
	} else {
		batchInfo, err := addBatchToJob(csvFilePath, jobInfo.Id)
		if err != nil {
			closeBulkJob(jobInfo.Id)
			ErrorAndExit(err.Error())
		} else {
			closeBulkJob(jobInfo.Id)
			fmt.Printf("Job created ( %s ) - for job status use\n force bulk batch %s %s\n", jobInfo.Id, jobInfo.Id, batchInfo.Id)
		}
	}
}

func createBulkUpdateJob(csvFilePath string, objectType string, format string) {
	jobInfo, err := createBulkJob(objectType, "update", format)
	if err != nil {
		ErrorAndExit(err.Error())
	} else {
		batchInfo, err := addBatchToJob(csvFilePath, jobInfo.Id)
		if err != nil {
			closeBulkJob(jobInfo.Id)
			ErrorAndExit(err.Error())
		} else {
			closeBulkJob(jobInfo.Id)
			fmt.Printf("Job created ( %s ) - for job status use\n force bulk batch %s %s\n", jobInfo.Id, jobInfo.Id, batchInfo.Id)
		}
	}
}

func addBatchToJob(csvFilePath string, jobId string) (result BatchInfo, err error) {

	force, _ := ActiveForce()

	filedata, err := ioutil.ReadFile(csvFilePath)
	batches := splitFileIntoBatches(filedata)
	for b := range batches {
		result, err = force.AddBatchToJob(batches[b], jobId)
		if err != nil {
			break
		} else {
			fmt.Printf("Batch %d of %d added with Id %s \n", b+1, len(batches), result.Id)
		}
	}
	return
}

func splitFileIntoBatches(filedata []byte) (batches []string) {
	batchsize := 10000
	rows := strings.Split(string(filedata), "\n")
	headerRow, rows := rows[0], rows[1:]
	for len(rows) > 1 {
		if len(rows) < batchsize {
			batchsize = len(rows)
		}
		batch := []string{headerRow}
		batch = append(batch, rows[0:batchsize]...)
		batches = append(batches, strings.Join(batch, "\n"))
		rows = rows[batchsize:]
	}
	return
}

func getBatchInfo(jobId string, batchId string) (batchInfo BatchInfo, err error) {
	force, _ := ActiveForce()
	batchInfo, err = force.GetBatchInfo(jobId, batchId)
	return
}

func createBulkJob(objectType string, operation string, fileFormat string) (jobInfo JobInfo, err error) {
	force, _ := ActiveForce()

	xml := `
	<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 		<operation>%s</operation>
 		<object>%s</object>
 		<contentType>%s</contentType>
	</jobInfo>
	`
	data := fmt.Sprintf(xml, operation, objectType, fileFormat)
	jobInfo, err = force.CreateBulkJob(data)
	return
}

func closeBulkJob(jobId string) (jobInfo JobInfo, err error) {
	force, _ := ActiveForce()

	xml := `
	<jobInfo xmlns="http://www.force.com/2009/06/asyncapi/dataload">
 		<state>Closed</state>
	</jobInfo>
	`
	jobInfo, err = force.CloseBulkJob(jobId, xml)
	if err != nil {
		ErrorAndExit(err.Error())
	}
	return
}
<|endoftext|>


<commit_before>package backoff

import (
	"math"
	"time"
)

type ExponentialBackoff struct {
	Retries    int
	MaxRetries int
	Delay      time.Duration
	Interval   time.Duration // time.Second, time.Millisecond, etc.
}

// Exponential creates a new instance of ExponentialBackoff.
func Exponential() *ExponentialBackoff {
	return &ExponentialBackoff{
		Retries:    0,
		MaxRetries: 5,
		Delay:      time.Duration(0),
		Interval:   time.Duration(1 * time.Second),
	}
}

/*
Next gets the next backoff delay. This method will increment the retries and check
if the maximum number of retries has been met. If this condition is satisfied, then
the function will return. Otherwise, the next backoff delay will be computed.

The exponential backoff delay is computed as follows:
`n = 2^c - 1` where `n` is the backoff delay and `c` is the number of retries.

Example, given a 1 second interval:

  Retry #        Backoff delay (in seconds)
    0                   0
    1                   1
    2                   3
    3                   7
    4                   15
*/
func (self *ExponentialBackoff) Next() bool {
	self.Retries++

	if self.Retries >= self.MaxRetries {
		return false
	}

	self.Delay = time.Duration(math.Pow(2, float64(self.Retries))-1) * self.Interval

	return true
}

/*
Retry will retry a function until the maximum number of retries is met. This method expects
the function `f` to return an error. If the failure condition is met, this method
will surface the error outputted from `f`, otherwise nil will be returned as normal.
*/
func (self *ExponentialBackoff) Retry(f func() error) error {
	err := f()

	if err == nil {
		return nil
	}

	for self.Next() {
		if err = f(); err == nil {
			return nil
		}

		time.Sleep(self.Delay)
	}

	return err
}

// Reset will reset the retry count and the backoff delay back to its initial state.
func (self *ExponentialBackoff) Reset() {
	self.Retries = 0
	self.Delay = time.Duration(0 * time.Second)
}
<commit_msg>Ensure lint passes for Exponential.<commit_after>package backoff

import (
	"math"
	"time"
)

type ExponentialBackoff struct {
	Retries    int
	MaxRetries int
	Delay      time.Duration
	Interval   time.Duration // time.Second, time.Millisecond, etc.
}

// Exponential creates a new instance of ExponentialBackoff.
func Exponential() *ExponentialBackoff {
	return &ExponentialBackoff{
		Retries:    0,
		MaxRetries: 5,
		Delay:      time.Duration(0),
		Interval:   time.Duration(1 * time.Second),
	}
}

/*
Next gets the next backoff delay. This method will increment the retries and check
if the maximum number of retries has been met. If this condition is satisfied, then
the function will return. Otherwise, the next backoff delay will be computed.

The exponential backoff delay is computed as follows:
`n = 2^c - 1` where `n` is the backoff delay and `c` is the number of retries.

Example, given a 1 second interval:

  Retry #        Backoff delay (in seconds)
    0                   0
    1                   1
    2                   3
    3                   7
    4                   15
*/
func (e *ExponentialBackoff) Next() bool {
	e.Retries++

	if e.Retries >= e.MaxRetries {
		return false
	}

	e.Delay = time.Duration(math.Pow(2, float64(e.Retries))-1) * e.Interval

	return true
}

/*
Retry will retry a function until the maximum number of retries is met. This method expects
the function `f` to return an error. If the failure condition is met, this method
will surface the error outputted from `f`, otherwise nil will be returned as normal.
*/
func (e *ExponentialBackoff) Retry(f func() error) error {
	err := f()

	if err == nil {
		return nil
	}

	for e.Next() {
		if err = f(); err == nil {
			return nil
		}

		time.Sleep(e.Delay)
	}

	return err
}

// Reset will reset the retry count and the backoff delay back to its initial state.
func (e *ExponentialBackoff) Reset() {
	e.Retries = 0
	e.Delay = time.Duration(0 * time.Second)
}
<|endoftext|>


<commit_before>package semaphore

import (
	"sync"
	"time"
)

type Semaphore struct {
	permits int
	avail   int
	channel chan int
	aMutex  *sync.Mutex
	rMutex  *sync.Mutex
}

func New(permits int) *Semaphore {
	if permits < 1 {
		panic("Invalid number of permits. Less than 1")
	}
	return &Semaphore{
		permits,
		permits,
		make(chan int, permits),
		&sync.Mutex{},
		&sync.Mutex{},
	}
}

//Acquire one permit, if its not available the goroutine will block till its available
func (s *Semaphore) Acquire() {
	s.aMutex.Lock()
	s.channel <- 1
	s.avail--
	s.aMutex.Unlock()
}

//Similar to Acquire() but for many permits
func (s *Semaphore) AcquireMany(n int) {
	if n > s.permits {
		panic("To many requested permits")
	}
	s.aMutex.Lock()
	s.avail -= n
	for ; n > 0; n-- {
		s.channel <- 1
	}
	s.avail += n
	s.aMutex.Unlock()
}

//Similar to AcquireMany() but cancels if duration elapse before getting the permits.
//Returns true if successful and false if timeout occurs.
func (s *Semaphore) AcquireWithin(n int, d time.Duration) bool {
	timeout := make(chan bool, 1)
	cancel := make(chan bool, 1)
	go func() {
		time.Sleep(d)
		timeout <- true
	}()
	go func() {
		s.AcquireMany(n)
		timeout <- false
		if <-cancel {
			s.ReleaseMany(n)
		}
	}()
	if <-timeout {
		cancel <- true
		return false
	}
	cancel <- false
	return true
}

//Release one permit
func (s *Semaphore) Release() {
	s.rMutex.Lock()
	<-s.channel
	s.avail++
	s.rMutex.Unlock()
}

//Release many permits
func (s *Semaphore) ReleaseMany(n int) {
	if n > s.permits {
		panic("Too many requested releases")
	}
	for ; n > 0; n-- {
		s.Release()
	}
}

//Number of available unacquired permits
func (s *Semaphore) AvailablePermits() int {
	if s.avail < 0 {
		return 0
	}
	return s.avail
}

//Acquire all available permits and return the number of permits acquired
func (s *Semaphore) DrainPermits() int {
	n := s.AvailablePermits()
	if n > 0 {
		s.AcquireMany(n)
	}
	return n
}
<commit_msg>Improve documentation<commit_after>package semaphore

import (
	"errors"
	"sync"
	"time"
)

// Semaphore is an implementation of semaphore.
type Semaphore struct {
	permits int
	avail   int
	channel chan struct{}
	aMutex  *sync.RWMutex
	rMutex  *sync.Mutex
}

// New creates a new Semaphore with specified number of permits.
func New(permits int) *Semaphore {
	if permits < 1 {
		panic("Invalid number of permits. Less than 1")
	}
	return &Semaphore{
		permits,
		permits,
		make(chan struct{}, permits),
		&sync.RWMutex{},
		&sync.Mutex{},
	}
}

// Acquire acquires one permit. If it is not available, the goroutine will block until it is available.
func (s *Semaphore) Acquire() {
	s.aMutex.Lock()
	defer s.aMutex.Unlock()

	s.channel <- struct{}{}
	s.avail--
}

// AcquireMany is similar to Acquire() but for many permits.
// An error is returned if n is greater number of permits in the semaphore.
func (s *Semaphore) AcquireMany(n int) error {
	if n > s.permits {
		return errors.New("Too many requested permits")
	}
	s.aMutex.Lock()
	defer s.aMutex.Unlock()

	s.avail -= n
	for ; n > 0; n-- {
		s.channel <- struct{}{}
	}
	s.avail += n
	return nil
}

// AcquireWithin is similar to AcquireMany() but cancels if duration elapses before getting the permits.
// Returns true if successful and false if timeout occurs.
func (s *Semaphore) AcquireWithin(n int, d time.Duration) bool {
	timeout := make(chan bool, 1)
	cancel := make(chan bool, 1)
	go func() {
		time.Sleep(d)
		timeout <- true
	}()
	go func() {
		s.AcquireMany(n)
		timeout <- false
		if <-cancel {
			s.ReleaseMany(n)
		}
	}()
	if <-timeout {
		cancel <- true
		return false
	}
	cancel <- false
	return true
}

// Release releases one permit.
func (s *Semaphore) Release() {
	s.rMutex.Lock()
	defer s.rMutex.Unlock()

	<-s.channel
	s.avail++
}

// ReleaseMany releases n permits.
func (s *Semaphore) ReleaseMany(n int) {
	if n > s.permits {
		panic("Too many requested releases")
	}
	for ; n > 0; n-- {
		s.Release()
	}
}

// AvailablePermits gives number of available unacquired permits.
func (s *Semaphore) AvailablePermits() int {
	s.aMutex.RLock()
	defer s.aMutex.RUnlock()

	if s.avail < 0 {
		return 0
	}
	return s.avail
}

// DrainPermits acquires all available permits and return the number of permits acquired.
func (s *Semaphore) DrainPermits() int {
	n := s.AvailablePermits()
	if n > 0 {
		s.AcquireMany(n)
	}
	return n
}
<|endoftext|>


<commit_before>package seq

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

// ID is a 64bits wide unsigned identifier.
type ID uint64

// IDStream is a read-only channel of `ID`s.
//
// An `IDStream` should always starts at 1.
type IDStream <-chan ID

// Next is a simple helper to get the next `ID` from an `IDStream`.
//
// It can be helpful if you cannot simply range on the stream for some reason.
//
// `Next` blocks if no `ID` is available, and returns `0` if the stream is
// already closed.
func (ids IDStream) Next() ID {
	select {
	case id := <-ids:
		return id
	}
}

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

// A Sequencer exposes methods to fetch sequential `ID` streams.
type Sequencer interface {
	// GetStream returns the `IDStream` associated with the `Sequencer`.
	GetStream() IDStream
	// Close stops the `Sequencer` and cleans associated resources.
	Close() error
}
<commit_msg>Implemented sort.Interface for ID slices<commit_after>package seq

import "sort"

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

// ID is a 64bits wide unsigned identifier.
type ID uint64

// IDSlice implements `sort.Interface` for slices of `ID`s.`
type IDSlice []ID

func (ids IDSlice) Len() int           { return len(ids) }
func (ids IDSlice) Less(i, j int) bool { return ids[i] < ids[j] }
func (ids IDSlice) Swap(i, j int)      { ids[i], ids[j] = ids[j], ids[i] }
func (ids IDSlice) Sort() IDSlice      { sort.Sort(ids); return ids }

// IDStream is a read-only channel of `ID`s.
//
// An `IDStream` should always starts at 1.
type IDStream <-chan ID

// Next is a simple helper to get the next `ID` from an `IDStream`.
//
// It can be helpful if you cannot simply range on the stream for some reason.
//
// `Next` blocks if no `ID` is available, and returns `0` if the stream is
// already closed.
func (ids IDStream) Next() ID {
	select {
	case id := <-ids:
		return id
	}
}

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

// A Sequencer exposes methods to fetch sequential `ID` streams.
type Sequencer interface {
	// GetStream returns the `IDStream` associated with the `Sequencer`.
	GetStream() IDStream
	// Close stops the `Sequencer` and cleans associated resources.
	Close() error
}
<|endoftext|>


<commit_before>// (C) 2014 Cybozu.  All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file.

package kintone

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"
)

// Record represens a record in an application.
//
// Fields is a mapping between field IDs and fields.
// Although field types are shown as interface{}, they are guaranteed
// to be one of a *Field type in this package.
type Record struct {
	id       uint64
	revision int64
	Fields   map[string]interface{}
}

// NewRecord creates an instance of Record.
//
// The revision number is initialized to -1.
func NewRecord(fields map[string]interface{}) *Record {
	return &Record{0, -1, fields}
}

// MarshalJSON marshals field data of a record into JSON.
func (rec Record) MarshalJSON() ([]byte, error) {
	return json.Marshal(rec.Fields)
}

// Id returns the record number.
//
// A record number is unique within an application.
func (rec Record) Id() uint64 {
	return rec.id
}

// Revision returns the record revision number.
func (rec Record) Revision() int64 {
	return rec.revision
}

// Assert string list.
func stringList(l []interface{}) []string {
	sl := make([]string, len(l))
	for i, v := range l {
		sl[i] = v.(string)
	}
	return sl
}

// Convert user list.
func userList(l []interface{}) ([]User, error) {
	b, err := json.Marshal(l)
	if err != nil {
		return nil, err
	}
	var ul []User
	err = json.Unmarshal(b, &ul)
	if err != nil {
		return nil, err
	}
	return ul, nil
}

// Convert string "record number" into an integer.
func numericId(id string) (uint64, error) {
	n := strings.LastIndex(id, "-")
	if n != -1 {
		id = id[(n + 1):]
	}
	nid, err := strconv.ParseUint(id, 10, 64)
	if err != nil {
		return 0, err
	}
	return nid, nil
}

type recordData map[string]struct {
	Type  string      `json:"type"`
	Value interface{} `json:"value"`
}

func decodeRecordData(data recordData) (*Record, error) {
	fields := make(map[string]interface{})
	rec := &Record{0, -1, fields}
	for key, v := range data {
		switch v.Type {
		case FT_SINGLE_LINE_TEXT:
			fields[key] = SingleLineTextField(v.Value.(string))
		case FT_MULTI_LINE_TEXT:
			fields[key] = MultiLineTextField(v.Value.(string))
		case FT_RICH_TEXT:
			fields[key] = RichTextField(v.Value.(string))
		case FT_DECIMAL:
			fields[key] = DecimalField(v.Value.(string))
		case FT_CALC:
			fields[key] = CalcField(v.Value.(string))
		case FT_CHECK_BOX:
			fields[key] = CheckBoxField(stringList(v.Value.([]interface{})))
		case FT_RADIO:
			fields[key] = RadioButtonField(v.Value.(string))
		case FT_SINGLE_SELECT:
			if v.Value == nil {
				fields[key] = SingleSelectField{Valid: false}
			} else {
				fields[key] = SingleSelectField{v.Value.(string), true}
			}
		case FT_MULTI_SELECT:
			fields[key] = MultiSelectField(stringList(v.Value.([]interface{})))
		case FT_FILE:
			b1, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var fl []File
			err = json.Unmarshal(b1, &fl)
			if err != nil {
				return nil, err
			}
			fields[key] = FileField(fl)
		case FT_LINK:
			fields[key] = LinkField(v.Value.(string))
		case FT_DATE:
			if v.Value == nil {
				fields[key] = DateField{Valid: false}
			} else {
				d, err := time.Parse("2006-01-02", v.Value.(string))
				if err != nil {
					return nil, fmt.Errorf("Invalid date: %v", v.Value)
				}
				fields[key] = DateField{d, true}
			}
		case FT_TIME:
			if v.Value == nil {
				fields[key] = TimeField{Valid: false}
			} else {
				t, err := time.Parse("15:04", v.Value.(string))
				if err != nil {
					t, err = time.Parse("15:04:05", v.Value.(string))
					if err != nil {
						return nil, fmt.Errorf("Invalid time: %v", v.Value)
					}
				}
				fields[key] = TimeField{t, true}
			}
		case FT_DATETIME:
			if s, ok := v.Value.(string); ok {
				dt, err := time.Parse(time.RFC3339, s)
				if err != nil {
					return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
				}
				fields[key] = DateTimeField(dt)
			}
		case FT_USER:
			ul, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = UserField(ul)
		case FT_CATEGORY:
			fields[key] = CategoryField(stringList(v.Value.([]interface{})))
		case FT_STATUS:
			fields[key] = StatusField(v.Value.(string))
		case FT_ASSIGNEE:
			al, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = AssigneeField(al)
		case FT_RECNUM:
			if nid, err := numericId(v.Value.(string)); err != nil {
				return nil, err
			} else {
				rec.id = nid
			}
			fields[key] = RecordNumberField(v.Value.(string))
		case FT_CREATOR:
			creator := v.Value.(map[string]interface{})
			fields[key] = CreatorField{
				creator["code"].(string),
				creator["name"].(string),
			}
		case FT_CTIME:
			var ctime time.Time
			if ctime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(ctime)
		case FT_MODIFIER:
			modifier := v.Value.(map[string]interface{})
			fields[key] = ModifierField{
				modifier["code"].(string),
				modifier["name"].(string),
			}
		case FT_MTIME:
			var mtime time.Time
			if mtime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(mtime)
		case FT_SUBTABLE:
			b2, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var stl []SubTableEntry
			err = json.Unmarshal(b2, &stl)
			if err != nil {
				return nil, err
			}
			fields[key] = SubTableField(stl)
		case FT_ID:
			id, err := strconv.ParseUint(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid record ID: %v", v.Value)
			}
			rec.id = id
		case FT_REVISION:
			revision, err := strconv.ParseInt(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid revision number: %v", v.Value)
			}
			rec.revision = revision
		default:
			return nil, fmt.Errorf("Invalid type: %v", v.Type)
		}
	}
	return rec, nil
}

// DecodeRecords decodes JSON response for multi-get API.
func DecodeRecords(b []byte) ([]*Record, error) {
	var t struct {
		Records []recordData `json:"records"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec_list := make([]*Record, len(t.Records))
	for i, rd := range t.Records {
		r, err := decodeRecordData(rd)
		if err != nil {
			return nil, err
		}
		rec_list[i] = r
	}
	return rec_list, nil
}

// DecodeRecord decodes JSON response for single-get API.
func DecodeRecord(b []byte) (*Record, error) {
	var t struct {
		RecordData recordData `json:"record"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec, err := decodeRecordData(t.RecordData)
	if err != nil {
		return nil, err
	}
	return rec, nil
}
<commit_msg>Ignore unknown field type.<commit_after>// (C) 2014 Cybozu.  All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file.

package kintone

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"
)

// Record represens a record in an application.
//
// Fields is a mapping between field IDs and fields.
// Although field types are shown as interface{}, they are guaranteed
// to be one of a *Field type in this package.
type Record struct {
	id       uint64
	revision int64
	Fields   map[string]interface{}
}

// NewRecord creates an instance of Record.
//
// The revision number is initialized to -1.
func NewRecord(fields map[string]interface{}) *Record {
	return &Record{0, -1, fields}
}

// MarshalJSON marshals field data of a record into JSON.
func (rec Record) MarshalJSON() ([]byte, error) {
	return json.Marshal(rec.Fields)
}

// Id returns the record number.
//
// A record number is unique within an application.
func (rec Record) Id() uint64 {
	return rec.id
}

// Revision returns the record revision number.
func (rec Record) Revision() int64 {
	return rec.revision
}

// Assert string list.
func stringList(l []interface{}) []string {
	sl := make([]string, len(l))
	for i, v := range l {
		sl[i] = v.(string)
	}
	return sl
}

// Convert user list.
func userList(l []interface{}) ([]User, error) {
	b, err := json.Marshal(l)
	if err != nil {
		return nil, err
	}
	var ul []User
	err = json.Unmarshal(b, &ul)
	if err != nil {
		return nil, err
	}
	return ul, nil
}

// Convert string "record number" into an integer.
func numericId(id string) (uint64, error) {
	n := strings.LastIndex(id, "-")
	if n != -1 {
		id = id[(n + 1):]
	}
	nid, err := strconv.ParseUint(id, 10, 64)
	if err != nil {
		return 0, err
	}
	return nid, nil
}

type recordData map[string]struct {
	Type  string      `json:"type"`
	Value interface{} `json:"value"`
}

func decodeRecordData(data recordData) (*Record, error) {
	fields := make(map[string]interface{})
	rec := &Record{0, -1, fields}
	for key, v := range data {
		switch v.Type {
		case FT_SINGLE_LINE_TEXT:
			fields[key] = SingleLineTextField(v.Value.(string))
		case FT_MULTI_LINE_TEXT:
			fields[key] = MultiLineTextField(v.Value.(string))
		case FT_RICH_TEXT:
			fields[key] = RichTextField(v.Value.(string))
		case FT_DECIMAL:
			fields[key] = DecimalField(v.Value.(string))
		case FT_CALC:
			fields[key] = CalcField(v.Value.(string))
		case FT_CHECK_BOX:
			fields[key] = CheckBoxField(stringList(v.Value.([]interface{})))
		case FT_RADIO:
			fields[key] = RadioButtonField(v.Value.(string))
		case FT_SINGLE_SELECT:
			if v.Value == nil {
				fields[key] = SingleSelectField{Valid: false}
			} else {
				fields[key] = SingleSelectField{v.Value.(string), true}
			}
		case FT_MULTI_SELECT:
			fields[key] = MultiSelectField(stringList(v.Value.([]interface{})))
		case FT_FILE:
			b1, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var fl []File
			err = json.Unmarshal(b1, &fl)
			if err != nil {
				return nil, err
			}
			fields[key] = FileField(fl)
		case FT_LINK:
			fields[key] = LinkField(v.Value.(string))
		case FT_DATE:
			if v.Value == nil {
				fields[key] = DateField{Valid: false}
			} else {
				d, err := time.Parse("2006-01-02", v.Value.(string))
				if err != nil {
					return nil, fmt.Errorf("Invalid date: %v", v.Value)
				}
				fields[key] = DateField{d, true}
			}
		case FT_TIME:
			if v.Value == nil {
				fields[key] = TimeField{Valid: false}
			} else {
				t, err := time.Parse("15:04", v.Value.(string))
				if err != nil {
					t, err = time.Parse("15:04:05", v.Value.(string))
					if err != nil {
						return nil, fmt.Errorf("Invalid time: %v", v.Value)
					}
				}
				fields[key] = TimeField{t, true}
			}
		case FT_DATETIME:
			if s, ok := v.Value.(string); ok {
				dt, err := time.Parse(time.RFC3339, s)
				if err != nil {
					return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
				}
				fields[key] = DateTimeField(dt)
			}
		case FT_USER:
			ul, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = UserField(ul)
		case FT_CATEGORY:
			fields[key] = CategoryField(stringList(v.Value.([]interface{})))
		case FT_STATUS:
			fields[key] = StatusField(v.Value.(string))
		case FT_ASSIGNEE:
			al, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = AssigneeField(al)
		case FT_RECNUM:
			if nid, err := numericId(v.Value.(string)); err != nil {
				return nil, err
			} else {
				rec.id = nid
			}
			fields[key] = RecordNumberField(v.Value.(string))
		case FT_CREATOR:
			creator := v.Value.(map[string]interface{})
			fields[key] = CreatorField{
				creator["code"].(string),
				creator["name"].(string),
			}
		case FT_CTIME:
			var ctime time.Time
			if ctime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(ctime)
		case FT_MODIFIER:
			modifier := v.Value.(map[string]interface{})
			fields[key] = ModifierField{
				modifier["code"].(string),
				modifier["name"].(string),
			}
		case FT_MTIME:
			var mtime time.Time
			if mtime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(mtime)
		case FT_SUBTABLE:
			b2, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var stl []SubTableEntry
			err = json.Unmarshal(b2, &stl)
			if err != nil {
				return nil, err
			}
			fields[key] = SubTableField(stl)
		case FT_ID:
			id, err := strconv.ParseUint(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid record ID: %v", v.Value)
			}
			rec.id = id
		case FT_REVISION:
			revision, err := strconv.ParseInt(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid revision number: %v", v.Value)
			}
			rec.revision = revision
		default:
			log.Printf("Unknown type: %v", v.Type)
		}
	}
	return rec, nil
}

// DecodeRecords decodes JSON response for multi-get API.
func DecodeRecords(b []byte) ([]*Record, error) {
	var t struct {
		Records []recordData `json:"records"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec_list := make([]*Record, len(t.Records))
	for i, rd := range t.Records {
		r, err := decodeRecordData(rd)
		if err != nil {
			return nil, err
		}
		rec_list[i] = r
	}
	return rec_list, nil
}

// DecodeRecord decodes JSON response for single-get API.
func DecodeRecord(b []byte) (*Record, error) {
	var t struct {
		RecordData recordData `json:"record"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec, err := decodeRecordData(t.RecordData)
	if err != nil {
		return nil, err
	}
	return rec, nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package core

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/loadimpact/k6/stats"
	"github.com/sirupsen/logrus"
	"gopkg.in/guregu/null.v3"
)

const (
	TickRate        = 1 * time.Millisecond
	MetricsRate     = 1 * time.Second
	CollectRate     = 50 * time.Millisecond
	ThresholdsRate  = 2 * time.Second
	ShutdownTimeout = 10 * time.Second

	BackoffAmount = 50 * time.Millisecond
	BackoffMax    = 10 * time.Second
)

// The Engine is the beating heart of k6.
type Engine struct {
	// TODO: Make most of the stuff here private! And think how to refactor the
	// engine to be less stateful... it's currently one big mess of moving
	// pieces, and you implicitly first have to call Init() and then Run() -
	// maybe we should refactor it so we have a `Session` dauther-object that
	// Init() returns? The only problem with doing this is the REST API - it
	// expects to be able to get information from the Engine and is initialized
	// before the Init() call...

	ExecutionScheduler lib.ExecutionScheduler
	executionState     *lib.ExecutionState

	Options       lib.Options
	Collectors    []lib.Collector
	NoThresholds  bool
	NoSummary     bool
	SummaryExport bool

	logger   *logrus.Entry
	stopChan chan struct{}

	Metrics     map[string]*stats.Metric
	MetricsLock sync.Mutex

	Samples chan stats.SampleContainer

	// Assigned to metrics upon first received sample.
	thresholds map[string]stats.Thresholds
	submetrics map[string][]*stats.Submetric

	// Are thresholds tainted?
	thresholdsTainted bool
}

// NewEngine instantiates a new Engine, without doing any heavy initialization.
func NewEngine(ex lib.ExecutionScheduler, o lib.Options, logger *logrus.Logger) (*Engine, error) {
	if ex == nil {
		return nil, errors.New("missing ExecutionScheduler instance")
	}

	e := &Engine{
		ExecutionScheduler: ex,
		executionState:     ex.GetState(),

		Options:  o,
		Metrics:  make(map[string]*stats.Metric),
		Samples:  make(chan stats.SampleContainer, o.MetricSamplesBufferSize.Int64),
		stopChan: make(chan struct{}),
		logger:   logger.WithField("component", "engine"),
	}

	e.thresholds = o.Thresholds
	e.submetrics = make(map[string][]*stats.Submetric)
	for name := range e.thresholds {
		if !strings.Contains(name, "{") {
			continue
		}

		parent, sm := stats.NewSubmetric(name)
		e.submetrics[parent] = append(e.submetrics[parent], sm)
	}

	return e, nil
}

// Init is used to initialize the execution scheduler and all metrics processing
// in the engine. The first is a costly operation, since it initializes all of
// the planned VUs and could potentially take a long time. It either returns an
// error immediately, or it returns test Run() and WindDown() functions.
//
// Things to note:
//  - The first lambda, Run(), synchronously executes the actual load test.
//  - It can be prematurely aborted by cancelling the runCtx - this won't stop
//    the metrics collection by the Engine.
//  - Stopping the metrics collection can be done at any time after Run() has
//    returned by cancelling the globalCtx
//  - The second returned lambda can be used to wait for that process to finish.
func (e *Engine) Init(globalCtx, runCtx context.Context) (run func() error, wait func(), err error) {
	e.logger.Debug("Initialization starting...")
	// TODO: if we ever need metrics processing in the init context, we can move
	// this below the other components... or even start them concurrently?
	if err := e.ExecutionScheduler.Init(runCtx, e.Samples); err != nil {
		return nil, nil, err
	}

	// TODO: move all of this in a separate struct? see main TODO above

	runSubCtx, runSubCancel := context.WithCancel(runCtx)

	resultCh := make(chan error)
	runFn := func() error {
		e.logger.Debug("Execution scheduler starting...")
		err := e.ExecutionScheduler.Run(globalCtx, runSubCtx, e.Samples)
		e.logger.WithError(err).Debug("Execution scheduler terminated")

		select {
		case <-runSubCtx.Done():
			// do nothing, the test run was aborted somehow
		default:
			resultCh <- err // we finished normally, so send the result
		}

		return err
	}
	waitFn := e.startBackgroundProcesses(globalCtx, runCtx, resultCh, runSubCancel)
	return runFn, waitFn, nil
}

// This starts a bunch of goroutines to process metrics, thresholds, and set the
// test run status when it ends. It returns a function that can be used after
// the provided context is called, to wait for the complete winding down of all
// started goroutines.
func (e *Engine) startBackgroundProcesses( //nolint:funlen
	globalCtx, runCtx context.Context, runResult <-chan error, runSubCancel func(),
) (wait func()) {
	processes := new(sync.WaitGroup)

	// Spin up all configured collectors
	for _, collector := range e.Collectors {
		processes.Add(1)
		go func(collector lib.Collector) {
			collector.Run(globalCtx)
			processes.Done()
		}(collector)
	}

	// Siphon and handle all produced metric samples
	processes.Add(1)
	go func() {
		defer processes.Done()
		e.processMetrics(globalCtx)
	}()

	// Run VU metrics emission, only while the test is running.
	// TODO: move? this seems like something the ExecutionScheduler should emit...
	processes.Add(1)
	go func() {
		defer processes.Done()
		e.logger.Debug("Starting emission of VU metrics...")
		e.runMetricsEmission(runCtx)
		e.logger.Debug("Metrics emission terminated")
	}()

	// Update the test run status when the test finishes
	processes.Add(1)
	thresholdAbortChan := make(chan struct{})
	go func() {
		defer processes.Done()
		select {
		case err := <-runResult:
			if err != nil {
				e.logger.WithError(err).Debug("run: execution scheduler returned an error")
				e.setRunStatus(lib.RunStatusAbortedSystem)
			} else {
				e.logger.Debug("run: execution scheduler terminated")
				e.setRunStatus(lib.RunStatusFinished)
			}
		case <-runCtx.Done():
			e.logger.Debug("run: context expired; exiting...")
			e.setRunStatus(lib.RunStatusAbortedUser)
		case <-e.stopChan:
			runSubCancel()
			e.logger.Debug("run: stopped by user; exiting...")
			e.setRunStatus(lib.RunStatusAbortedUser)
		case <-thresholdAbortChan:
			e.logger.Debug("run: stopped by thresholds; exiting...")
			runSubCancel()
			e.setRunStatus(lib.RunStatusAbortedThreshold)
		}
	}()

	// Run thresholds, if not disabled.
	if !e.NoThresholds {
		processes.Add(1)
		go func() {
			defer processes.Done()
			defer e.logger.Debug("Engine: Thresholds terminated")
			ticker := time.NewTicker(ThresholdsRate)
			defer ticker.Stop()

			for {
				select {
				case <-ticker.C:
					if e.processThresholds() {
						close(thresholdAbortChan)
						return
					}
				case <-runCtx.Done():
					return
				}
			}
		}()
	}

	return processes.Wait
}

func (e *Engine) processMetrics(globalCtx context.Context) {
	sampleContainers := []stats.SampleContainer{}

	defer func() {
		// Process any remaining metrics in the pipeline, by this point Run()
		// has already finished and nothing else should be producing metrics.
		e.logger.Debug("Metrics processing winding down...")

		close(e.Samples)
		for sc := range e.Samples {
			sampleContainers = append(sampleContainers, sc)
		}
		e.processSamples(sampleContainers)

		if !e.NoThresholds {
			e.processThresholds() // Process the thresholds one final time
		}
	}()

	ticker := time.NewTicker(CollectRate)
	defer ticker.Stop()

	e.logger.Debug("Metrics processing started...")
	for {
		select {
		case <-ticker.C:
			if len(sampleContainers) > 0 {
				e.processSamples(sampleContainers)
				sampleContainers = []stats.SampleContainer{} // TODO: optimize?
			}
		case sc := <-e.Samples:
			sampleContainers = append(sampleContainers, sc)
		case <-globalCtx.Done():
			return
		}
	}
}

func (e *Engine) setRunStatus(status lib.RunStatus) {
	for _, c := range e.Collectors {
		c.SetRunStatus(status)
	}
}

func (e *Engine) IsTainted() bool {
	return e.thresholdsTainted
}

// Stop closes a signal channel, forcing a running Engine to return
func (e *Engine) Stop() {
	close(e.stopChan)
}

// IsStopped returns a bool indicating whether the Engine has been stopped
func (e *Engine) IsStopped() bool {
	select {
	case <-e.stopChan:
		return true
	default:
		return false
	}
}

func (e *Engine) runMetricsEmission(ctx context.Context) {
	ticker := time.NewTicker(MetricsRate)
	for {
		select {
		case <-ticker.C:
			e.emitMetrics()
		case <-ctx.Done():
			return
		}
	}
}

func (e *Engine) emitMetrics() {
	t := time.Now()

	executionState := e.ExecutionScheduler.GetState()
	e.processSamples([]stats.SampleContainer{stats.ConnectedSamples{
		Samples: []stats.Sample{
			{
				Time:   t,
				Metric: metrics.VUs,
				Value:  float64(executionState.GetCurrentlyActiveVUsCount()),
				Tags:   e.Options.RunTags,
			}, {
				Time:   t,
				Metric: metrics.VUsMax,
				Value:  float64(executionState.GetInitializedVUsCount()),
				Tags:   e.Options.RunTags,
			},
		},
		Tags: e.Options.RunTags,
		Time: t,
	}})
}

func (e *Engine) processThresholds() (shouldAbort bool) {
	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	t := e.executionState.GetCurrentTestRunDuration()

	e.thresholdsTainted = false
	for _, m := range e.Metrics {
		if len(m.Thresholds.Thresholds) == 0 {
			continue
		}
		m.Tainted = null.BoolFrom(false)

		e.logger.WithField("m", m.Name).Debug("running thresholds")
		succ, err := m.Thresholds.Run(m.Sink, t)
		if err != nil {
			e.logger.WithField("m", m.Name).WithError(err).Error("Threshold error")
			continue
		}
		if !succ {
			e.logger.WithField("m", m.Name).Debug("Thresholds failed")
			m.Tainted = null.BoolFrom(true)
			e.thresholdsTainted = true
			if m.Thresholds.Abort {
				shouldAbort = true
			}
		}
	}

	return shouldAbort
}

func (e *Engine) processSamplesForMetrics(sampleCointainers []stats.SampleContainer) {
	for _, sampleCointainer := range sampleCointainers {
		samples := sampleCointainer.GetSamples()

		if len(samples) == 0 {
			continue
		}

		for _, sample := range samples {
			m, ok := e.Metrics[sample.Metric.Name]
			if !ok {
				m = stats.New(sample.Metric.Name, sample.Metric.Type, sample.Metric.Contains)
				m.Thresholds = e.thresholds[m.Name]
				m.Submetrics = e.submetrics[m.Name]
				e.Metrics[m.Name] = m
			}
			m.Sink.Add(sample)

			for _, sm := range m.Submetrics {
				if !sample.Tags.Contains(sm.Tags) {
					continue
				}

				if sm.Metric == nil {
					sm.Metric = stats.New(sm.Name, sample.Metric.Type, sample.Metric.Contains)
					sm.Metric.Sub = *sm
					sm.Metric.Thresholds = e.thresholds[sm.Name]
					e.Metrics[sm.Name] = sm.Metric
				}
				sm.Metric.Sink.Add(sample)
			}
		}
	}
}

func (e *Engine) processSamples(sampleContainers []stats.SampleContainer) {
	if len(sampleContainers) == 0 {
		return
	}

	// TODO: optimize this...
	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	// TODO: run this and the below code in goroutines?
	if !(e.NoSummary && e.NoThresholds && !e.SummaryExport) {
		e.processSamplesForMetrics(sampleContainers)
	}

	for _, collector := range e.Collectors {
		collector.Collect(sampleContainers)
	}
}
<commit_msg>Address code review comments<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package core

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/loadimpact/k6/stats"
	"github.com/sirupsen/logrus"
	"gopkg.in/guregu/null.v3"
)

const (
	metricsRate    = 1 * time.Second
	collectRate    = 50 * time.Millisecond
	thresholdsRate = 2 * time.Second
)

// The Engine is the beating heart of k6.
type Engine struct {
	// TODO: Make most of the stuff here private! And think how to refactor the
	// engine to be less stateful... it's currently one big mess of moving
	// pieces, and you implicitly first have to call Init() and then Run() -
	// maybe we should refactor it so we have a `Session` dauther-object that
	// Init() returns? The only problem with doing this is the REST API - it
	// expects to be able to get information from the Engine and is initialized
	// before the Init() call...

	ExecutionScheduler lib.ExecutionScheduler
	executionState     *lib.ExecutionState

	Options       lib.Options
	Collectors    []lib.Collector
	NoThresholds  bool
	NoSummary     bool
	SummaryExport bool

	logger   *logrus.Entry
	stopChan chan struct{}

	Metrics     map[string]*stats.Metric
	MetricsLock sync.Mutex

	Samples chan stats.SampleContainer

	// Assigned to metrics upon first received sample.
	thresholds map[string]stats.Thresholds
	submetrics map[string][]*stats.Submetric

	// Are thresholds tainted?
	thresholdsTainted bool
}

// NewEngine instantiates a new Engine, without doing any heavy initialization.
func NewEngine(ex lib.ExecutionScheduler, o lib.Options, logger *logrus.Logger) (*Engine, error) {
	if ex == nil {
		return nil, errors.New("missing ExecutionScheduler instance")
	}

	e := &Engine{
		ExecutionScheduler: ex,
		executionState:     ex.GetState(),

		Options:  o,
		Metrics:  make(map[string]*stats.Metric),
		Samples:  make(chan stats.SampleContainer, o.MetricSamplesBufferSize.Int64),
		stopChan: make(chan struct{}),
		logger:   logger.WithField("component", "engine"),
	}

	e.thresholds = o.Thresholds
	e.submetrics = make(map[string][]*stats.Submetric)
	for name := range e.thresholds {
		if !strings.Contains(name, "{") {
			continue
		}

		parent, sm := stats.NewSubmetric(name)
		e.submetrics[parent] = append(e.submetrics[parent], sm)
	}

	return e, nil
}

// Init is used to initialize the execution scheduler and all metrics processing
// in the engine. The first is a costly operation, since it initializes all of
// the planned VUs and could potentially take a long time. It either returns an
// error immediately, or it returns test run() and wait() functions.
//
// Things to note:
//  - The first lambda, Run(), synchronously executes the actual load test.
//  - It can be prematurely aborted by cancelling the runCtx - this won't stop
//    the metrics collection by the Engine.
//  - Stopping the metrics collection can be done at any time after Run() has
//    returned by cancelling the globalCtx
//  - The second returned lambda can be used to wait for that process to finish.
func (e *Engine) Init(globalCtx, runCtx context.Context) (run func() error, wait func(), err error) {
	e.logger.Debug("Initialization starting...")
	// TODO: if we ever need metrics processing in the init context, we can move
	// this below the other components... or even start them concurrently?
	if err := e.ExecutionScheduler.Init(runCtx, e.Samples); err != nil {
		return nil, nil, err
	}

	// TODO: move all of this in a separate struct? see main TODO above

	runSubCtx, runSubCancel := context.WithCancel(runCtx)

	resultCh := make(chan error)
	runFn := func() error {
		e.logger.Debug("Execution scheduler starting...")
		err := e.ExecutionScheduler.Run(globalCtx, runSubCtx, e.Samples)
		e.logger.WithError(err).Debug("Execution scheduler terminated")

		select {
		case <-runSubCtx.Done():
			// do nothing, the test run was aborted somehow
		default:
			resultCh <- err // we finished normally, so send the result
		}

		return err
	}
	waitFn := e.startBackgroundProcesses(globalCtx, runCtx, resultCh, runSubCancel)
	return runFn, waitFn, nil
}

// This starts a bunch of goroutines to process metrics, thresholds, and set the
// test run status when it ends. It returns a function that can be used after
// the provided context is called, to wait for the complete winding down of all
// started goroutines.
func (e *Engine) startBackgroundProcesses( //nolint:funlen
	globalCtx, runCtx context.Context, runResult <-chan error, runSubCancel func(),
) (wait func()) {
	processes := new(sync.WaitGroup)

	// Spin up all configured collectors
	for _, collector := range e.Collectors {
		processes.Add(1)
		go func(collector lib.Collector) {
			collector.Run(globalCtx)
			processes.Done()
		}(collector)
	}

	// Siphon and handle all produced metric samples
	processes.Add(1)
	go func() {
		defer processes.Done()
		e.processMetrics(globalCtx)
	}()

	// Run VU metrics emission, only while the test is running.
	// TODO: move? this seems like something the ExecutionScheduler should emit...
	processes.Add(1)
	go func() {
		defer processes.Done()
		e.logger.Debug("Starting emission of VU metrics...")
		e.runMetricsEmission(runCtx)
		e.logger.Debug("Metrics emission terminated")
	}()

	// Update the test run status when the test finishes
	processes.Add(1)
	thresholdAbortChan := make(chan struct{})
	go func() {
		defer processes.Done()
		select {
		case err := <-runResult:
			if err != nil {
				e.logger.WithError(err).Debug("run: execution scheduler returned an error")
				e.setRunStatus(lib.RunStatusAbortedSystem)
			} else {
				e.logger.Debug("run: execution scheduler terminated")
				e.setRunStatus(lib.RunStatusFinished)
			}
		case <-runCtx.Done():
			e.logger.Debug("run: context expired; exiting...")
			e.setRunStatus(lib.RunStatusAbortedUser)
		case <-e.stopChan:
			runSubCancel()
			e.logger.Debug("run: stopped by user; exiting...")
			e.setRunStatus(lib.RunStatusAbortedUser)
		case <-thresholdAbortChan:
			e.logger.Debug("run: stopped by thresholds; exiting...")
			runSubCancel()
			e.setRunStatus(lib.RunStatusAbortedThreshold)
		}
	}()

	// Run thresholds, if not disabled.
	if !e.NoThresholds {
		processes.Add(1)
		go func() {
			defer processes.Done()
			defer e.logger.Debug("Engine: Thresholds terminated")
			ticker := time.NewTicker(thresholdsRate)
			defer ticker.Stop()

			for {
				select {
				case <-ticker.C:
					if e.processThresholds() {
						close(thresholdAbortChan)
						return
					}
				case <-runCtx.Done():
					return
				}
			}
		}()
	}

	return processes.Wait
}

func (e *Engine) processMetrics(globalCtx context.Context) {
	sampleContainers := []stats.SampleContainer{}

	defer func() {
		// Process any remaining metrics in the pipeline, by this point Run()
		// has already finished and nothing else should be producing metrics.
		e.logger.Debug("Metrics processing winding down...")

		close(e.Samples)
		for sc := range e.Samples {
			sampleContainers = append(sampleContainers, sc)
		}
		e.processSamples(sampleContainers)

		if !e.NoThresholds {
			e.processThresholds() // Process the thresholds one final time
		}
	}()

	ticker := time.NewTicker(collectRate)
	defer ticker.Stop()

	e.logger.Debug("Metrics processing started...")
	for {
		select {
		case <-ticker.C:
			if len(sampleContainers) > 0 {
				e.processSamples(sampleContainers)
				sampleContainers = []stats.SampleContainer{} // TODO: optimize?
			}
		case sc := <-e.Samples:
			sampleContainers = append(sampleContainers, sc)
		case <-globalCtx.Done():
			return
		}
	}
}

func (e *Engine) setRunStatus(status lib.RunStatus) {
	for _, c := range e.Collectors {
		c.SetRunStatus(status)
	}
}

func (e *Engine) IsTainted() bool {
	return e.thresholdsTainted
}

// Stop closes a signal channel, forcing a running Engine to return
func (e *Engine) Stop() {
	close(e.stopChan)
}

// IsStopped returns a bool indicating whether the Engine has been stopped
func (e *Engine) IsStopped() bool {
	select {
	case <-e.stopChan:
		return true
	default:
		return false
	}
}

func (e *Engine) runMetricsEmission(ctx context.Context) {
	ticker := time.NewTicker(metricsRate)
	for {
		select {
		case <-ticker.C:
			e.emitMetrics()
		case <-ctx.Done():
			return
		}
	}
}

func (e *Engine) emitMetrics() {
	t := time.Now()

	executionState := e.ExecutionScheduler.GetState()
	e.Samples <- stats.ConnectedSamples{
		Samples: []stats.Sample{
			{
				Time:   t,
				Metric: metrics.VUs,
				Value:  float64(executionState.GetCurrentlyActiveVUsCount()),
				Tags:   e.Options.RunTags,
			}, {
				Time:   t,
				Metric: metrics.VUsMax,
				Value:  float64(executionState.GetInitializedVUsCount()),
				Tags:   e.Options.RunTags,
			},
		},
		Tags: e.Options.RunTags,
		Time: t,
	}
}

func (e *Engine) processThresholds() (shouldAbort bool) {
	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	t := e.executionState.GetCurrentTestRunDuration()

	e.thresholdsTainted = false
	for _, m := range e.Metrics {
		if len(m.Thresholds.Thresholds) == 0 {
			continue
		}
		m.Tainted = null.BoolFrom(false)

		e.logger.WithField("m", m.Name).Debug("running thresholds")
		succ, err := m.Thresholds.Run(m.Sink, t)
		if err != nil {
			e.logger.WithField("m", m.Name).WithError(err).Error("Threshold error")
			continue
		}
		if !succ {
			e.logger.WithField("m", m.Name).Debug("Thresholds failed")
			m.Tainted = null.BoolFrom(true)
			e.thresholdsTainted = true
			if m.Thresholds.Abort {
				shouldAbort = true
			}
		}
	}

	return shouldAbort
}

func (e *Engine) processSamplesForMetrics(sampleCointainers []stats.SampleContainer) {
	for _, sampleCointainer := range sampleCointainers {
		samples := sampleCointainer.GetSamples()

		if len(samples) == 0 {
			continue
		}

		for _, sample := range samples {
			m, ok := e.Metrics[sample.Metric.Name]
			if !ok {
				m = stats.New(sample.Metric.Name, sample.Metric.Type, sample.Metric.Contains)
				m.Thresholds = e.thresholds[m.Name]
				m.Submetrics = e.submetrics[m.Name]
				e.Metrics[m.Name] = m
			}
			m.Sink.Add(sample)

			for _, sm := range m.Submetrics {
				if !sample.Tags.Contains(sm.Tags) {
					continue
				}

				if sm.Metric == nil {
					sm.Metric = stats.New(sm.Name, sample.Metric.Type, sample.Metric.Contains)
					sm.Metric.Sub = *sm
					sm.Metric.Thresholds = e.thresholds[sm.Name]
					e.Metrics[sm.Name] = sm.Metric
				}
				sm.Metric.Sink.Add(sample)
			}
		}
	}
}

func (e *Engine) processSamples(sampleContainers []stats.SampleContainer) {
	if len(sampleContainers) == 0 {
		return
	}

	// TODO: optimize this...
	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	// TODO: run this and the below code in goroutines?
	if !(e.NoSummary && e.NoThresholds && !e.SummaryExport) {
		e.processSamplesForMetrics(sampleContainers)
	}

	for _, collector := range e.Collectors {
		collector.Collect(sampleContainers)
	}
}
<|endoftext|>