<commit_before>package cache

import (
	"bytes"
	"compress/zlib"
	"gondola/log"
	"io"
	"net/url"
	"strconv"
)

type BackendInitializer func(*url.URL) Backend

var (
	defaultCacheUrl string
	backends        = map[string]BackendInitializer{}
	codecs          = map[string]*Codec{}
)

type Backend interface {
	Set(key string, b []byte, timeout int) error
	Get(key string) ([]byte, error)
	GetMulti(keys []string) (map[string][]byte, error)
	Delete(key string) error
	Close() error
}

type Codec struct {
	Encode func(v interface{}) ([]byte, error)
	Decode func(data []byte, v interface{}) error
}

type Cache struct {
	Prefix            string
	Backend           Backend
	Codec             *Codec
	MinCompressLength int
	CompressLevel     int
}

func (c *Cache) manipulatesKeys() bool {
	return c.Prefix != ""
}

func (c *Cache) backendKey(key string) string {
	return c.Prefix + key
}

func (c *Cache) frontendKey(key string) string {
	if c.Prefix != "" {
		return key[len(c.Prefix):]
	}
	return key
}

// Set stores the given object in the cache associated with the
// given key. Timeout is the number of seconds until the item
// expires. If the timeout is 0, the item is never expired and
// might be only purged from cache when running out of space
func (c *Cache) Set(key string, object interface{}, timeout int) {
	b, err := c.Codec.Encode(&object)
	if err != nil {
		log.Errorf("Error encoding object for key %s: %s", key, err)
		return
	}
	c.SetBytes(key, b, timeout)
}

// Get returns the item assocciated with the given key, wrapped
// in an interface{}
func (c *Cache) Get(key string) interface{} {
	var obj interface{}
	if c.GetObject(key, &obj) {
		return obj
	}
	return nil
}

// GetObject returns the item associated with the given key
// using the passed in interface{} (which should be a pointer to
// the same struct type that was stored with Set). Returns true
// if the object could be succesfully decoded.
func (c *Cache) GetObject(key string, obj interface{}) bool {
	b := c.GetBytes(key)
	if b != nil {
		err := c.Codec.Decode(b, obj)
		if err == nil {
			return true
		}
		log.Errorf("Error decoding object for key %s: %s", key, err)
	}
	return false
}

// GetMulti returns several objects as a map[string]interface{}
// in only one roundtrip to the cache
func (c *Cache) GetMulti(keys []string) map[string]interface{} {
	if c.manipulatesKeys() {
		k := make([]string, len(keys))
		for ii, v := range keys {
			k[ii] = c.backendKey(v)
		}
		keys = k
	}
	data, err := c.Backend.GetMulti(keys)
	if err != nil {
		log.Errorf("Error querying cache for keys %v: %v", keys, err)
		return nil
	}
	objects := make(map[string]interface{}, len(data))
	if c.manipulatesKeys() {
		for k, v := range data {
			var object interface{}
			err := c.Codec.Decode(v, &object)
			if err != nil {
				log.Errorf("Error decoding object for key %s: %s", k, err)
				continue
			}
			objects[c.frontendKey(k)] = object
		}
	} else {
		for k, v := range data {
			var object interface{}
			err := c.Codec.Decode(v, &object)
			if err != nil {
				log.Errorf("Error decoding object for key %s: %s", k, err)
				continue
			}
			objects[k] = object
		}
	}
	return objects
}

// SetBytes stores the given byte array assocciated with
// the given key. See the documentation for Set for an
// explanation of the timeout parameter
func (c *Cache) SetBytes(key string, b []byte, timeout int) {
	if c.MinCompressLength >= 0 {
		if l := len(b); l > c.MinCompressLength {
			var buf bytes.Buffer
			w, err := zlib.NewWriterLevel(&buf, c.CompressLevel)
			if err == nil {
				w.Write(b)
				w.Close()
				cb := buf.Bytes()
				if cl := len(cb); cl < l {
					b = cb
					log.Debugf("Compressed key %s from %d bytes to %d bytes", key, l, cl)
				}
			} else {
				log.Warningf("Error opening zlib writer: %s", err)
			}
		}
	}
	log.Debugf("Setting key %s (%d bytes)", c.backendKey(key), len(b))
	err := c.Backend.Set(c.backendKey(key), b, timeout)
	if err != nil {
		log.Errorf("Error setting cache key %s: %s", key, err)
	}
}

// GetBytes returns the byte array assocciated with the given key
func (c *Cache) GetBytes(key string) []byte {
	b, err := c.Backend.Get(c.backendKey(key))
	if err != nil {
		log.Errorf("Error getting cache key %s: %s", key, err)
		return nil
	}
	if c.MinCompressLength >= 0 && len(b) > 0 {
		log.Debugf("Decompressing key %s", key)
		r, err := zlib.NewReader(bytes.NewBuffer(b))
		if err == nil {
			var buf bytes.Buffer
			_, err = io.Copy(&buf, r)
			if err == nil {
				if r.Close() == nil {
					b = buf.Bytes()
				}
			} else {
				r.Close()
			}
		}
	}
	return b
}

func (c *Cache) Delete(key string) {
	err := c.Backend.Delete(c.backendKey(key))
	if err != nil {
		log.Errorf("Error deleting cache key %s: %s", key, err)
	}
}

func (c *Cache) Close() {
	c.Backend.Close()
}

func RegisterBackend(scheme string, f BackendInitializer) {
	backends[scheme] = f
}

func RegisterCodec(name string, codec *Codec) {
	codecs[name] = codec
}

func SetDefaultUrl(url string) {
	defaultCacheUrl = url
}

func DefaultUrl() string {
	return defaultCacheUrl
}

func New(cacheUrl string) *Cache {
	cache := &Cache{
		MinCompressLength: -1,
		CompressLevel:     zlib.DefaultCompression,
	}
	var query url.Values
	u, err := url.Parse(cacheUrl)
	if err != nil {
		if err != nil && cacheUrl != "" {
			log.Errorf("Invalid cache URL %q: %s\n", cacheUrl, err)
		}
	} else {
		query = u.Query()
	}
	var codec *Codec = nil
	if query != nil {
		codecName := query.Get("codec")
		if codecName != "" {
			if c, ok := codecs[codecName]; ok {
				codec = c
			} else {
				log.Errorf("Unknown cache codec name %q\n", codecName)
			}
		}
		cache.Prefix = query.Get("prefix")
		if mcl := query.Get("min_compress"); mcl != "" {
			val, err := strconv.Atoi(mcl)
			if err == nil {
				cache.MinCompressLength = val
			} else {
				log.Warningf("Invalid min_compress value %q (must be integer)", mcl)
			}
		}
		if cl := query.Get("compress_level"); cl != "" {
			val, err := strconv.Atoi(cl)
			if err == nil {
				if val >= zlib.NoCompression && val <= zlib.BestCompression {
					cache.CompressLevel = val
				} else {
					log.Warningf("Invalid compress_level value %d (must be between %d and %d)",
						val, zlib.NoCompression, zlib.BestCompression)
				}
			} else {
				log.Warningf("Invalid compress_level value %q (must be integer)", cl)
			}
		}
	}
	if codec == nil {
		codec = &GobEncoder
	}
	cache.Codec = codec
	var backendInitializer BackendInitializer
	if u != nil {
		backendInitializer = backends[u.Scheme]
		if backendInitializer == nil && cacheUrl != "" {
			log.Errorf("Unknown cache backend type %q\n", u.Scheme)
		}
	}
	if backendInitializer == nil {
		backendInitializer = InitializeDummyBackend
	}
	cache.Backend = backendInitializer(u)
	return cache
}

func NewDefault() *Cache {
	return New(defaultCacheUrl)
}
<commit_msg>Move debug message to indicate when a key has been decompressed<commit_after>package cache

import (
	"bytes"
	"compress/zlib"
	"gondola/log"
	"io"
	"net/url"
	"strconv"
)

type BackendInitializer func(*url.URL) Backend

var (
	defaultCacheUrl string
	backends        = map[string]BackendInitializer{}
	codecs          = map[string]*Codec{}
)

type Backend interface {
	Set(key string, b []byte, timeout int) error
	Get(key string) ([]byte, error)
	GetMulti(keys []string) (map[string][]byte, error)
	Delete(key string) error
	Close() error
}

type Codec struct {
	Encode func(v interface{}) ([]byte, error)
	Decode func(data []byte, v interface{}) error
}

type Cache struct {
	Prefix            string
	Backend           Backend
	Codec             *Codec
	MinCompressLength int
	CompressLevel     int
}

func (c *Cache) manipulatesKeys() bool {
	return c.Prefix != ""
}

func (c *Cache) backendKey(key string) string {
	return c.Prefix + key
}

func (c *Cache) frontendKey(key string) string {
	if c.Prefix != "" {
		return key[len(c.Prefix):]
	}
	return key
}

// Set stores the given object in the cache associated with the
// given key. Timeout is the number of seconds until the item
// expires. If the timeout is 0, the item is never expired and
// might be only purged from cache when running out of space
func (c *Cache) Set(key string, object interface{}, timeout int) {
	b, err := c.Codec.Encode(&object)
	if err != nil {
		log.Errorf("Error encoding object for key %s: %s", key, err)
		return
	}
	c.SetBytes(key, b, timeout)
}

// Get returns the item assocciated with the given key, wrapped
// in an interface{}
func (c *Cache) Get(key string) interface{} {
	var obj interface{}
	if c.GetObject(key, &obj) {
		return obj
	}
	return nil
}

// GetObject returns the item associated with the given key
// using the passed in interface{} (which should be a pointer to
// the same struct type that was stored with Set). Returns true
// if the object could be succesfully decoded.
func (c *Cache) GetObject(key string, obj interface{}) bool {
	b := c.GetBytes(key)
	if b != nil {
		err := c.Codec.Decode(b, obj)
		if err == nil {
			return true
		}
		log.Errorf("Error decoding object for key %s: %s", key, err)
	}
	return false
}

// GetMulti returns several objects as a map[string]interface{}
// in only one roundtrip to the cache
func (c *Cache) GetMulti(keys []string) map[string]interface{} {
	if c.manipulatesKeys() {
		k := make([]string, len(keys))
		for ii, v := range keys {
			k[ii] = c.backendKey(v)
		}
		keys = k
	}
	data, err := c.Backend.GetMulti(keys)
	if err != nil {
		log.Errorf("Error querying cache for keys %v: %v", keys, err)
		return nil
	}
	objects := make(map[string]interface{}, len(data))
	if c.manipulatesKeys() {
		for k, v := range data {
			var object interface{}
			err := c.Codec.Decode(v, &object)
			if err != nil {
				log.Errorf("Error decoding object for key %s: %s", k, err)
				continue
			}
			objects[c.frontendKey(k)] = object
		}
	} else {
		for k, v := range data {
			var object interface{}
			err := c.Codec.Decode(v, &object)
			if err != nil {
				log.Errorf("Error decoding object for key %s: %s", k, err)
				continue
			}
			objects[k] = object
		}
	}
	return objects
}

// SetBytes stores the given byte array assocciated with
// the given key. See the documentation for Set for an
// explanation of the timeout parameter
func (c *Cache) SetBytes(key string, b []byte, timeout int) {
	if c.MinCompressLength >= 0 {
		if l := len(b); l > c.MinCompressLength {
			var buf bytes.Buffer
			w, err := zlib.NewWriterLevel(&buf, c.CompressLevel)
			if err == nil {
				w.Write(b)
				w.Close()
				cb := buf.Bytes()
				if cl := len(cb); cl < l {
					b = cb
					log.Debugf("Compressed key %s from %d bytes to %d bytes", key, l, cl)
				}
			} else {
				log.Warningf("Error opening zlib writer: %s", err)
			}
		}
	}
	log.Debugf("Setting key %s (%d bytes)", c.backendKey(key), len(b))
	err := c.Backend.Set(c.backendKey(key), b, timeout)
	if err != nil {
		log.Errorf("Error setting cache key %s: %s", key, err)
	}
}

// GetBytes returns the byte array assocciated with the given key
func (c *Cache) GetBytes(key string) []byte {
	b, err := c.Backend.Get(c.backendKey(key))
	if err != nil {
		log.Errorf("Error getting cache key %s: %s", key, err)
		return nil
	}
	if c.MinCompressLength >= 0 && len(b) > 0 {
		r, err := zlib.NewReader(bytes.NewBuffer(b))
		if err == nil {
			var buf bytes.Buffer
			_, err = io.Copy(&buf, r)
			if err == nil {
				if r.Close() == nil {
					log.Debugf("Decompressed key %s", key)
					b = buf.Bytes()
				}
			} else {
				r.Close()
			}
		}
	}
	return b
}

func (c *Cache) Delete(key string) {
	err := c.Backend.Delete(c.backendKey(key))
	if err != nil {
		log.Errorf("Error deleting cache key %s: %s", key, err)
	}
}

func (c *Cache) Close() {
	c.Backend.Close()
}

func RegisterBackend(scheme string, f BackendInitializer) {
	backends[scheme] = f
}

func RegisterCodec(name string, codec *Codec) {
	codecs[name] = codec
}

func SetDefaultUrl(url string) {
	defaultCacheUrl = url
}

func DefaultUrl() string {
	return defaultCacheUrl
}

func New(cacheUrl string) *Cache {
	cache := &Cache{
		MinCompressLength: -1,
		CompressLevel:     zlib.DefaultCompression,
	}
	var query url.Values
	u, err := url.Parse(cacheUrl)
	if err != nil {
		if err != nil && cacheUrl != "" {
			log.Errorf("Invalid cache URL %q: %s\n", cacheUrl, err)
		}
	} else {
		query = u.Query()
	}
	var codec *Codec = nil
	if query != nil {
		codecName := query.Get("codec")
		if codecName != "" {
			if c, ok := codecs[codecName]; ok {
				codec = c
			} else {
				log.Errorf("Unknown cache codec name %q\n", codecName)
			}
		}
		cache.Prefix = query.Get("prefix")
		if mcl := query.Get("min_compress"); mcl != "" {
			val, err := strconv.Atoi(mcl)
			if err == nil {
				cache.MinCompressLength = val
			} else {
				log.Warningf("Invalid min_compress value %q (must be integer)", mcl)
			}
		}
		if cl := query.Get("compress_level"); cl != "" {
			val, err := strconv.Atoi(cl)
			if err == nil {
				if val >= zlib.NoCompression && val <= zlib.BestCompression {
					cache.CompressLevel = val
				} else {
					log.Warningf("Invalid compress_level value %d (must be between %d and %d)",
						val, zlib.NoCompression, zlib.BestCompression)
				}
			} else {
				log.Warningf("Invalid compress_level value %q (must be integer)", cl)
			}
		}
	}
	if codec == nil {
		codec = &GobEncoder
	}
	cache.Codec = codec
	var backendInitializer BackendInitializer
	if u != nil {
		backendInitializer = backends[u.Scheme]
		if backendInitializer == nil && cacheUrl != "" {
			log.Errorf("Unknown cache backend type %q\n", u.Scheme)
		}
	}
	if backendInitializer == nil {
		backendInitializer = InitializeDummyBackend
	}
	cache.Backend = backendInitializer(u)
	return cache
}

func NewDefault() *Cache {
	return New(defaultCacheUrl)
}
<|endoftext|>


<commit_before>package testing

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"strings"
	"testing"

	"github.com/cockroachdb/examples-orms/go/gorm/model"
)

const (
	customersTable     = "customers"
	ordersTable        = "orders"
	productsTable      = "products"
	orderProductsTable = "order_products"
)

// These need to be variables so that their address can be taken.
var (
	customerName1 = "Billy"

	productName1       = "Ice Cream"
	productPrice1      = "123.40"
	productPrice1Float = 123.40
)

// parallelTestGroup maps a set of names to test functions, and will run each
// entry as a subtest in parallel by passing its T method to t.Run.
type parallelTestGroup map[string]func(t *testing.T)

func (ptg parallelTestGroup) T(t *testing.T) {
	for name, f := range ptg {
		t.Run(name, func(subT *testing.T) {
			subT.Parallel()
			f(subT)
		})
	}
}

// testDriver holds testing state and provides a suite of test methods that
// incrementally stress ORM functionality.
type testDriver struct {
	db     *sql.DB
	dbName string
	api    apiHandler
}

func (td testDriver) TestGeneratedTables(t *testing.T) {
	exp := []string{
		customersTable,
		orderProductsTable,
		ordersTable,
		productsTable,
	}

	actual := make(map[string]interface{}, len(exp))
	tables := td.query(t, `
SELECT table_name
FROM information_schema.tables
-- support both the legacy and the new information_schema structures. The former returned
-- the string 'def' as the table_catalog value for all rows. The latter returns 'public' as
-- the table_schema value for all user-created tables.
WHERE (table_catalog = 'def' AND table_schema = $1) OR (table_catalog = $1 AND table_schema = 'public')
ORDER BY 1`, td.dbName)
	for i := range tables {
		actual[tables[i]] = nil
	}

	for i := range exp {
		if _, ok := actual[exp[i]]; !ok {
			t.Fatalf("table %s is missing from generated schema", exp[i])
		}
	}
}

func (td testDriver) TestGeneratedCustomersTableColumns(t *testing.T) {
	exp := []string{"id", "name"}
	td.testGeneratedColumnsForTable(t, customersTable, exp)
}
func (td testDriver) TestGeneratedOrdersTableColumns(t *testing.T) {
	exp := []string{"customer_id", "id", "subtotal"}
	td.testGeneratedColumnsForTable(t, ordersTable, exp)
}
func (td testDriver) TestGeneratedProductsTableColumns(t *testing.T) {
	exp := []string{"id", "name", "price"}
	td.testGeneratedColumnsForTable(t, productsTable, exp)
}
func (td testDriver) TestGeneratedOrderProductsTableColumns(t *testing.T) {
	exp := []string{"order_id", "product_id"}
	td.testGeneratedColumnsForTable(t, orderProductsTable, exp)
}
func (td testDriver) testGeneratedColumnsForTable(t *testing.T, table string, columns []string) {
	td.queryAndAssert(t, columns, `
SELECT column_name
FROM information_schema.columns
-- see above about supporting both the legacy and the new information_schema structures.
WHERE ((table_catalog = 'def' AND table_schema = $1) OR (table_catalog = $1 AND table_schema = 'public'))
  AND table_name = $2
ORDER BY 1`, td.dbName, table)
}

func (td testDriver) TestCustomersEmpty(t *testing.T) {
	td.testTableEmpty(t, productsTable)
}
func (td testDriver) TestOrdersTableEmpty(t *testing.T) {
	td.testTableEmpty(t, customersTable)
}
func (td testDriver) TestProductsTableEmpty(t *testing.T) {
	td.testTableEmpty(t, ordersTable)
}
func (td testDriver) TestOrderProductsTableEmpty(t *testing.T) {
	td.testTableEmpty(t, orderProductsTable)
}
func (td testDriver) testTableEmpty(t *testing.T, table string) {
	td.queryAndAssert(t, []string{"0"}, fmt.Sprintf(`SELECT COUNT(*) FROM %s`, table))
}

func (td testDriver) TestRetrieveCustomersBeforeCreation(t *testing.T) {
	found, err := td.api.queryCustomers()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Customer{}
	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting customers from api before creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveProductsBeforeCreation(t *testing.T) {
	found, err := td.api.queryProducts()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Product{}
	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting products from api before creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveOrdersBeforeCreation(t *testing.T) {
	found, err := td.api.queryOrders()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Order{}
	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting orders from api before creation to be %v, found %v", expected, found)
	}
}

func (td testDriver) TestCreateCustomer(t *testing.T) {
	if err := td.api.createCustomer(customerName1); err != nil {
		t.Fatalf("error creating customer: %v", err)
	}
	td.queryAndAssert(t, []string{customerName1}, fmt.Sprintf(`SELECT name FROM %s`, customersTable))
}
func (td testDriver) TestCreateProduct(t *testing.T) {
	if err := td.api.createProduct(productName1, productPrice1Float); err != nil {
		t.Fatalf("error creating product: %v", err)
	}
	td.queryAndAssert(t, []string{row(productName1, productPrice1)}, fmt.Sprintf(`SELECT name, price FROM %s`, productsTable))
}

func (td testDriver) TestCreateOrder(t *testing.T) {
	// Get the single customer ID.
	customerIDs, err := td.queryIDs(t, customersTable)
	if err != nil {
		t.Fatal(err)
	}
	if len(customerIDs) != 1 {
		t.Fatalf("expected a single customer ID, found %v", customerIDs)
	}
	customerID := customerIDs[0]

	// Get the single product.
	productIDs, err := td.queryIDs(t, productsTable)
	if err != nil {
		t.Fatal(err)
	}
	if len(productIDs) != 1 {
		t.Fatalf("expected a single product ID, found %v", productIDs)
	}
	productID := productIDs[0]

	if err := td.api.createOrder(customerID, productID, productPrice1Float); err != nil {
		t.Fatalf("error creating order: %v", err)
	}
	td.queryAndAssert(t, []string{row(productPrice1)}, fmt.Sprintf(`SELECT subtotal FROM %s`, ordersTable))
}

func (td testDriver) TestRetrieveCustomerAfterCreation(t *testing.T) {
	found, err := td.api.queryCustomers()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Customer{
		{Name: &customerName1},
	}
	if !reflect.DeepEqual(expected, cleanCustomers(found)) {
		t.Fatalf("expecting customers from api after creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveProductAfterCreation(t *testing.T) {
	found, err := td.api.queryProducts()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Product{
		{Name: &productName1, Price: productPrice1Float},
	}
	if !reflect.DeepEqual(expected, cleanProducts(found)) {
		t.Fatalf("expecting products from api after creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveOrderAfterCreation(t *testing.T) {
	found, err := td.api.queryOrders()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Order{
		{Subtotal: productPrice1Float},
	}
	if !reflect.DeepEqual(expected, cleanOrders(found)) {
		t.Fatalf("expecting orders from api after creation to be %v, found %v", expected, found)
	}
}

func (td testDriver) queryIDs(t *testing.T, table string) ([]int, error) {
	rows, err := td.db.Query(fmt.Sprintf("SELECT id FROM %s", table))
	if err != nil {
		t.Fatal(err)
	}

	var ids []int
	for rows.Next() {
		var id int
		rows.Scan(&id)
		ids = append(ids, id)
	}
	if err := rows.Err(); err != nil {
		t.Fatal(err)
	}
	return ids, nil
}

func (td testDriver) query(t *testing.T, query string, args ...interface{}) []string {
	rows, err := td.db.Query(query, args...)
	if err != nil {
		t.Fatal(err)
	}

	found, err := rowsToStringSlice(rows)
	if err != nil {
		t.Fatal(err)
	}
	return found
}

func (td testDriver) queryAndAssert(t *testing.T, expected []string, query string, args ...interface{}) {
	found := td.query(t, query, args...)

	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting rows for query %q with args %+v to be %v, found %v", query, args, expected, found)
	}
}

func rowsToStringSlice(rows *sql.Rows) ([]string, error) {
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	vals := make([]interface{}, len(cols))
	strs := make([]string, len(cols))
	for i := range vals {
		vals[i] = &strs[i]
	}

	var s []string
	for rows.Next() {
		rows.Scan(vals...)
		s = append(s, strings.Join(strs, ", "))
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return s, nil
}

func row(vals ...interface{}) string {
	var b bytes.Buffer
	for i, val := range vals {
		if i > 0 {
			b.WriteString(", ")
		}
		fmt.Fprintf(&b, "%v", val)
	}
	return b.String()
}
<commit_msg>Skip over rowid when testing schema<commit_after>package testing

import (
	"bytes"
	"database/sql"
	"fmt"
	"reflect"
	"strings"
	"testing"

	"github.com/cockroachdb/examples-orms/go/gorm/model"
)

const (
	customersTable     = "customers"
	ordersTable        = "orders"
	productsTable      = "products"
	orderProductsTable = "order_products"
)

// These need to be variables so that their address can be taken.
var (
	customerName1 = "Billy"

	productName1       = "Ice Cream"
	productPrice1      = "123.40"
	productPrice1Float = 123.40
)

// parallelTestGroup maps a set of names to test functions, and will run each
// entry as a subtest in parallel by passing its T method to t.Run.
type parallelTestGroup map[string]func(t *testing.T)

func (ptg parallelTestGroup) T(t *testing.T) {
	for name, f := range ptg {
		t.Run(name, func(subT *testing.T) {
			subT.Parallel()
			f(subT)
		})
	}
}

// testDriver holds testing state and provides a suite of test methods that
// incrementally stress ORM functionality.
type testDriver struct {
	db     *sql.DB
	dbName string
	api    apiHandler
}

func (td testDriver) TestGeneratedTables(t *testing.T) {
	exp := []string{
		customersTable,
		orderProductsTable,
		ordersTable,
		productsTable,
	}

	actual := make(map[string]interface{}, len(exp))
	tables := td.query(t, `
SELECT table_name
FROM information_schema.tables
-- support both the legacy and the new information_schema structures. The former returned
-- the string 'def' as the table_catalog value for all rows. The latter returns 'public' as
-- the table_schema value for all user-created tables.
WHERE (table_catalog = 'def' AND table_schema = $1) OR (table_catalog = $1 AND table_schema = 'public')
ORDER BY 1`, td.dbName)
	for i := range tables {
		actual[tables[i]] = nil
	}

	for i := range exp {
		if _, ok := actual[exp[i]]; !ok {
			t.Fatalf("table %s is missing from generated schema", exp[i])
		}
	}
}

func (td testDriver) TestGeneratedCustomersTableColumns(t *testing.T) {
	exp := []string{"id", "name"}
	td.testGeneratedColumnsForTable(t, customersTable, exp)
}
func (td testDriver) TestGeneratedOrdersTableColumns(t *testing.T) {
	exp := []string{"customer_id", "id", "subtotal"}
	td.testGeneratedColumnsForTable(t, ordersTable, exp)
}
func (td testDriver) TestGeneratedProductsTableColumns(t *testing.T) {
	exp := []string{"id", "name", "price"}
	td.testGeneratedColumnsForTable(t, productsTable, exp)
}
func (td testDriver) TestGeneratedOrderProductsTableColumns(t *testing.T) {
	exp := []string{"order_id", "product_id"}
	td.testGeneratedColumnsForTable(t, orderProductsTable, exp)
}
func (td testDriver) testGeneratedColumnsForTable(t *testing.T, table string, columns []string) {
	td.queryAndAssert(t, columns, `
SELECT column_name
FROM information_schema.columns
-- see above about supporting both the legacy and the new information_schema structures.
WHERE ((table_catalog = 'def' AND table_schema = $1) OR (table_catalog = $1 AND table_schema = 'public'))
  AND table_name = $2
  AND column_name != 'rowid'
ORDER BY 1`, td.dbName, table)
}

func (td testDriver) TestCustomersEmpty(t *testing.T) {
	td.testTableEmpty(t, productsTable)
}
func (td testDriver) TestOrdersTableEmpty(t *testing.T) {
	td.testTableEmpty(t, customersTable)
}
func (td testDriver) TestProductsTableEmpty(t *testing.T) {
	td.testTableEmpty(t, ordersTable)
}
func (td testDriver) TestOrderProductsTableEmpty(t *testing.T) {
	td.testTableEmpty(t, orderProductsTable)
}
func (td testDriver) testTableEmpty(t *testing.T, table string) {
	td.queryAndAssert(t, []string{"0"}, fmt.Sprintf(`SELECT COUNT(*) FROM %s`, table))
}

func (td testDriver) TestRetrieveCustomersBeforeCreation(t *testing.T) {
	found, err := td.api.queryCustomers()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Customer{}
	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting customers from api before creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveProductsBeforeCreation(t *testing.T) {
	found, err := td.api.queryProducts()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Product{}
	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting products from api before creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveOrdersBeforeCreation(t *testing.T) {
	found, err := td.api.queryOrders()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Order{}
	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting orders from api before creation to be %v, found %v", expected, found)
	}
}

func (td testDriver) TestCreateCustomer(t *testing.T) {
	if err := td.api.createCustomer(customerName1); err != nil {
		t.Fatalf("error creating customer: %v", err)
	}
	td.queryAndAssert(t, []string{customerName1}, fmt.Sprintf(`SELECT name FROM %s`, customersTable))
}
func (td testDriver) TestCreateProduct(t *testing.T) {
	if err := td.api.createProduct(productName1, productPrice1Float); err != nil {
		t.Fatalf("error creating product: %v", err)
	}
	td.queryAndAssert(t, []string{row(productName1, productPrice1)}, fmt.Sprintf(`SELECT name, price FROM %s`, productsTable))
}

func (td testDriver) TestCreateOrder(t *testing.T) {
	// Get the single customer ID.
	customerIDs, err := td.queryIDs(t, customersTable)
	if err != nil {
		t.Fatal(err)
	}
	if len(customerIDs) != 1 {
		t.Fatalf("expected a single customer ID, found %v", customerIDs)
	}
	customerID := customerIDs[0]

	// Get the single product.
	productIDs, err := td.queryIDs(t, productsTable)
	if err != nil {
		t.Fatal(err)
	}
	if len(productIDs) != 1 {
		t.Fatalf("expected a single product ID, found %v", productIDs)
	}
	productID := productIDs[0]

	if err := td.api.createOrder(customerID, productID, productPrice1Float); err != nil {
		t.Fatalf("error creating order: %v", err)
	}
	td.queryAndAssert(t, []string{row(productPrice1)}, fmt.Sprintf(`SELECT subtotal FROM %s`, ordersTable))
}

func (td testDriver) TestRetrieveCustomerAfterCreation(t *testing.T) {
	found, err := td.api.queryCustomers()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Customer{
		{Name: &customerName1},
	}
	if !reflect.DeepEqual(expected, cleanCustomers(found)) {
		t.Fatalf("expecting customers from api after creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveProductAfterCreation(t *testing.T) {
	found, err := td.api.queryProducts()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Product{
		{Name: &productName1, Price: productPrice1Float},
	}
	if !reflect.DeepEqual(expected, cleanProducts(found)) {
		t.Fatalf("expecting products from api after creation to be %v, found %v", expected, found)
	}
}
func (td testDriver) TestRetrieveOrderAfterCreation(t *testing.T) {
	found, err := td.api.queryOrders()
	if err != nil {
		t.Fatal(err)
	}

	expected := []model.Order{
		{Subtotal: productPrice1Float},
	}
	if !reflect.DeepEqual(expected, cleanOrders(found)) {
		t.Fatalf("expecting orders from api after creation to be %v, found %v", expected, found)
	}
}

func (td testDriver) queryIDs(t *testing.T, table string) ([]int, error) {
	rows, err := td.db.Query(fmt.Sprintf("SELECT id FROM %s", table))
	if err != nil {
		t.Fatal(err)
	}

	var ids []int
	for rows.Next() {
		var id int
		rows.Scan(&id)
		ids = append(ids, id)
	}
	if err := rows.Err(); err != nil {
		t.Fatal(err)
	}
	return ids, nil
}

func (td testDriver) query(t *testing.T, query string, args ...interface{}) []string {
	rows, err := td.db.Query(query, args...)
	if err != nil {
		t.Fatal(err)
	}

	found, err := rowsToStringSlice(rows)
	if err != nil {
		t.Fatal(err)
	}
	return found
}

func (td testDriver) queryAndAssert(t *testing.T, expected []string, query string, args ...interface{}) {
	found := td.query(t, query, args...)

	if !reflect.DeepEqual(expected, found) {
		t.Fatalf("expecting rows for query %q with args %+v to be %v, found %v", query, args, expected, found)
	}
}

func rowsToStringSlice(rows *sql.Rows) ([]string, error) {
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	vals := make([]interface{}, len(cols))
	strs := make([]string, len(cols))
	for i := range vals {
		vals[i] = &strs[i]
	}

	var s []string
	for rows.Next() {
		rows.Scan(vals...)
		s = append(s, strings.Join(strs, ", "))
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return s, nil
}

func row(vals ...interface{}) string {
	var b bytes.Buffer
	for i, val := range vals {
		if i > 0 {
			b.WriteString(", ")
		}
		fmt.Fprintf(&b, "%v", val)
	}
	return b.String()
}
<|endoftext|>


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

// +build ignore

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/pkg/errors"
)

var (
	old      = flag.Bool("old", false, "display old file data to stdout")
	new      = flag.Bool("new", false, "display new file data to stdout")
	manifest = flag.Bool("manifest", false, "display new plugin manifest to stdout")
	write    = flag.Bool("w", false, "write specs to file instead of stdout")
)

func main() {
	// Define usage
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [-old|-new|-manifest|-w] plugin_name\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	// Required first arg of plugin_name
	pluginName := flag.Arg(0)
	if pluginName == "" || flag.NFlag() > 1 {
		flag.Usage()
		os.Exit(2)
	}

	// Search gb binary path
	gbBin, err := exec.LookPath("gb")
	if err != nil {
		err = errors.Wrap(err, "does not exists gb binary")
		log.Fatal(err)
	}

	// Get the gb project directory root
	gbCmd := exec.Command(gbBin)
	gbCmd.Args = append(gbCmd.Args, "env", "GB_PROJECT_DIR")
	gbResult, err := gbCmd.Output()
	if err != nil {
		err = errors.Wrap(err, "cannot get gb project directory")
		log.Fatal(err)
	}
	prjDir := strings.TrimSpace(string(gbResult))

	// Get new plugin manifest
	manifestsCmd := exec.Command(filepath.Join(prjDir, "bin", pluginName), "-manifest", pluginName)
	newManifest, err := manifestsCmd.Output()
	if err != nil {
		panic(err)
	}
	newManifest = append(newManifest, '\n')

	if strings.Contains(pluginName, "-race") {
		pluginName = strings.TrimSuffix(pluginName, "-race")
	}
	// Get vim file information from the "./plugin" directory
	plugFile, err := os.OpenFile(filepath.Join(prjDir, "plugin", pluginName+".vim"), os.O_RDWR, os.ModeAppend)
	if err != nil {
		panic(err)
	}
	defer plugFile.Close()

	// Read plugin vim file
	oldData, err := ioutil.ReadAll(plugFile)
	if err != nil {
		panic(err)
	}

	re := regexp.MustCompile(`(?s)call remote#host#RegisterPlugin.+`)
	// Replace the old specs to the latest specs
	newData := re.ReplaceAll(oldData, newManifest)

	// Output result
	switch {
	case *old:
		fmt.Printf("%v", string(oldData))
		return
	case *new:
		fmt.Printf("%v", string(newData))
		return
	case *manifest:
		// Trim last newline for output to stdout
		fmt.Printf("%v", string(newManifest[:len(newManifest)-1]))
		return
	case *write:
		data := bytes.TrimSpace(newData)
		if bytes.Contains(data, []byte("-race")) {
			data = bytes.Replace(data, []byte("-race"), nil, -1)
		}
		data = append(data, byte('\n'))
		if _, err := plugFile.WriteAt(data, 0); err != nil {
			panic(err)
		}
		return
	}
}
<commit_msg>plugin/manifest: remove pkg/errors dependency<commit_after>// Copyright 2016 The nvim-go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ignore

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

var (
	old      = flag.Bool("old", false, "display old file data to stdout")
	new      = flag.Bool("new", false, "display new file data to stdout")
	manifest = flag.Bool("manifest", false, "display new plugin manifest to stdout")
	write    = flag.Bool("w", false, "write specs to file instead of stdout")
)

func main() {
	// Define usage
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [-old|-new|-manifest|-w] plugin_name\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	// Required first arg of plugin_name
	pluginName := flag.Arg(0)
	if pluginName == "" || flag.NFlag() > 1 {
		flag.Usage()
		os.Exit(2)
	}

	// Search gb binary path
	gbBin, err := exec.LookPath("gb")
	if err != nil {
		err = fmt.Errorf("does not exists gb binary: %v", err)
		log.Fatal(err)
	}

	// Get the gb project directory root
	gbCmd := exec.Command(gbBin)
	gbCmd.Args = append(gbCmd.Args, "env", "GB_PROJECT_DIR")
	gbResult, err := gbCmd.Output()
	if err != nil {
		err = fmt.Errorf("cannot get gb project directory: %v", err)
		log.Fatal(err)
	}
	prjDir := strings.TrimSpace(string(gbResult))

	// Get new plugin manifest
	manifestsCmd := exec.Command(filepath.Join(prjDir, "bin", pluginName), "-manifest", pluginName)
	newManifest, err := manifestsCmd.Output()
	if err != nil {
		panic(err)
	}
	newManifest = append(newManifest, '\n')

	if strings.Contains(pluginName, "-race") {
		pluginName = strings.TrimSuffix(pluginName, "-race")
	}
	// Get vim file information from the "./plugin" directory
	plugFile, err := os.OpenFile(filepath.Join(prjDir, "plugin", pluginName+".vim"), os.O_RDWR, os.ModeAppend)
	if err != nil {
		panic(err)
	}
	defer plugFile.Close()

	// Read plugin vim file
	oldData, err := ioutil.ReadAll(plugFile)
	if err != nil {
		panic(err)
	}

	re := regexp.MustCompile(`(?s)call remote#host#RegisterPlugin.+`)
	// Replace the old specs to the latest specs
	newData := re.ReplaceAll(oldData, newManifest)

	// Output result
	switch {
	case *old:
		fmt.Printf("%v", string(oldData))
		return
	case *new:
		fmt.Printf("%v", string(newData))
		return
	case *manifest:
		// Trim last newline for output to stdout
		fmt.Printf("%v", string(newManifest[:len(newManifest)-1]))
		return
	case *write:
		data := bytes.TrimSpace(newData)
		if bytes.Contains(data, []byte("-race")) {
			data = bytes.Replace(data, []byte("-race"), nil, -1)
		}
		data = append(data, byte('\n'))
		if _, err := plugFile.WriteAt(data, 0); err != nil {
			panic(err)
		}
		return
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Tigera, Inc. All rights reserved.

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

package client

import (
	goerrors "errors"
	"fmt"
	"math/big"
	"math/rand"
	"net"
	"reflect"
	"time"

	"github.com/golang/glog"
	"github.com/tigera/libcalico-go/lib/api"
	"github.com/tigera/libcalico-go/lib/backend/model"
	"github.com/tigera/libcalico-go/lib/errors"
	cnet "github.com/tigera/libcalico-go/lib/net"
)

type blockReaderWriter struct {
	client *Client
}

func (rw blockReaderWriter) getAffineBlocks(host string, ver ipVersion, pool *cnet.IPNet) ([]cnet.IPNet, error) {
	// Lookup all blocks by providing an empty BlockListOptions
	// to the List operation.
	opts := model.BlockAffinityListOptions{Host: host, IPVersion: ver.Number}
	datastoreObjs, err := rw.client.backend.List(opts)
	if err != nil {
		if _, ok := err.(errors.ErrorResourceDoesNotExist); ok {
			// The block path does not exist yet.  This is OK - it means
			// there are no affine blocks.
			return []cnet.IPNet{}, nil

		} else {
			glog.Errorf("Error getting affine blocks: %s", err)
			return nil, err
		}
	}

	// Iterate through and extract the block CIDRs.
	ids := []cnet.IPNet{}
	for _, o := range datastoreObjs {
		k := o.Key.(model.BlockAffinityKey)
		ids = append(ids, k.CIDR)
	}
	return ids, nil
}

func (rw blockReaderWriter) claimNewAffineBlock(
	host string, version ipVersion, pool *cnet.IPNet, config IPAMConfig) (*cnet.IPNet, error) {

	// If pool is not nil, use the given pool.  Otherwise, default to
	// all configured pools.
	var pools []cnet.IPNet
	if pool != nil {
		// Validate the given pool is actually configured and matches the version.
		if !rw.isConfiguredPool(*pool) {
			estr := fmt.Sprintf("The given pool (%s) does not exist", pool.String())
			return nil, goerrors.New(estr)
		} else if version.Number != pool.Version() {
			estr := fmt.Sprintf("The given pool (%s) does not match IP version %d", pool.String(), version.Number)
			return nil, goerrors.New(estr)
		}
		pools = []cnet.IPNet{*pool}
	} else {
		// Default to all configured pools.
		allPools, err := rw.client.Pools().List(api.PoolMetadata{})
		if err != nil {
			glog.Errorf("Error reading configured pools: %s", err)
			return nil, err
		}

		// Grab all the IP networks in these pools.
		for _, p := range allPools.Items {
			// Don't include disabled pools or pools that don't match
			// the requested IP version.
			if !p.Spec.Disabled && version.Number == p.Metadata.CIDR.Version() {
				pools = append(pools, p.Metadata.CIDR)
			}
		}
	}

	// If there are no pools, we cannot assign addresses.
	if len(pools) == 0 {
		return nil, goerrors.New("No configured Calico pools")
	}

	// Iterate through pools to find a new block.
	glog.V(2).Infof("Claiming a new affine block for host '%s'", host)
	for _, pool := range pools {
		// Use a block generator to iterate through all of the blocks
		// that fall within the pool.
		blocks := randomBlockGenerator(pool)
		for subnet := blocks(); subnet != nil; subnet = blocks() {
			// Check if a block already exists for this subnet.
			glog.V(4).Infof("Getting block: %s", subnet.String())
			key := model.BlockKey{CIDR: subnet}
			_, err := rw.client.backend.Get(key)
			if err != nil {
				if _, ok := err.(errors.ErrorResourceDoesNotExist); ok {
					// The block does not yet exist in etcd.  Try to grab it.
					glog.V(3).Infof("Found free block: %+v", *subnet)
					err = rw.claimBlockAffinity(*subnet, host, config)
					return subnet, err
				} else {
					glog.Errorf("Error getting block: %s", err)
					return nil, err
				}
			}
		}
	}
	return nil, noFreeBlocksError("No Free Blocks")
}

func (rw blockReaderWriter) claimBlockAffinity(subnet cnet.IPNet, host string, config IPAMConfig) error {
	// Claim the block affinity for this host.
	glog.V(2).Infof("Host %s claiming block affinity for %s", host, subnet)
	obj := model.KVPair{
		Key:   model.BlockAffinityKey{Host: host, CIDR: subnet},
		Value: model.BlockAffinity{},
	}
	_, err := rw.client.backend.Create(&obj)

	// Create the new block.
	block := newBlock(subnet)
	block.HostAffinity = &host
	block.StrictAffinity = config.StrictAffinity

	// Create the new block in the datastore.
	o := model.KVPair{
		Key:   model.BlockKey{block.CIDR},
		Value: block.AllocationBlock,
	}
	_, err = rw.client.backend.Create(&o)
	if err != nil {
		if _, ok := err.(errors.ErrorResourceAlreadyExists); ok {
			// Block already exists, check affinity.
			glog.Warningf("Problem claiming block affinity:", err)
			obj, err := rw.client.backend.Get(model.BlockKey{subnet})
			if err != nil {
				glog.Errorf("Error reading block:", err)
				return err
			}

			// Pull out the allocationBlock object.
			b := allocationBlock{obj.Value.(model.AllocationBlock)}

			if b.HostAffinity != nil && *b.HostAffinity == host {
				// Block has affinity to this host, meaning another
				// process on this host claimed it.
				glog.V(3).Infof("Block %s already claimed by us.  Success", subnet)
				return nil
			}

			// Some other host beat us to this block.  Cleanup and return error.
			err = rw.client.backend.Delete(&model.KVPair{
				Key: model.BlockAffinityKey{Host: host, CIDR: b.CIDR},
			})
			if err != nil {
				glog.Errorf("Error cleaning up block affinity: %s", err)
				return err
			}
			return affinityClaimedError{Block: b}
		} else {
			return err
		}
	}
	return nil
}

func (rw blockReaderWriter) releaseBlockAffinity(host string, blockCIDR cnet.IPNet) error {
	for i := 0; i < ipamEtcdRetries; i++ {
		// Read the model.KVPair containing the block
		// and pull out the allocationBlock object.  We need to hold on to this
		// so that we can pass it back to the datastore on Update.
		obj, err := rw.client.backend.Get(model.BlockKey{CIDR: blockCIDR})
		if err != nil {
			glog.Errorf("Error getting block %s: %s", blockCIDR.String(), err)
			return err
		}
		b := allocationBlock{obj.Value.(model.AllocationBlock)}

		// Check that the block affinity matches the given affinity.
		if b.HostAffinity != nil && *b.HostAffinity != host {
			glog.Errorf("Mismatched affinity: %s != %s", *b.HostAffinity, host)
			return affinityClaimedError{Block: b}
		}

		if b.empty() {
			// If the block is empty, we can delete it.
			err := rw.client.backend.Delete(&model.KVPair{
				Key: model.BlockKey{CIDR: b.CIDR},
			})
			if err != nil {
				// Return the error unless the block didn't exist.
				if _, ok := err.(errors.ErrorResourceDoesNotExist); !ok {
					glog.Errorf("Error deleting block: %s", err)
					return err
				}
			}
		} else {
			// Otherwise, we need to remove affinity from it.
			// This prevents the host from automatically assigning
			// from this block unless we're allowed to overflow into
			// non-affine blocks.
			b.HostAffinity = nil

			// Pass back the original KVPair with the new
			// block information so we can do a CAS.
			obj.Value = b
			_, err = rw.client.backend.Update(obj)
			if err != nil {
				if _, ok := err.(errors.ErrorResourceUpdateConflict); ok {
					// CASError - continue.
					continue
				} else {
					return err
				}
			}
		}

		// We've removed / updated the block, so update the host config
		// to remove the CIDR.
		err = rw.client.backend.Delete(&model.KVPair{
			Key: model.BlockAffinityKey{Host: host, CIDR: b.CIDR},
		})
		if err != nil {
			// Return the error unless the affinity didn't exist.
			if _, ok := err.(errors.ErrorResourceDoesNotExist); !ok {
				glog.Errorf("Error deleting block affinity: %s", err)
				return err
			}
		}
		return nil

	}
	return goerrors.New("Max retries hit")
}

// withinConfiguredPools returns true if the given IP is within a configured
// Calico pool, and false otherwise.
func (rw blockReaderWriter) withinConfiguredPools(ip cnet.IP) bool {
	allPools, _ := rw.client.Pools().List(api.PoolMetadata{})
	for _, p := range allPools.Items {
		// Compare any enabled pools.
		if !p.Spec.Disabled && p.Metadata.CIDR.Contains(ip.IP) {
			return true
		}
	}
	return false
}

// isConfiguredPool returns true if the given IPNet is a configured
// Calico pool, and false otherwise.
func (rw blockReaderWriter) isConfiguredPool(cidr cnet.IPNet) bool {
	allPools, _ := rw.client.Pools().List(api.PoolMetadata{})
	for _, p := range allPools.Items {
		// Compare any enabled pools.
		if !p.Spec.Disabled && reflect.DeepEqual(p.Metadata.CIDR, cidr) {
			return true
		}
	}
	return false
}

// Generator to get list of block CIDRs which
// fall within the given pool. Returns nil when no more
// blocks can be generated.
func blockGenerator(pool cnet.IPNet) func() *cnet.IPNet {
	// Determine the IP type to use.
	version := getIPVersion(cnet.IP{pool.IP})
	ip := cnet.IP{pool.IP}
	return func() *cnet.IPNet {
		returnIP := ip
		ip = incrementIP(ip, big.NewInt(blockSize))
		if pool.Contains(ip.IP) {
			ipnet := net.IPNet{returnIP.IP, version.BlockPrefixMask}
			cidr := cnet.IPNet{ipnet}
			return &cidr
		} else {
			return nil
		}
	}
}

// Returns a generator that, when called, returns a random
// block from the given pool.  When there are no blocks left,
// the it returns nil.
func randomBlockGenerator(pool cnet.IPNet) func() *cnet.IPNet {

	// Determine the IP type to use.
	version := getIPVersion(cnet.IP{pool.IP})
	baseIP := cnet.IP{pool.IP}

	// Determine the number of blocks within this pool.
	ones, size := pool.Mask.Size()
	prefixLen := size - ones
	numIP := new(big.Int).Exp(big.NewInt(2), big.NewInt(int64(prefixLen)), nil)
	numBlocks := new(big.Int)
	numBlocks.Div(numIP, big.NewInt(blockSize))

	// Start at a random offset index
	source := rand.NewSource(time.Now().UnixNano())
	randm := rand.New(source)

	// initialIndex keeps track of the random starting point
	initialIndex := new(big.Int)
	initialIndex.Rand(randm, numBlocks)

	// i keeps track of current index while walking the blocks in a pool
	i := initialIndex

	// numReturned keeps track of number of blocks returned
	numReturned := big.NewInt(0)

	// numDiff = numBlocks - i
	numDiff := new(big.Int)

	return func() *cnet.IPNet {
		ip := incrementIP(baseIP, i.Mul(i, big.NewInt(blockSize)))
		ipnet := net.IPNet{ip.IP, version.BlockPrefixMask}

		numDiff.Sub(numBlocks, i)

		if numDiff.Cmp(big.NewInt(1)) <= 0 {
			// Index has reached end of the blocks;
			// Loop back to beginning of pool rather than
			// increment, because incrementing would put us outside of the pool.
			i = big.NewInt(0)
		} else {
			// Increment to the next block
			i.Add(i, big.NewInt(1))
		}

		if numReturned.Cmp(numBlocks) >= 0 {
			// Index finished one full circle across the blocks
			// Used all of the blocks in this pool.
			return nil
		}
		numReturned.Add(numReturned, big.NewInt(1))

		// Return the block from this pool that corresponds with the index.
		return &cnet.IPNet{ipnet}
	}
}
<commit_msg>bugfix for transient-missing-a-block issue<commit_after>// Copyright (c) 2016 Tigera, Inc. All rights reserved.

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

package client

import (
	goerrors "errors"
	"fmt"
	"math/big"
	"math/rand"
	"net"
	"reflect"
	"time"

	"github.com/golang/glog"
	"github.com/tigera/libcalico-go/lib/api"
	"github.com/tigera/libcalico-go/lib/backend/model"
	"github.com/tigera/libcalico-go/lib/errors"
	cnet "github.com/tigera/libcalico-go/lib/net"
)

type blockReaderWriter struct {
	client *Client
}

func (rw blockReaderWriter) getAffineBlocks(host string, ver ipVersion, pool *cnet.IPNet) ([]cnet.IPNet, error) {
	// Lookup all blocks by providing an empty BlockListOptions
	// to the List operation.
	opts := model.BlockAffinityListOptions{Host: host, IPVersion: ver.Number}
	datastoreObjs, err := rw.client.backend.List(opts)
	if err != nil {
		if _, ok := err.(errors.ErrorResourceDoesNotExist); ok {
			// The block path does not exist yet.  This is OK - it means
			// there are no affine blocks.
			return []cnet.IPNet{}, nil

		} else {
			glog.Errorf("Error getting affine blocks: %s", err)
			return nil, err
		}
	}

	// Iterate through and extract the block CIDRs.
	ids := []cnet.IPNet{}
	for _, o := range datastoreObjs {
		k := o.Key.(model.BlockAffinityKey)
		ids = append(ids, k.CIDR)
	}
	return ids, nil
}

func (rw blockReaderWriter) claimNewAffineBlock(
	host string, version ipVersion, pool *cnet.IPNet, config IPAMConfig) (*cnet.IPNet, error) {

	// If pool is not nil, use the given pool.  Otherwise, default to
	// all configured pools.
	var pools []cnet.IPNet
	if pool != nil {
		// Validate the given pool is actually configured and matches the version.
		if !rw.isConfiguredPool(*pool) {
			estr := fmt.Sprintf("The given pool (%s) does not exist", pool.String())
			return nil, goerrors.New(estr)
		} else if version.Number != pool.Version() {
			estr := fmt.Sprintf("The given pool (%s) does not match IP version %d", pool.String(), version.Number)
			return nil, goerrors.New(estr)
		}
		pools = []cnet.IPNet{*pool}
	} else {
		// Default to all configured pools.
		allPools, err := rw.client.Pools().List(api.PoolMetadata{})
		if err != nil {
			glog.Errorf("Error reading configured pools: %s", err)
			return nil, err
		}

		// Grab all the IP networks in these pools.
		for _, p := range allPools.Items {
			// Don't include disabled pools or pools that don't match
			// the requested IP version.
			if !p.Spec.Disabled && version.Number == p.Metadata.CIDR.Version() {
				pools = append(pools, p.Metadata.CIDR)
			}
		}
	}

	// If there are no pools, we cannot assign addresses.
	if len(pools) == 0 {
		return nil, goerrors.New("No configured Calico pools")
	}

	// Iterate through pools to find a new block.
	glog.V(2).Infof("Claiming a new affine block for host '%s'", host)
	for _, pool := range pools {
		// Use a block generator to iterate through all of the blocks
		// that fall within the pool.
		blocks := randomBlockGenerator(pool)
		for subnet := blocks(); subnet != nil; subnet = blocks() {
			// Check if a block already exists for this subnet.
			glog.V(4).Infof("Getting block: %s", subnet.String())
			key := model.BlockKey{CIDR: subnet}
			_, err := rw.client.backend.Get(key)
			if err != nil {
				if _, ok := err.(errors.ErrorResourceDoesNotExist); ok {
					// The block does not yet exist in etcd.  Try to grab it.
					glog.V(3).Infof("Found free block: %+v", *subnet)
					err = rw.claimBlockAffinity(*subnet, host, config)
					return subnet, err
				} else {
					glog.Errorf("Error getting block: %s", err)
					return nil, err
				}
			}
		}
	}
	return nil, noFreeBlocksError("No Free Blocks")
}

func (rw blockReaderWriter) claimBlockAffinity(subnet cnet.IPNet, host string, config IPAMConfig) error {
	// Claim the block affinity for this host.
	glog.V(2).Infof("Host %s claiming block affinity for %s", host, subnet)
	obj := model.KVPair{
		Key:   model.BlockAffinityKey{Host: host, CIDR: subnet},
		Value: model.BlockAffinity{},
	}
	_, err := rw.client.backend.Create(&obj)

	// Create the new block.
	block := newBlock(subnet)
	block.HostAffinity = &host
	block.StrictAffinity = config.StrictAffinity

	// Create the new block in the datastore.
	o := model.KVPair{
		Key:   model.BlockKey{block.CIDR},
		Value: block.AllocationBlock,
	}
	_, err = rw.client.backend.Create(&o)
	if err != nil {
		if _, ok := err.(errors.ErrorResourceAlreadyExists); ok {
			// Block already exists, check affinity.
			glog.Warningf("Problem claiming block affinity:", err)
			obj, err := rw.client.backend.Get(model.BlockKey{subnet})
			if err != nil {
				glog.Errorf("Error reading block:", err)
				return err
			}

			// Pull out the allocationBlock object.
			b := allocationBlock{obj.Value.(model.AllocationBlock)}

			if b.HostAffinity != nil && *b.HostAffinity == host {
				// Block has affinity to this host, meaning another
				// process on this host claimed it.
				glog.V(3).Infof("Block %s already claimed by us.  Success", subnet)
				return nil
			}

			// Some other host beat us to this block.  Cleanup and return error.
			err = rw.client.backend.Delete(&model.KVPair{
				Key: model.BlockAffinityKey{Host: host, CIDR: b.CIDR},
			})
			if err != nil {
				glog.Errorf("Error cleaning up block affinity: %s", err)
				return err
			}
			return affinityClaimedError{Block: b}
		} else {
			return err
		}
	}
	return nil
}

func (rw blockReaderWriter) releaseBlockAffinity(host string, blockCIDR cnet.IPNet) error {
	for i := 0; i < ipamEtcdRetries; i++ {
		// Read the model.KVPair containing the block
		// and pull out the allocationBlock object.  We need to hold on to this
		// so that we can pass it back to the datastore on Update.
		obj, err := rw.client.backend.Get(model.BlockKey{CIDR: blockCIDR})
		if err != nil {
			glog.Errorf("Error getting block %s: %s", blockCIDR.String(), err)
			return err
		}
		b := allocationBlock{obj.Value.(model.AllocationBlock)}

		// Check that the block affinity matches the given affinity.
		if b.HostAffinity != nil && *b.HostAffinity != host {
			glog.Errorf("Mismatched affinity: %s != %s", *b.HostAffinity, host)
			return affinityClaimedError{Block: b}
		}

		if b.empty() {
			// If the block is empty, we can delete it.
			err := rw.client.backend.Delete(&model.KVPair{
				Key: model.BlockKey{CIDR: b.CIDR},
			})
			if err != nil {
				// Return the error unless the block didn't exist.
				if _, ok := err.(errors.ErrorResourceDoesNotExist); !ok {
					glog.Errorf("Error deleting block: %s", err)
					return err
				}
			}
		} else {
			// Otherwise, we need to remove affinity from it.
			// This prevents the host from automatically assigning
			// from this block unless we're allowed to overflow into
			// non-affine blocks.
			b.HostAffinity = nil

			// Pass back the original KVPair with the new
			// block information so we can do a CAS.
			obj.Value = b
			_, err = rw.client.backend.Update(obj)
			if err != nil {
				if _, ok := err.(errors.ErrorResourceUpdateConflict); ok {
					// CASError - continue.
					continue
				} else {
					return err
				}
			}
		}

		// We've removed / updated the block, so update the host config
		// to remove the CIDR.
		err = rw.client.backend.Delete(&model.KVPair{
			Key: model.BlockAffinityKey{Host: host, CIDR: b.CIDR},
		})
		if err != nil {
			// Return the error unless the affinity didn't exist.
			if _, ok := err.(errors.ErrorResourceDoesNotExist); !ok {
				glog.Errorf("Error deleting block affinity: %s", err)
				return err
			}
		}
		return nil

	}
	return goerrors.New("Max retries hit")
}

// withinConfiguredPools returns true if the given IP is within a configured
// Calico pool, and false otherwise.
func (rw blockReaderWriter) withinConfiguredPools(ip cnet.IP) bool {
	allPools, _ := rw.client.Pools().List(api.PoolMetadata{})
	for _, p := range allPools.Items {
		// Compare any enabled pools.
		if !p.Spec.Disabled && p.Metadata.CIDR.Contains(ip.IP) {
			return true
		}
	}
	return false
}

// isConfiguredPool returns true if the given IPNet is a configured
// Calico pool, and false otherwise.
func (rw blockReaderWriter) isConfiguredPool(cidr cnet.IPNet) bool {
	allPools, _ := rw.client.Pools().List(api.PoolMetadata{})
	for _, p := range allPools.Items {
		// Compare any enabled pools.
		if !p.Spec.Disabled && reflect.DeepEqual(p.Metadata.CIDR, cidr) {
			return true
		}
	}
	return false
}

// Generator to get list of block CIDRs which
// fall within the given pool. Returns nil when no more
// blocks can be generated.
func blockGenerator(pool cnet.IPNet) func() *cnet.IPNet {
	// Determine the IP type to use.
	version := getIPVersion(cnet.IP{pool.IP})
	ip := cnet.IP{pool.IP}
	return func() *cnet.IPNet {
		returnIP := ip
		ip = incrementIP(ip, big.NewInt(blockSize))
		if pool.Contains(ip.IP) {
			ipnet := net.IPNet{returnIP.IP, version.BlockPrefixMask}
			cidr := cnet.IPNet{ipnet}
			return &cidr
		} else {
			return nil
		}
	}
}

// Returns a generator that, when called, returns a random
// block from the given pool.  When there are no blocks left,
// the it returns nil.
func randomBlockGenerator(pool cnet.IPNet) func() *cnet.IPNet {

	// Determine the IP type to use.
	version := getIPVersion(cnet.IP{pool.IP})
	baseIP := cnet.IP{pool.IP}

	// Determine the number of blocks within this pool.
	ones, size := pool.Mask.Size()
	prefixLen := size - ones
	numIP := new(big.Int).Exp(big.NewInt(2), big.NewInt(int64(prefixLen)), nil)
	numBlocks := new(big.Int)
	numBlocks.Div(numIP, big.NewInt(blockSize))

	// Start at a random offset index
	source := rand.NewSource(time.Now().UnixNano())
	randm := rand.New(source)

	// initialIndex keeps track of the random starting point
	initialIndex := new(big.Int)
	initialIndex.Rand(randm, numBlocks)

	// i keeps track of current index while walking the blocks in a pool
	i := initialIndex

	// numReturned keeps track of number of blocks returned
	numReturned := big.NewInt(0)

	// numDiff = numBlocks - i
	numDiff := new(big.Int)

	return func() *cnet.IPNet {
		// The `big.NewInt(0)` part creates a temp variable and assigns the result of multiplication of `i` and `big.NewInt(blockSize)`
		// Note: we are not using `i.Mul()` because that will assign the result of the multiplication to `i`, which will cause unexpected issues
		ip := incrementIP(baseIP, big.NewInt(0).Mul(i, big.NewInt(blockSize)))
		ipnet := net.IPNet{ip.IP, version.BlockPrefixMask}

		numDiff.Sub(numBlocks, i)

		if numDiff.Cmp(big.NewInt(1)) <= 0 {
			// Index has reached end of the blocks;
			// Loop back to beginning of pool rather than
			// increment, because incrementing would put us outside of the pool.
			i = big.NewInt(0)
		} else {
			// Increment to the next block
			i.Add(i, big.NewInt(1))
		}

		if numReturned.Cmp(numBlocks) >= 0 {
			// Index finished one full circle across the blocks
			// Used all of the blocks in this pool.
			return nil
		}
		numReturned.Add(numReturned, big.NewInt(1))

		// Return the block from this pool that corresponds with the index.
		return &cnet.IPNet{ipnet}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	l "log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Shopify/sarama"
	"github.com/grafana/globalconf"
	"github.com/grafana/metrictank/api"
	"github.com/grafana/metrictank/cluster"
	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/idx/memory"
	"github.com/grafana/metrictank/input"
	inKafka "github.com/grafana/metrictank/input/kafkamdm"
	"github.com/grafana/metrictank/logger"
	"github.com/grafana/metrictank/mdata"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/schema/msg"
	cassandraStore "github.com/grafana/metrictank/store/cassandra"
	"github.com/raintank/dur"
	log "github.com/sirupsen/logrus"
)

var (
	// metrictank
	aggMetrics  *mdata.AggMetrics
	metricIndex idx.MetricIndex
	inputKafka  input.Plugin
	store       mdata.Store

	// config file
	confFile = flag.String("config", "/etc/metrictank/metrictank.ini", "config file path")

	// Data: the following configs are the same with the normal metrictank configs:
	chunkMaxStaleStr  = flag.String("chunk-max-stale", "1m", "chunk max stale age.")
	metricMaxStaleStr = flag.String("metric-max-stale", "5m", "metric max stale age.")
	gcIntervalStr     = flag.String("gc-interval", "2m", "gc interval.")
	publicOrg         = flag.Int("public-org", 0, "org Id")
	timeout           = flag.Int("timeout", 10, "the tool will exit if no kafka message is received during this interval ")
	logLevel          = flag.String("log-level", "info", "log level. panic|fatal|error|warning|info|debug")

	// backfilling
	lastRcvTime int64        // epoch time when the previous kafka message was received
	rcvdPoints  int64        // number of received datapoints
	mux         sync.Mutex   // mutex to protect lastRcvTime
	handler     inputHandler // input message handler to track the last kafka receive event and handles kafka messages
)

// a kafka message handler that implements the input.Handler interface
type inputHandler struct {
	handler  input.DefaultHandler // default handler that processes metric metadata and points
	finished chan bool
}

func newInputHandler(metrics mdata.Metrics, metricIndex idx.MetricIndex, pluginName string) inputHandler {
	dh := input.NewDefaultHandler(metrics, metricIndex, pluginName)
	return inputHandler{
		handler:  dh,
		finished: make(chan bool),
	}
}

// input.Handler interface
func (ih inputHandler) ProcessMetricData(metric *schema.MetricData, partition int32) {
	ih.handler.ProcessMetricData(metric, partition)
	mux.Lock()
	defer mux.Unlock()
	lastRcvTime = int64(time.Now().Unix())
	rcvdPoints++
	if rcvdPoints%10000 == 0 {
		log.Infof("Received %d datapoints", rcvdPoints)
	}
}

// input.Handler interface
func (ih inputHandler) ProcessMetricPoint(point schema.MetricPoint, format msg.Format, partition int32) {
	ih.handler.ProcessMetricPoint(point, format, partition)
	mux.Lock()
	defer mux.Unlock()
	lastRcvTime = int64(time.Now().Unix())
	rcvdPoints++
	if rcvdPoints%10000 == 0 {
		log.Infof("Received %d datapoints", rcvdPoints)
	}
}

func main() {

	flag.Usage = func() {
		fmt.Println("mt-backfill")
		fmt.Println("Consumes data from Kafka and backfills chunks to Cassandra (only supports cassandra).")
		fmt.Println("Does not update the index table. Useful when existing series in Metrictank need historical data.")
		fmt.Println()
		fmt.Printf("Parameters:\n\n")
		flag.PrintDefaults()
		fmt.Println()
		fmt.Println("Config file supports same elements as `metrictank` command, but only supports Kafka in and Cassandra out.")
		fmt.Println()
		fmt.Printf("Example:\n\n")
		fmt.Printf("  mt-backfill -config /etc/metrictank/backfill.ini -timeout 600\n\n")
	}

	flag.Parse()

	// logger
	formatter := &logger.TextFormatter{}
	formatter.TimestampFormat = "2006-01-02 15:04:05.000"
	log.SetFormatter(formatter)
	lvl, err := log.ParseLevel(*logLevel)
	if err != nil {
		log.Fatalf("failed to parse log-level, %s", err.Error())
	}
	log.SetLevel(lvl)
	log.Infof("logging level set to '%s'", *logLevel)
	log.Infof("metrictank backfilling")

	// Only try and parse the conf file if it exists
	path := ""
	if _, err := os.Stat(*confFile); err == nil {
		path = *confFile
	}
	config, err := globalconf.NewWithOptions(&globalconf.Options{
		Filename:  path,
		EnvPrefix: "MT_",
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "configuration file error: %s", err)
		os.Exit(1)
	}

	// load configs
	inKafka.ConfigSetup()
	memory.ConfigSetup()
	api.ConfigSetup()
	cluster.ConfigSetup()
	mdata.ConfigSetup()
	cassandraStore.ConfigSetup()
	config.ParseAll()

	// cluster is required because of aggMetric.add()
	// this should be configured as single mode
	api.ConfigProcess()
	cluster.ConfigProcess()
	addrParts := strings.Split(api.Addr, ":")
	port, err := strconv.ParseInt(addrParts[len(addrParts)-1], 10, 64)
	if err != nil {
		log.Fatalf("Could not parse port from listenAddr. %s", api.Addr)
	}
	cluster.Init("backfill", "none", time.Now(), "http", int(port))

	// other settings
	inKafka.ConfigProcess("backfill")
	mdata.ConfigProcess()
	memory.ConfigProcess()

	// interrupt signals
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// backend store
	cassandraStore.CliConfig.Enabled = true
	store, err = cassandraStore.NewCassandraStore(cassandraStore.CliConfig, mdata.TTLs(), uint32(cassandraStore.CliConfig.MaxChunkSpan.Seconds()))
	if err != nil {
		log.Fatalf("failed to initialize cassandra store. %s", err)
	}

	// memory store
	chunkMaxStale := dur.MustParseNDuration("chunk-max-stale", *chunkMaxStaleStr)
	metricMaxStale := dur.MustParseNDuration("metric-max-stale", *metricMaxStaleStr)
	gcInterval := time.Duration(dur.MustParseNDuration("gc-interval", *gcIntervalStr)) * time.Second
	aggMetrics = mdata.NewAggMetrics(store, nil, false, nil /*todo*/, chunkMaxStale, metricMaxStale, gcInterval)

	// input
	// we use kafkamdm as input
	inKafka.Enabled = true
	sarama.Logger = l.New(os.Stdout, "[Sarama] ", l.LstdFlags)
	inputKafka = inKafka.New()

	// cluster manager
	cluster.Start()

	if *publicOrg < 0 {
		log.Fatal("public-org cannot be <0")
	}
	idx.OrgIdPublic = uint32(*publicOrg)

	memory.Enabled = true
	metricIndex = memory.New()

	// load index entries
	err = metricIndex.Init()
	if err != nil {
		log.Fatalf("Failed to initialize metricIndex: %s", err.Error())
	}

	// start input
	ctx, cancel := context.WithCancel(context.Background())
	handler = newInputHandler(aggMetrics, metricIndex, "kafkamdm")
	err = inputKafka.Start(handler, cancel)
	if err != nil {
		shutdown()
		log.Warn("Cannot start input.")
		return
	}
	inputKafka.MaintainPriority()
	lastRcvTime = int64(time.Now().Unix())
	go handlerTimeout()

	cluster.Manager.SetReady()

	// wait for shutdown
	select {
	case sig := <-sigChan:
		log.Infof("Received signal %q. Shutting down", sig)
	case <-ctx.Done():
		log.Info("The input plugin signalled a fatal error. Shutting down")
	case <-handler.finished:
		log.Infof("Received finished signal from input handler. Shutting down")
	}
	shutdown()

	defer cancel()
}

// if there is no update in 3 mins, shut down the tool
func handlerTimeout() {
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case now := <-ticker.C:
			mux.Lock()
			prevRcv := lastRcvTime
			mux.Unlock()
			// wait for timeout seconds, if there is no more input, shut down
			if now.Unix()-prevRcv > int64(*timeout) {
				log.Infof("Handler timeout, shutting down")
				handler.finished <- true
				ticker.Stop()
				return
			}
		}
	}
}

// normal shutdown
func shutdown() {
	cluster.Stop()
	timer := time.NewTimer(time.Second * 20)
	kafkaStopped := make(chan bool)
	go func() {
		log.Infof("Shutting down kafka consumer")
		inputKafka.Stop()
		log.Infof("kafka consumer finished shutdown")
		kafkaStopped <- true
	}()

	select {
	case <-timer.C:
		log.Warn("Plugin shutdown timeout.")
	case <-kafkaStopped:
		timer.Stop()
	}
	store.Stop()
	log.Info("shutting down.")
}
<commit_msg>Update interface<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	l "log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Shopify/sarama"
	"github.com/grafana/globalconf"
	"github.com/grafana/metrictank/api"
	"github.com/grafana/metrictank/cluster"
	"github.com/grafana/metrictank/idx"
	"github.com/grafana/metrictank/idx/memory"
	"github.com/grafana/metrictank/input"
	inKafka "github.com/grafana/metrictank/input/kafkamdm"
	"github.com/grafana/metrictank/logger"
	"github.com/grafana/metrictank/mdata"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/schema/msg"
	cassandraStore "github.com/grafana/metrictank/store/cassandra"
	"github.com/raintank/dur"
	log "github.com/sirupsen/logrus"
)

var (
	// metrictank
	aggMetrics  *mdata.AggMetrics
	metricIndex idx.MetricIndex
	inputKafka  input.Plugin
	store       mdata.Store

	// config file
	confFile = flag.String("config", "/etc/metrictank/metrictank.ini", "config file path")

	// Data: the following configs are the same with the normal metrictank configs:
	chunkMaxStaleStr  = flag.String("chunk-max-stale", "1m", "chunk max stale age.")
	metricMaxStaleStr = flag.String("metric-max-stale", "5m", "metric max stale age.")
	gcIntervalStr     = flag.String("gc-interval", "2m", "gc interval.")
	publicOrg         = flag.Int("public-org", 0, "org Id")
	timeout           = flag.Int("timeout", 10, "the tool will exit if no kafka message is received during this interval ")
	logLevel          = flag.String("log-level", "info", "log level. panic|fatal|error|warning|info|debug")

	// backfilling
	lastRcvTime int64        // epoch time when the previous kafka message was received
	rcvdPoints  int64        // number of received datapoints
	mux         sync.Mutex   // mutex to protect lastRcvTime
	handler     inputHandler // input message handler to track the last kafka receive event and handles kafka messages
)

// a kafka message handler that implements the input.Handler interface
type inputHandler struct {
	handler  input.DefaultHandler // default handler that processes metric metadata and points
	finished chan bool
}

func newInputHandler(metrics mdata.Metrics, metricIndex idx.MetricIndex, pluginName string) inputHandler {
	dh := input.NewDefaultHandler(metrics, metricIndex, pluginName)
	return inputHandler{
		handler:  dh,
		finished: make(chan bool),
	}
}

// input.Handler interface
func (ih inputHandler) ProcessMetricData(metric *schema.MetricData, partition int32) {
	ih.handler.ProcessMetricData(metric, partition)
	mux.Lock()
	defer mux.Unlock()
	lastRcvTime = int64(time.Now().Unix())
	rcvdPoints++
	if rcvdPoints%10000 == 0 {
		log.Infof("Received %d datapoints", rcvdPoints)
	}
}

// input.Handler interface
func (ih inputHandler) ProcessMetricPoint(point schema.MetricPoint, format msg.Format, partition int32) {
	ih.handler.ProcessMetricPoint(point, format, partition)
	mux.Lock()
	defer mux.Unlock()
	lastRcvTime = int64(time.Now().Unix())
	rcvdPoints++
	if rcvdPoints%10000 == 0 {
		log.Infof("Received %d datapoints", rcvdPoints)
	}
}

// input.Handler interface
func (ih inputHandler) ProcessIndexControlMsg(msg schema.ControlMsg, partition int32) {
	// Ignore index messages
}

func main() {

	flag.Usage = func() {
		fmt.Println("mt-backfill")
		fmt.Println("Consumes data from Kafka and backfills chunks to Cassandra (only supports cassandra).")
		fmt.Println("Does not update the index table. Useful when existing series in Metrictank need historical data.")
		fmt.Println()
		fmt.Printf("Parameters:\n\n")
		flag.PrintDefaults()
		fmt.Println()
		fmt.Println("Config file supports same elements as `metrictank` command, but only supports Kafka in and Cassandra out.")
		fmt.Println()
		fmt.Printf("Example:\n\n")
		fmt.Printf("  mt-backfill -config /etc/metrictank/backfill.ini -timeout 600\n\n")
	}

	flag.Parse()

	// logger
	formatter := &logger.TextFormatter{}
	formatter.TimestampFormat = "2006-01-02 15:04:05.000"
	log.SetFormatter(formatter)
	lvl, err := log.ParseLevel(*logLevel)
	if err != nil {
		log.Fatalf("failed to parse log-level, %s", err.Error())
	}
	log.SetLevel(lvl)
	log.Infof("logging level set to '%s'", *logLevel)
	log.Infof("metrictank backfilling")

	// Only try and parse the conf file if it exists
	path := ""
	if _, err := os.Stat(*confFile); err == nil {
		path = *confFile
	}
	config, err := globalconf.NewWithOptions(&globalconf.Options{
		Filename:  path,
		EnvPrefix: "MT_",
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "configuration file error: %s", err)
		os.Exit(1)
	}

	// load configs
	inKafka.ConfigSetup()
	memory.ConfigSetup()
	api.ConfigSetup()
	cluster.ConfigSetup()
	mdata.ConfigSetup()
	cassandraStore.ConfigSetup()
	config.ParseAll()

	// cluster is required because of aggMetric.add()
	// this should be configured as single mode
	api.ConfigProcess()
	cluster.ConfigProcess()
	addrParts := strings.Split(api.Addr, ":")
	port, err := strconv.ParseInt(addrParts[len(addrParts)-1], 10, 64)
	if err != nil {
		log.Fatalf("Could not parse port from listenAddr. %s", api.Addr)
	}
	cluster.Init("backfill", "none", time.Now(), "http", int(port))

	// other settings
	inKafka.ConfigProcess("backfill")
	mdata.ConfigProcess()
	memory.ConfigProcess()

	// interrupt signals
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// backend store
	cassandraStore.CliConfig.Enabled = true
	store, err = cassandraStore.NewCassandraStore(cassandraStore.CliConfig, mdata.TTLs(), uint32(cassandraStore.CliConfig.MaxChunkSpan.Seconds()))
	if err != nil {
		log.Fatalf("failed to initialize cassandra store. %s", err)
	}

	// memory store
	chunkMaxStale := dur.MustParseNDuration("chunk-max-stale", *chunkMaxStaleStr)
	metricMaxStale := dur.MustParseNDuration("metric-max-stale", *metricMaxStaleStr)
	gcInterval := time.Duration(dur.MustParseNDuration("gc-interval", *gcIntervalStr)) * time.Second
	aggMetrics = mdata.NewAggMetrics(store, nil, false, nil /*todo*/, chunkMaxStale, metricMaxStale, gcInterval)

	// input
	// we use kafkamdm as input
	inKafka.Enabled = true
	sarama.Logger = l.New(os.Stdout, "[Sarama] ", l.LstdFlags)
	inputKafka = inKafka.New()

	// cluster manager
	cluster.Start()

	if *publicOrg < 0 {
		log.Fatal("public-org cannot be <0")
	}
	idx.OrgIdPublic = uint32(*publicOrg)

	memory.Enabled = true
	metricIndex = memory.New()

	// load index entries
	err = metricIndex.Init()
	if err != nil {
		log.Fatalf("Failed to initialize metricIndex: %s", err.Error())
	}

	// start input
	ctx, cancel := context.WithCancel(context.Background())
	handler = newInputHandler(aggMetrics, metricIndex, "kafkamdm")
	err = inputKafka.Start(handler, cancel)
	if err != nil {
		shutdown()
		log.Warn("Cannot start input.")
		return
	}
	inputKafka.MaintainPriority()
	lastRcvTime = int64(time.Now().Unix())
	go handlerTimeout()

	cluster.Manager.SetReady()

	// wait for shutdown
	select {
	case sig := <-sigChan:
		log.Infof("Received signal %q. Shutting down", sig)
	case <-ctx.Done():
		log.Info("The input plugin signalled a fatal error. Shutting down")
	case <-handler.finished:
		log.Infof("Received finished signal from input handler. Shutting down")
	}
	shutdown()

	defer cancel()
}

// if there is no update in 3 mins, shut down the tool
func handlerTimeout() {
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case now := <-ticker.C:
			mux.Lock()
			prevRcv := lastRcvTime
			mux.Unlock()
			// wait for timeout seconds, if there is no more input, shut down
			if now.Unix()-prevRcv > int64(*timeout) {
				log.Infof("Handler timeout, shutting down")
				handler.finished <- true
				ticker.Stop()
				return
			}
		}
	}
}

// normal shutdown
func shutdown() {
	cluster.Stop()
	timer := time.NewTimer(time.Second * 20)
	kafkaStopped := make(chan bool)
	go func() {
		log.Infof("Shutting down kafka consumer")
		inputKafka.Stop()
		log.Infof("kafka consumer finished shutdown")
		kafkaStopped <- true
	}()

	select {
	case <-timer.C:
		log.Warn("Plugin shutdown timeout.")
	case <-kafkaStopped:
		timer.Stop()
	}
	store.Stop()
	log.Info("shutting down.")
}
<|endoftext|>


<commit_before>package middleware

import (
	"fmt"
	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gin-gonic/gin"
	"net/http"

	"github.com/techjanitor/pram-post/config"
	e "github.com/techjanitor/pram-post/errors"
	u "github.com/techjanitor/pram-post/utils"
)

type Permission uint

const (
	All Permission = 1 << iota
	Registered
	Moderators
	Admins
)

// checks for session cookie and handles permissions
func Auth(perm Permission) gin.HandlerFunc {
	return func(c *gin.Context) {

		// set default anonymous user
		user := u.User{
			Id:              1,
			Group:           uint(All),
			IsAuthenticated: false,
		}

		// parse jwt token if its there
		token, err := jwt.ParseFromRequest(c.Request, func(token *jwt.Token) (interface{}, error) {

			// check alg to make sure its hmac
			_, ok := token.Method.(*jwt.SigningMethodHMAC)
			if !ok {
				return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
			}

			// compare with secret from settings
			return []byte(config.Settings.Session.Secret), nil
		})
		if err != nil && err != jwt.ErrNoTokenInRequest {
			// if theres some jwt error then return unauth
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(err)
			c.Abort()
			return
		}

		// process token
		if token != nil {
			// if the token is valid set the data
			if err == nil && token.Valid {

				// get uid from jwt, cast to float
				uid, ok := token.Claims["user_id"].(float64)
				if !ok {
					c.JSON(e.ErrorMessage(e.ErrInternalError))
					c.Error(err)
					c.Abort()
					return
				}

				// set user id in user struct
				user.Id = uint(uid)

				// get the rest of the user info
				err = user.Info()
				if err != nil {
					c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
					c.Error(err)
					c.Abort()
					return
				}

			} else {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err)
				c.Abort()
				return
			}

		}

		// check if user meets set permissions
		if user.Group < uint(perm) {
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(e.ErrUnauthorized)
			c.Abort()
			return
		}

		// set user data
		c.Set("userdata", user)

		c.Next()

	}

}
<commit_msg>perm iota<commit_after>package middleware

import (
	"fmt"
	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gin-gonic/gin"
	"net/http"

	"github.com/techjanitor/pram-post/config"
	e "github.com/techjanitor/pram-post/errors"
	u "github.com/techjanitor/pram-post/utils"
)

type Permission uint

const (
	All Permission = 1 + iota
	Registered
	Moderators
	Admins
)

// checks for session cookie and handles permissions
func Auth(perm Permission) gin.HandlerFunc {
	return func(c *gin.Context) {

		// set default anonymous user
		user := u.User{
			Id:              1,
			Group:           uint(All),
			IsAuthenticated: false,
		}

		// parse jwt token if its there
		token, err := jwt.ParseFromRequest(c.Request, func(token *jwt.Token) (interface{}, error) {

			// check alg to make sure its hmac
			_, ok := token.Method.(*jwt.SigningMethodHMAC)
			if !ok {
				return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
			}

			// compare with secret from settings
			return []byte(config.Settings.Session.Secret), nil
		})
		if err != nil && err != jwt.ErrNoTokenInRequest {
			// if theres some jwt error then return unauth
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(err)
			c.Abort()
			return
		}

		// process token
		if token != nil {
			// if the token is valid set the data
			if err == nil && token.Valid {

				// get uid from jwt, cast to float
				uid, ok := token.Claims["user_id"].(float64)
				if !ok {
					c.JSON(e.ErrorMessage(e.ErrInternalError))
					c.Error(err)
					c.Abort()
					return
				}

				// set user id in user struct
				user.Id = uint(uid)

				// get the rest of the user info
				err = user.Info()
				if err != nil {
					c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
					c.Error(err)
					c.Abort()
					return
				}

			} else {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err)
				c.Abort()
				return
			}

		}

		// check if user meets set permissions
		if user.Group < uint(perm) {
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(e.ErrUnauthorized)
			c.Abort()
			return
		}

		// set user data
		c.Set("userdata", user)

		c.Next()

	}

}
<|endoftext|>


<commit_before>package rollbar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"hash/adler32"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"time"
)

const (
	NAME    = "go-rollbar"
	VERSION = "0.3.1"

	// Severity levels
	CRIT  = "critical"
	ERR   = "error"
	WARN  = "warning"
	INFO  = "info"
	DEBUG = "debug"

	FILTERED = "[FILTERED]"
)

var (
	// Rollbar access token. If this is blank, no errors will be reported to
	// Rollbar.
	Token = ""

	// All errors and messages will be submitted under this environment.
	Environment = "development"

	// Platform, default to OS, but could be change ('client' for instance)
	Platform = runtime.GOOS

	// API endpoint for Rollbar.
	Endpoint = "https://api.rollbar.com/api/1/item/"

	// Maximum number of errors allowed in the sending queue before we start
	// dropping new errors on the floor.
	Buffer = 1000

	// Filter GET and POST parameters from being sent to Rollbar.
	FilterFields = regexp.MustCompile("password|secret|token")

	// Output of error, by default stderr
	ErrorWriter = os.Stderr

	// Queue of messages to be sent.
	bodyChannel chan map[string]interface{}
	waitGroup   sync.WaitGroup
)

// Fields can be used to pass arbitrary data to the Rollbar API.
type Field struct {
	Name string
	Data interface{}
}

// -- Setup

func init() {
	bodyChannel = make(chan map[string]interface{}, Buffer)

	go func() {
		for body := range bodyChannel {
			post(body)
			waitGroup.Done()
		}
	}()
}

// -- Error reporting

// Error asynchronously sends an error to Rollbar with the given severity
// level. You can pass, optionally, custom Fields to be passed on to Rollbar.
func Error(level string, err error, fields ...*Field) {
	ErrorWithStackSkip(level, err, 1, fields...)
}

// ErrorWithStackSkip asynchronously sends an error to Rollbar with the given
// severity level and a given number of stack trace frames skipped. You can
// pass, optionally, custom Fields to be passed on to Rollbar.
func ErrorWithStackSkip(level string, err error, skip int, fields ...*Field) {
	stack := BuildStack(2 + skip)
	ErrorWithStack(level, err, stack, fields...)
}

// ErrorWithStack asynchronously sends and error to Rollbar with the given
// stacktrace and (optionally) custom Fields to be passed on to Rollbar.
func ErrorWithStack(level string, err error, stack Stack, fields ...*Field) {
	buildAndPushError(level, err, stack, fields...)
}

// RequestError asynchronously sends an error to Rollbar with the given
// severity level and request-specific information. You can pass, optionally,
// custom Fields to be passed on to Rollbar.
func RequestError(level string, r *http.Request, err error, fields ...*Field) {
	RequestErrorWithStackSkip(level, r, err, 1, fields...)
}

// RequestErrorWithStackSkip asynchronously sends an error to Rollbar with the
// given severity level and a given number of stack trace frames skipped, in
// addition to extra request-specific information. You can pass, optionally,
// custom Fields to be passed on to Rollbar.
func RequestErrorWithStackSkip(level string, r *http.Request, err error, skip int, fields ...*Field) {
	stack := BuildStack(2 + skip)
	RequestErrorWithStack(level, r, err, stack, fields...)
}

// RequestErrorWithStack asynchronously sends an error to Rollbar with the
// given severity level, request-specific information provided by the given
// http.Request, and a custom Stack. You You can pass, optionally, custom
// Fields to be passed on to Rollbar.
func RequestErrorWithStack(level string, r *http.Request, err error, stack Stack, fields ...*Field) {
	buildAndPushError(level, err, stack, &Field{Name: "request", Data: errorRequest(r)})
}

func buildError(level string, err error, stack Stack, fields ...*Field) map[string]interface{} {
	body := buildBody(level, err.Error())
	data := body["data"].(map[string]interface{})
	errBody, fingerprint := errorBody(err, stack)
	data["body"] = errBody
	data["fingerprint"] = fingerprint

	for _, field := range fields {
		data[field.Name] = field.Data
	}

	return body
}

func buildAndPushError(level string, err error, stack Stack, fields ...*Field) {
	push(buildError(level, err, stack, fields...))
}

// -- Message reporting

// Message asynchronously sends a message to Rollbar with the given severity
// level.
func Message(level string, msg string) {
	body := buildBody(level, msg)
	data := body["data"].(map[string]interface{})
	data["body"] = messageBody(msg)

	push(body)
}

// -- Misc.

// Wait will block until the queue of errors / messages is empty. This allows
// you to ensure that errors / messages are sent to Rollbar before exiting an
// application.
func Wait() {
	waitGroup.Wait()
}

// Build the main JSON structure that will be sent to Rollbar with the
// appropriate metadata.
func buildBody(level, title string) map[string]interface{} {
	timestamp := time.Now().Unix()
	hostname, _ := os.Hostname()

	return map[string]interface{}{
		"access_token": Token,
		"data": map[string]interface{}{
			"environment": Environment,
			"title":       title,
			"level":       level,
			"timestamp":   timestamp,
			"platform":    Platform,
			"language":    "go",
			"server": map[string]interface{}{
				"host": hostname,
			},
			"notifier": map[string]interface{}{
				"name":    NAME,
				"version": VERSION,
			},
		},
	}
}

// errorBody generates a Rollbar error body with a given stack trace.
func errorBody(err error, stack Stack) (map[string]interface{}, string) {
	fingerprint := stack.Fingerprint()
	errBody := map[string]interface{}{
		"trace": map[string]interface{}{
			"frames": stack,
			"exception": map[string]interface{}{
				"class":   errorClass(err),
				"message": err.Error(),
			},
		},
	}
	return errBody, fingerprint
}

// errorRequest extracts details from a Request in a format that Rollbar
// accepts.
func errorRequest(r *http.Request) map[string]interface{} {
	cleanQuery := filterParams(r.URL.Query())

	return map[string]interface{}{
		"url":     r.URL.String(),
		"method":  r.Method,
		"headers": flattenValues(r.Header),

		// GET params
		"query_string": url.Values(cleanQuery).Encode(),
		"GET":          flattenValues(cleanQuery),

		// POST / PUT params
		"POST": flattenValues(filterParams(r.Form)),
	}
}

// filterParams filters sensitive information like passwords from being sent to
// Rollbar.
func filterParams(values map[string][]string) map[string][]string {
	for key, _ := range values {
		if FilterFields.Match([]byte(key)) {
			values[key] = []string{FILTERED}
		}
	}

	return values
}

func flattenValues(values map[string][]string) map[string]interface{} {
	result := make(map[string]interface{})

	for k, v := range values {
		if len(v) == 1 {
			result[k] = v[0]
		} else {
			result[k] = v
		}
	}

	return result
}

// Build a message inner-body for the given message string.
func messageBody(s string) map[string]interface{} {
	return map[string]interface{}{
		"message": map[string]interface{}{
			"body": s,
		},
	}
}

func errorClass(err error) string {
	class := reflect.TypeOf(err).String()
	if class == "" {
		return "panic"
	} else if class == "*errors.errorString" {
		checksum := adler32.Checksum([]byte(err.Error()))
		return fmt.Sprintf("{%x}", checksum)
	} else {
		return strings.TrimPrefix(class, "*")
	}
}

// -- POST handling

// Queue the given JSON body to be POSTed to Rollbar.
func push(body map[string]interface{}) {
	if len(bodyChannel) < Buffer {
		waitGroup.Add(1)
		bodyChannel <- body
	} else {
		stderr("buffer full, dropping error on the floor")
	}
}

// POST the given JSON body to Rollbar synchronously.
func post(body map[string]interface{}) {
	if len(Token) == 0 {
		stderr("empty token")
		return
	}

	jsonBody, err := json.Marshal(body)
	if err != nil {
		stderr("failed to encode payload: %s", err.Error())
		return
	}

	resp, err := http.Post(Endpoint, "application/json", bytes.NewReader(jsonBody))
	if err != nil {
		stderr("POST failed: %s", err.Error())
	} else if resp.StatusCode != 200 {
		stderr("received response: %s", resp.Status)
	}
	if resp != nil {
		resp.Body.Close()
	}
}

// -- stderr
func stderr(format string, args ...interface{}) {
	if ErrorWriter != nil {
		format = "Rollbar error: " + format + "\n"
		fmt.Fprintf(ErrorWriter, format, args...)
	}
}
<commit_msg>Fields are no longer overwritten<commit_after>package rollbar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"hash/adler32"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"time"
)

const (
	NAME    = "go-rollbar"
	VERSION = "0.3.1"

	// Severity levels
	CRIT  = "critical"
	ERR   = "error"
	WARN  = "warning"
	INFO  = "info"
	DEBUG = "debug"

	FILTERED = "[FILTERED]"
)

var (
	// Rollbar access token. If this is blank, no errors will be reported to
	// Rollbar.
	Token = ""

	// All errors and messages will be submitted under this environment.
	Environment = "development"

	// Platform, default to OS, but could be change ('client' for instance)
	Platform = runtime.GOOS

	// API endpoint for Rollbar.
	Endpoint = "https://api.rollbar.com/api/1/item/"

	// Maximum number of errors allowed in the sending queue before we start
	// dropping new errors on the floor.
	Buffer = 1000

	// Filter GET and POST parameters from being sent to Rollbar.
	FilterFields = regexp.MustCompile("password|secret|token")

	// Output of error, by default stderr
	ErrorWriter = os.Stderr

	// Queue of messages to be sent.
	bodyChannel chan map[string]interface{}
	waitGroup   sync.WaitGroup
)

// Fields can be used to pass arbitrary data to the Rollbar API.
type Field struct {
	Name string
	Data interface{}
}

// -- Setup

func init() {
	bodyChannel = make(chan map[string]interface{}, Buffer)

	go func() {
		for body := range bodyChannel {
			post(body)
			waitGroup.Done()
		}
	}()
}

// -- Error reporting

// Error asynchronously sends an error to Rollbar with the given severity
// level. You can pass, optionally, custom Fields to be passed on to Rollbar.
func Error(level string, err error, fields ...*Field) {
	ErrorWithStackSkip(level, err, 1, fields...)
}

// ErrorWithStackSkip asynchronously sends an error to Rollbar with the given
// severity level and a given number of stack trace frames skipped. You can
// pass, optionally, custom Fields to be passed on to Rollbar.
func ErrorWithStackSkip(level string, err error, skip int, fields ...*Field) {
	stack := BuildStack(2 + skip)
	ErrorWithStack(level, err, stack, fields...)
}

// ErrorWithStack asynchronously sends and error to Rollbar with the given
// stacktrace and (optionally) custom Fields to be passed on to Rollbar.
func ErrorWithStack(level string, err error, stack Stack, fields ...*Field) {
	buildAndPushError(level, err, stack, fields...)
}

// RequestError asynchronously sends an error to Rollbar with the given
// severity level and request-specific information. You can pass, optionally,
// custom Fields to be passed on to Rollbar.
func RequestError(level string, r *http.Request, err error, fields ...*Field) {
	RequestErrorWithStackSkip(level, r, err, 1, fields...)
}

// RequestErrorWithStackSkip asynchronously sends an error to Rollbar with the
// given severity level and a given number of stack trace frames skipped, in
// addition to extra request-specific information. You can pass, optionally,
// custom Fields to be passed on to Rollbar.
func RequestErrorWithStackSkip(level string, r *http.Request, err error, skip int, fields ...*Field) {
	stack := BuildStack(2 + skip)
	RequestErrorWithStack(level, r, err, stack, fields...)
}

// RequestErrorWithStack asynchronously sends an error to Rollbar with the
// given severity level, request-specific information provided by the given
// http.Request, and a custom Stack. You You can pass, optionally, custom
// Fields to be passed on to Rollbar.
func RequestErrorWithStack(level string, r *http.Request, err error, stack Stack, fields ...*Field) {
	buildAndPushError(level, err, stack, append(fields, &Field{Name: "request", Data: errorRequest(r)})...)
}

func buildError(level string, err error, stack Stack, fields ...*Field) map[string]interface{} {
	body := buildBody(level, err.Error())
	data := body["data"].(map[string]interface{})
	errBody, fingerprint := errorBody(err, stack)
	data["body"] = errBody
	data["fingerprint"] = fingerprint

	for _, field := range fields {
		data[field.Name] = field.Data
	}

	return body
}

func buildAndPushError(level string, err error, stack Stack, fields ...*Field) {
	push(buildError(level, err, stack, fields...))
}

// -- Message reporting

// Message asynchronously sends a message to Rollbar with the given severity
// level.
func Message(level string, msg string) {
	body := buildBody(level, msg)
	data := body["data"].(map[string]interface{})
	data["body"] = messageBody(msg)

	push(body)
}

// -- Misc.

// Wait will block until the queue of errors / messages is empty. This allows
// you to ensure that errors / messages are sent to Rollbar before exiting an
// application.
func Wait() {
	waitGroup.Wait()
}

// Build the main JSON structure that will be sent to Rollbar with the
// appropriate metadata.
func buildBody(level, title string) map[string]interface{} {
	timestamp := time.Now().Unix()
	hostname, _ := os.Hostname()

	return map[string]interface{}{
		"access_token": Token,
		"data": map[string]interface{}{
			"environment": Environment,
			"title":       title,
			"level":       level,
			"timestamp":   timestamp,
			"platform":    Platform,
			"language":    "go",
			"server": map[string]interface{}{
				"host": hostname,
			},
			"notifier": map[string]interface{}{
				"name":    NAME,
				"version": VERSION,
			},
		},
	}
}

// errorBody generates a Rollbar error body with a given stack trace.
func errorBody(err error, stack Stack) (map[string]interface{}, string) {
	fingerprint := stack.Fingerprint()
	errBody := map[string]interface{}{
		"trace": map[string]interface{}{
			"frames": stack,
			"exception": map[string]interface{}{
				"class":   errorClass(err),
				"message": err.Error(),
			},
		},
	}
	return errBody, fingerprint
}

// errorRequest extracts details from a Request in a format that Rollbar
// accepts.
func errorRequest(r *http.Request) map[string]interface{} {
	cleanQuery := filterParams(r.URL.Query())

	return map[string]interface{}{
		"url":     r.URL.String(),
		"method":  r.Method,
		"headers": flattenValues(r.Header),

		// GET params
		"query_string": url.Values(cleanQuery).Encode(),
		"GET":          flattenValues(cleanQuery),

		// POST / PUT params
		"POST":    flattenValues(filterParams(r.Form)),
		"user_ip": r.RemoteAddr,
	}
}

// filterParams filters sensitive information like passwords from being sent to
// Rollbar.
func filterParams(values map[string][]string) map[string][]string {
	for key, _ := range values {
		if FilterFields.Match([]byte(key)) {
			values[key] = []string{FILTERED}
		}
	}

	return values
}

func flattenValues(values map[string][]string) map[string]interface{} {
	result := make(map[string]interface{})

	for k, v := range values {
		if len(v) == 1 {
			result[k] = v[0]
		} else {
			result[k] = v
		}
	}

	return result
}

// Build a message inner-body for the given message string.
func messageBody(s string) map[string]interface{} {
	return map[string]interface{}{
		"message": map[string]interface{}{
			"body": s,
		},
	}
}

func errorClass(err error) string {
	class := reflect.TypeOf(err).String()
	if class == "" {
		return "panic"
	} else if class == "*errors.errorString" {
		checksum := adler32.Checksum([]byte(err.Error()))
		return fmt.Sprintf("{%x}", checksum)
	} else {
		return strings.TrimPrefix(class, "*")
	}
}

// -- POST handling

// Queue the given JSON body to be POSTed to Rollbar.
func push(body map[string]interface{}) {
	if len(bodyChannel) < Buffer {
		waitGroup.Add(1)
		bodyChannel <- body
	} else {
		stderr("buffer full, dropping error on the floor")
	}
}

// POST the given JSON body to Rollbar synchronously.
func post(body map[string]interface{}) {
	if len(Token) == 0 {
		stderr("empty token")
		return
	}

	jsonBody, err := json.Marshal(body)
	if err != nil {
		stderr("failed to encode payload: %s", err.Error())
		return
	}

	resp, err := http.Post(Endpoint, "application/json", bytes.NewReader(jsonBody))
	if err != nil {
		stderr("POST failed: %s", err.Error())
	} else if resp.StatusCode != 200 {
		stderr("received response: %s", resp.Status)
	}
	if resp != nil {
		resp.Body.Close()
	}
}

// -- stderr
func stderr(format string, args ...interface{}) {
	if ErrorWriter != nil {
		format = "Rollbar error: " + format + "\n"
		fmt.Fprintf(ErrorWriter, format, args...)
	}
}
<|endoftext|>