<commit_before>package resty

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"testing"
)

var (
	LogRequests = false
)

type TestRequest struct {
	Method string
	Path   string
	Data   interface{}

	ExpectedStatus int
	ExpectedData   interface{}
}

func (tr *TestRequest) String() string {
	return tr.Method + " " + tr.Path
}

func (tr *TestRequest) Run(t *testing.T, c *Client) {
	r := c.Do(tr.Method, tr.Path, tr.Data, nil)
	if LogRequests {
		t.Logf("%s: %s", tr.String(), r.Value)
	}

	switch {
	case r.Err != nil:
		t.Fatalf("%s: error: %v, status: %d, resp: %s", tr.String(), r.Err, r.Status, r.Value)
	case tr.ExpectedStatus == 0 && r.Status != 200, r.Status != tr.ExpectedStatus:
		t.Fatalf("%s: wanted %d, got %d: %s", tr.String(), tr.ExpectedStatus, r.Status, r.Value)
	case tr.ExpectedData != nil:
		if err := compareRes(r.Value, getVal(tr.ExpectedData)); err != nil {
			t.Fatalf("%s: %v", tr.String(), err)
		}
	}
}

// a == result, b == expected
func compareRes(a, b []byte) error {
	var am, bm map[string]interface{}
	if err := json.Unmarshal(a, &am); err != nil {
		return fmt.Errorf("%s: %v", a, err)
	}
	if err := json.Unmarshal(b, &bm); err != nil {
		return fmt.Errorf("%s: %v", b, err)
	}
	if len(am) == len(bm) {
		if !reflect.DeepEqual(am, bm) {
			return fmt.Errorf("%s != %s", a, b)
		}
	} else {
		for k, v := range bm {
			if ov := am[k]; !reflect.DeepEqual(v, ov) {
				return fmt.Errorf("wanted %v, got %v", v, ov)
			}
		}
	}
	return nil
}

func getVal(v interface{}) []byte {
	switch v := v.(type) {
	case []byte:
		return v
	case string:
		return []byte(v)
	case io.Reader:
		b, _ := ioutil.ReadAll(v)
		return b
	case nil:
		return nil
	}
	j, _ := json.Marshal(v)
	return j
}
<commit_msg>simplify logic<commit_after>package resty

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"testing"
)

var (
	LogRequests = false
)

type TestRequest struct {
	Method string
	Path   string
	Data   interface{}

	ExpectedStatus int
	ExpectedData   interface{}
}

func (tr *TestRequest) String() string {
	return tr.Method + " " + tr.Path
}

func (tr *TestRequest) Run(t *testing.T, c *Client) {
	r := c.Do(tr.Method, tr.Path, tr.Data, nil)
	if LogRequests {
		t.Logf("%s: %s", tr.String(), r.Value)
	}

	switch {
	case r.Err != nil:
		t.Fatalf("%s: error: %v, status: %d, resp: %s", tr.String(), r.Err, r.Status, r.Value)
	case tr.ExpectedStatus == 0 && r.Status != 200, r.Status != tr.ExpectedStatus:
		t.Fatalf("%s: wanted %d, got %d: %s", tr.String(), tr.ExpectedStatus, r.Status, r.Value)
	case tr.ExpectedData != nil:
		if err := compareRes(r.Value, getVal(tr.ExpectedData)); err != nil {
			t.Fatalf("%s: %v", tr.String(), err)
		}
	}
}

// a == result, b == expected
func compareRes(a, b []byte) error {
	var am, bm map[string]interface{}
	if err := json.Unmarshal(a, &am); err != nil {
		return fmt.Errorf("%s: %v", a, err)
	}
	if err := json.Unmarshal(b, &bm); err != nil {
		return fmt.Errorf("%s: %v", b, err)
	}

	for k, v := range bm {
		if ov := am[k]; !reflect.DeepEqual(v, ov) {
			return fmt.Errorf("wanted %v, got %v", v, ov)
		}
	}
	return nil
}

func getVal(v interface{}) []byte {
	switch v := v.(type) {
	case []byte:
		return v
	case string:
		return []byte(v)
	case io.Reader:
		b, _ := ioutil.ReadAll(v)
		return b
	case nil:
		return nil
	}
	j, _ := json.Marshal(v)
	return j
}
<|endoftext|>


<commit_before>package config_test

import (
	"testing"

	clusterConfig "github.com/lxc/lxd/lxd/cluster/config"
	"github.com/lxc/lxd/lxd/db"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// The server configuration is initially empty.
func TestConfigLoad_Initial(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)

	require.NoError(t, err)
	assert.Equal(t, map[string]any{}, config.Dump())

	assert.Equal(t, float64(20), config.OfflineThreshold().Seconds())
}

// If the database contains invalid keys, they are ignored.
func TestConfigLoad_IgnoreInvalidKeys(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	err := tx.UpdateClusterConfig(map[string]string{
		"foo":             "garbage",
		"core.proxy_http": "foo.bar",
	})
	require.NoError(t, err)

	config, err := clusterConfig.Load(tx)

	require.NoError(t, err)
	values := map[string]any{"core.proxy_http": "foo.bar"}
	assert.Equal(t, values, config.Dump())
}

// Triggers can be specified to execute custom code on config key changes.
func TestConfigLoad_Triggers(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)

	require.NoError(t, err)
	assert.Equal(t, map[string]any{}, config.Dump())
}

// Offline threshold must be greater than the heartbeat interval.
func TestConfigLoad_OfflineThresholdValidator(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	_, err = config.Patch(map[string]any{"cluster.offline_threshold": "2"})
	require.EqualError(t, err, "cannot set 'cluster.offline_threshold' to '2': Value must be greater than '10'")

}

// Max number of voters must be odd.
func TestConfigLoad_MaxVotersValidator(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	_, err = config.Patch(map[string]any{"cluster.max_voters": "4"})
	require.EqualError(t, err, "cannot set 'cluster.max_voters' to '4': Value must be an odd number equal to or higher than 3")

}

// If some previously set values are missing from the ones passed to Replace(),
// they are deleted from the configuration.
func TestConfig_ReplaceDeleteValues(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	changed, err := config.Replace(map[string]any{"core.proxy_http": "foo.bar"})
	assert.NoError(t, err)
	assert.Equal(t, map[string]string{"core.proxy_http": "foo.bar"}, changed)

	_, err = config.Replace(map[string]any{})
	assert.NoError(t, err)

	assert.Equal(t, "", config.ProxyHTTP())

	values, err := tx.Config()
	require.NoError(t, err)
	assert.Equal(t, map[string]string{}, values)
}

// If some previously set values are missing from the ones passed to Patch(),
// they are kept as they are.
func TestConfig_PatchKeepsValues(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	_, err = config.Replace(map[string]any{"core.proxy_http": "foo.bar"})
	assert.NoError(t, err)

	_, err = config.Patch(map[string]any{})
	assert.NoError(t, err)

	assert.Equal(t, "foo.bar", config.ProxyHTTP())

	values, err := tx.Config()
	require.NoError(t, err)
	assert.Equal(t, map[string]string{"core.proxy_http": "foo.bar"}, values)
}
<commit_msg>lxd/cluster/config: Removes unnecessary whitespace.<commit_after>package config_test

import (
	"testing"

	clusterConfig "github.com/lxc/lxd/lxd/cluster/config"
	"github.com/lxc/lxd/lxd/db"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// The server configuration is initially empty.
func TestConfigLoad_Initial(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)

	require.NoError(t, err)
	assert.Equal(t, map[string]any{}, config.Dump())

	assert.Equal(t, float64(20), config.OfflineThreshold().Seconds())
}

// If the database contains invalid keys, they are ignored.
func TestConfigLoad_IgnoreInvalidKeys(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	err := tx.UpdateClusterConfig(map[string]string{
		"foo":             "garbage",
		"core.proxy_http": "foo.bar",
	})
	require.NoError(t, err)

	config, err := clusterConfig.Load(tx)

	require.NoError(t, err)
	values := map[string]any{"core.proxy_http": "foo.bar"}
	assert.Equal(t, values, config.Dump())
}

// Triggers can be specified to execute custom code on config key changes.
func TestConfigLoad_Triggers(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)

	require.NoError(t, err)
	assert.Equal(t, map[string]any{}, config.Dump())
}

// Offline threshold must be greater than the heartbeat interval.
func TestConfigLoad_OfflineThresholdValidator(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	_, err = config.Patch(map[string]any{"cluster.offline_threshold": "2"})
	require.EqualError(t, err, "cannot set 'cluster.offline_threshold' to '2': Value must be greater than '10'")
}

// Max number of voters must be odd.
func TestConfigLoad_MaxVotersValidator(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	_, err = config.Patch(map[string]any{"cluster.max_voters": "4"})
	require.EqualError(t, err, "cannot set 'cluster.max_voters' to '4': Value must be an odd number equal to or higher than 3")
}

// If some previously set values are missing from the ones passed to Replace(),
// they are deleted from the configuration.
func TestConfig_ReplaceDeleteValues(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	changed, err := config.Replace(map[string]any{"core.proxy_http": "foo.bar"})
	assert.NoError(t, err)
	assert.Equal(t, map[string]string{"core.proxy_http": "foo.bar"}, changed)

	_, err = config.Replace(map[string]any{})
	assert.NoError(t, err)

	assert.Equal(t, "", config.ProxyHTTP())

	values, err := tx.Config()
	require.NoError(t, err)
	assert.Equal(t, map[string]string{}, values)
}

// If some previously set values are missing from the ones passed to Patch(),
// they are kept as they are.
func TestConfig_PatchKeepsValues(t *testing.T) {
	tx, cleanup := db.NewTestClusterTx(t)
	defer cleanup()

	config, err := clusterConfig.Load(tx)
	require.NoError(t, err)

	_, err = config.Replace(map[string]any{"core.proxy_http": "foo.bar"})
	assert.NoError(t, err)

	_, err = config.Patch(map[string]any{})
	assert.NoError(t, err)

	assert.Equal(t, "foo.bar", config.ProxyHTTP())

	values, err := tx.Config()
	require.NoError(t, err)
	assert.Equal(t, map[string]string{"core.proxy_http": "foo.bar"}, values)
}
<|endoftext|>


<commit_before>package drivers

import (
	"fmt"
	"path/filepath"
	"strings"

	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

type dir struct {
	common
}

// Info returns info about the driver and its environment.
func (d *dir) Info() Info {
	return Info{
		Name:                  "dir",
		Version:               "1",
		OptimizedImages:       false,
		PreservesInodes:       false,
		Remote:                false,
		VolumeTypes:           []VolumeType{VolumeTypeCustom, VolumeTypeImage, VolumeTypeContainer, VolumeTypeVM},
		BlockBacking:          false,
		RunningQuotaResize:    true,
		RunningSnapshotFreeze: true,
	}
}

// Create is called during pool creation and is effectively using an empty driver struct.
// WARNING: The Create() function cannot rely on any of the struct attributes being set.
func (d *dir) Create() error {
	// Set default source if missing.
	if d.config["source"] == "" {
		d.config["source"] = GetPoolMountPath(d.name)
	}

	if !shared.PathExists(d.config["source"]) {
		return fmt.Errorf("Source path '%s' doesn't exist", d.config["source"])
	}

	// Check that if within LXD_DIR, we're at our expected spot.
	cleanSource := filepath.Clean(d.config["source"])
	if strings.HasPrefix(cleanSource, shared.VarPath()) && cleanSource != GetPoolMountPath(d.name) {
		return fmt.Errorf("Source path '%s' is within the LXD directory", d.config["source"])
	}

	// Check that the path is currently empty.
	isEmpty, err := shared.PathIsEmpty(d.config["source"])
	if err != nil {
		return err
	}

	if !isEmpty {
		return fmt.Errorf("Source path '%s' isn't empty", d.config["source"])
	}

	return nil
}

// Delete removes the storage pool from the storage device.
func (d *dir) Delete(op *operations.Operation) error {
	// On delete, wipe everything in the directory.
	err := wipeDirectory(GetPoolMountPath(d.name))
	if err != nil {
		return err
	}

	// Unmount the path.
	_, err = d.Unmount()
	if err != nil {
		return err
	}

	return nil
}

// Validate checks that all provide keys are supported and that no conflicting or missing configuration is present.
func (d *dir) Validate(config map[string]string) error {
	return nil
}

// Update applies any driver changes required from a configuration change.
func (d *dir) Update(changedConfig map[string]string) error {
	return nil
}

// Mount mounts the storage pool.
func (d *dir) Mount() (bool, error) {
	path := GetPoolMountPath(d.name)

	// Check if we're dealing with an external mount.
	if d.config["source"] == path {
		return false, nil
	}

	// Check if already mounted.
	if sameMount(d.config["source"], path) {
		return false, nil
	}

	// Setup the bind-mount.
	err := tryMount(d.config["source"], path, "none", unix.MS_BIND, "")
	if err != nil {
		return false, err
	}

	return true, nil
}

// Unmount unmounts the storage pool.
func (d *dir) Unmount() (bool, error) {
	path := GetPoolMountPath(d.name)

	// Check if we're dealing with an external mount.
	if d.config["source"] == path {
		return false, nil
	}

	// Unmount until nothing is left mounted.
	return forceUnmount(path)
}

// GetResources returns the pool resource usage information.
func (d *dir) GetResources() (*api.ResourcesStoragePool, error) {
	return d.vfsGetResources()
}
<commit_msg>lxd/storage/drivers/driver/dir: TryMount usage<commit_after>package drivers

import (
	"fmt"
	"path/filepath"
	"strings"

	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

type dir struct {
	common
}

// Info returns info about the driver and its environment.
func (d *dir) Info() Info {
	return Info{
		Name:                  "dir",
		Version:               "1",
		OptimizedImages:       false,
		PreservesInodes:       false,
		Remote:                false,
		VolumeTypes:           []VolumeType{VolumeTypeCustom, VolumeTypeImage, VolumeTypeContainer, VolumeTypeVM},
		BlockBacking:          false,
		RunningQuotaResize:    true,
		RunningSnapshotFreeze: true,
	}
}

// Create is called during pool creation and is effectively using an empty driver struct.
// WARNING: The Create() function cannot rely on any of the struct attributes being set.
func (d *dir) Create() error {
	// Set default source if missing.
	if d.config["source"] == "" {
		d.config["source"] = GetPoolMountPath(d.name)
	}

	if !shared.PathExists(d.config["source"]) {
		return fmt.Errorf("Source path '%s' doesn't exist", d.config["source"])
	}

	// Check that if within LXD_DIR, we're at our expected spot.
	cleanSource := filepath.Clean(d.config["source"])
	if strings.HasPrefix(cleanSource, shared.VarPath()) && cleanSource != GetPoolMountPath(d.name) {
		return fmt.Errorf("Source path '%s' is within the LXD directory", d.config["source"])
	}

	// Check that the path is currently empty.
	isEmpty, err := shared.PathIsEmpty(d.config["source"])
	if err != nil {
		return err
	}

	if !isEmpty {
		return fmt.Errorf("Source path '%s' isn't empty", d.config["source"])
	}

	return nil
}

// Delete removes the storage pool from the storage device.
func (d *dir) Delete(op *operations.Operation) error {
	// On delete, wipe everything in the directory.
	err := wipeDirectory(GetPoolMountPath(d.name))
	if err != nil {
		return err
	}

	// Unmount the path.
	_, err = d.Unmount()
	if err != nil {
		return err
	}

	return nil
}

// Validate checks that all provide keys are supported and that no conflicting or missing configuration is present.
func (d *dir) Validate(config map[string]string) error {
	return nil
}

// Update applies any driver changes required from a configuration change.
func (d *dir) Update(changedConfig map[string]string) error {
	return nil
}

// Mount mounts the storage pool.
func (d *dir) Mount() (bool, error) {
	path := GetPoolMountPath(d.name)

	// Check if we're dealing with an external mount.
	if d.config["source"] == path {
		return false, nil
	}

	// Check if already mounted.
	if sameMount(d.config["source"], path) {
		return false, nil
	}

	// Setup the bind-mount.
	err := TryMount(d.config["source"], path, "none", unix.MS_BIND, "")
	if err != nil {
		return false, err
	}

	return true, nil
}

// Unmount unmounts the storage pool.
func (d *dir) Unmount() (bool, error) {
	path := GetPoolMountPath(d.name)

	// Check if we're dealing with an external mount.
	if d.config["source"] == path {
		return false, nil
	}

	// Unmount until nothing is left mounted.
	return forceUnmount(path)
}

// GetResources returns the pool resource usage information.
func (d *dir) GetResources() (*api.ResourcesStoragePool, error) {
	return d.vfsGetResources()
}
<|endoftext|>


<commit_before>package cache

import (
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"testing"
	"time"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/test"
	"github.com/raintank/schema"
)

func generateChunks(b testing.TB, startAt, count, step uint32) []chunk.IterGen {
	res := make([]chunk.IterGen, 0, count)

	values := make([]uint32, step)
	for t0 := startAt; t0 < startAt+(step*uint32(count)); t0 += step {
		c := getItgen(b, values, t0, true)
		res = append(res, c)
	}
	return res
}

func getCCM() (schema.AMKey, *CCacheMetric) {
	amkey, _ := schema.AMKeyFromString("1.12345678901234567890123456789012")
	ccm := NewCCacheMetric(amkey.MKey)
	return amkey, ccm
}

// TestAddAsc tests adding ascending timestamp chunks individually
func TestAddAsc(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(1)
		for _, chunk := range chunks {
			ccm.Add(prev, chunk)
			prev = chunk.Ts
		}
	})
}

// TestAddDesc1 tests adding chunks that are all descending
func TestAddDesc1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		for i := len(chunks) - 1; i >= 0; i-- {
			ccm.Add(0, chunks[i])
		}
	})
}

// TestAddDesc4 tests adding chunks that are globally descending
// but in groups 4 are ascending
func TestAddDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.Add(0, chunks[2])
		ccm.Add(0, chunks[3])
		ccm.Add(0, chunks[4])
		ccm.Add(0, chunks[5])
		ccm.Add(0, chunks[0])
		ccm.Add(0, chunks[1])
	})
}

// TestAddRange tests adding a contiguous range at once
func TestAddRange(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(10)
		ccm.AddRange(prev, chunks)
	})
}

// TestAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups of 4 are ascending. those groups are added via 1 AddRange.
func TestAddRangeDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[2:6])
		ccm.AddRange(0, chunks[0:2])
	})
}

// test executes the run function
// run should generate chunks and add them to the CCacheMetric however it likes,
// but in a way so that the result will be as expected
func testRun(t *testing.T, run func(*CCacheMetric)) {
	amkey, ccm := getCCM()

	run(ccm)

	res := CCSearchResult{}
	ccm.Search(test.NewContext(), amkey, &res, 25, 45)
	if res.Complete != true {
		t.Fatalf("Expected result to be complete, but it was not")
	}

	if res.Start[0].Ts != 20 {
		t.Fatalf("Expected result to start at 20, but had %d", res.Start[0].Ts)
	}

	if res.Start[len(res.Start)-1].Ts != 40 {
		t.Fatalf("Expected result to start at 40, but had %d", res.Start[len(res.Start)-1].Ts)
	}
}

// BenchmarkAddAsc benchmarks adding ascending timestamp chunks individually
func BenchmarkAddAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	for _, chunk := range chunks {
		ccm.Add(prev, chunk)
		prev = chunk.Ts
	}
}

// BenchmarkAddDesc1 benchmarks adding chunks that are all descending
func BenchmarkAddDesc1(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 1; i >= 0; i-- {
		ccm.Add(0, chunks[i])
	}
}

// BenchmarkAddDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending
func BenchmarkAddDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.Add(0, chunks[i])
		ccm.Add(0, chunks[i+1])
		ccm.Add(0, chunks[i+2])
		ccm.Add(0, chunks[i+3])
	}
}

// BenchmarkAddDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending
func BenchmarkAddDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		for offset := 0; offset < 64; offset += 1 {
			ccm.Add(0, chunks[i+offset])
		}
	}

}

// BenchmarkAddRangeAsc benchmarks adding a contiguous range at once
func BenchmarkAddRangeAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	ccm.AddRange(prev, chunks)
}

// BenchmarkAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.AddRange(0, chunks[i:i+4])
	}
}

// BenchmarkAddRangeDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		ccm.AddRange(0, chunks[i:i+64])
	}
}

func TestCorruptionCase1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[3:6])
		ccm.AddRange(0, chunks[0:4])
		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	})
}

func getRandomNumber(min, max int) int {
	return rand.Intn(max-min) + min
}

// getRandomRange returns a range start-end so that
// end >= start and both numbers drawn from [min, max)
func getRandomRange(min, max int) (int, int) {
	start := getRandomNumber(min, max)
	end := getRandomNumber(start, max)
	return start, end
}

func TestCorruptionCase2(t *testing.T) {
	rand.Seed(time.Now().Unix())
	_, ccm := getCCM()
	iterations := 100000

	// 100 chunks, first t0=t10, last is t0=1000
	chunks := generateChunks(t, 10, 100, 10)

	var opAdd, opAddRange, opDel, opDelRange int
	var adds, dels int

	for i := 0; i < iterations; i++ {
		// 0 = Add
		// 1 = AddRange
		// 2 = Del
		// 2 = Del range (via multi del cals)
		action := getRandomNumber(0, 4)
		switch action {
		case 0:
			chunk := getRandomNumber(0, 100)
			t.Logf("adding chunk %d", chunk)
			ccm.Add(0, chunks[chunk])
			opAdd++
			adds++
		case 1:
			from, to := getRandomRange(0, 100)
			t.Logf("adding range %d-%d", from, to)
			ccm.AddRange(0, chunks[from:to])
			adds += (to - from)
			opAddRange++
		case 2:
			chunk := getRandomNumber(0, 100)
			t.Logf("deleting chunk %d", chunk)
			ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			opDel++
			dels++
		case 3:
			from, to := getRandomRange(0, 100)
			for chunk := from; chunk < to; chunk++ {
				t.Logf("deleting chunk %d", chunk)
				ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			}
			opDelRange++
			dels += (to - from)
		}

		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	}

	fmt.Printf("operations: add %d - addRange %d - del %d - delRange %d\n", opAdd, opAddRange, opDel, opDelRange)
	fmt.Printf("total chunk adds %d - total chunk deletes %d\n", adds, dels)
}

// verifyCcm verifies the integrity of a CCacheMetric
// it assumes that all itergens are span-aware
func verifyCcm(ccm *CCacheMetric) error {
	var chunk *CCacheChunk
	var ok bool

	if len(ccm.chunks) != len(ccm.keys) {
		return errors.New("Length of ccm.chunks does not match ccm.keys")
	}

	if !sort.IsSorted(accnt.Uint32Asc(ccm.keys)) {
		return errors.New("keys are not sorted")
	}

	for i, ts := range ccm.keys {
		if chunk, ok = ccm.chunks[ts]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ts)
		}

		if i == 0 {
			if chunk.Prev != 0 {
				return errors.New("First chunk has Prev != 0")
			}
		} else {
			if chunk.Prev == 0 {
				if ccm.chunks[ccm.keys[i-1]].Ts == chunk.Ts-chunk.Itgen.Span {
					return fmt.Errorf("Chunk of ts %d has Prev == 0, but the previous chunk is present", ts)
				}
			} else {
				if ccm.chunks[ccm.keys[i-1]].Ts != chunk.Prev {
					return fmt.Errorf("Chunk of ts %d has Prev set to wrong ts %d but should be %d", ts, chunk.Prev, ccm.chunks[ccm.keys[i-1]].Ts)
				}
			}
		}

		if i == len(ccm.keys)-1 {
			if chunk.Next != 0 {
				return fmt.Errorf("Next of last chunk should be 0, but it's %d", chunk.Next)
			}

			// all checks completed
			break
		}

		var nextChunk *CCacheChunk
		if nextChunk, ok = ccm.chunks[ccm.keys[i+1]]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ccm.keys[i+1])
		}

		if chunk.Next == 0 {
			if chunk.Ts+chunk.Itgen.Span == nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to 0, but the next chunk is present", ts)
			}
		} else {
			if chunk.Next != nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to %d, but it should be %d", ts, chunk.Next, nextChunk.Ts)
			}
		}
	}
	return nil
}
<commit_msg>minor improvements<commit_after>package cache

import (
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"testing"
	"time"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/test"
	"github.com/raintank/schema"
)

func generateChunks(b testing.TB, startAt, count, step uint32) []chunk.IterGen {
	res := make([]chunk.IterGen, 0, count)

	values := make([]uint32, step)
	for t0 := startAt; t0 < startAt+(step*uint32(count)); t0 += step {
		c := getItgen(b, values, t0, true)
		res = append(res, c)
	}
	return res
}

func getCCM() (schema.AMKey, *CCacheMetric) {
	amkey, _ := schema.AMKeyFromString("1.12345678901234567890123456789012")
	ccm := NewCCacheMetric(amkey.MKey)
	return amkey, ccm
}

// TestAddAsc tests adding ascending timestamp chunks individually
func TestAddAsc(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(1)
		for _, chunk := range chunks {
			ccm.Add(prev, chunk)
			prev = chunk.Ts
		}
	})
}

// TestAddDesc1 tests adding chunks that are all descending
func TestAddDesc1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		for i := len(chunks) - 1; i >= 0; i-- {
			ccm.Add(0, chunks[i])
		}
	})
}

// TestAddDesc4 tests adding chunks that are globally descending
// but in groups 4 are ascending
func TestAddDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.Add(0, chunks[2])
		ccm.Add(0, chunks[3])
		ccm.Add(0, chunks[4])
		ccm.Add(0, chunks[5])
		ccm.Add(0, chunks[0])
		ccm.Add(0, chunks[1])
	})
}

// TestAddRange tests adding a contiguous range at once
func TestAddRange(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(10)
		ccm.AddRange(prev, chunks)
	})
}

// TestAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups of 4 are ascending. those groups are added via 1 AddRange.
func TestAddRangeDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[2:6])
		ccm.AddRange(0, chunks[0:2])
	})
}

// test executes the run function
// run should generate chunks and add them to the CCacheMetric however it likes,
// but in a way so that the result will be as expected
func testRun(t *testing.T, run func(*CCacheMetric)) {
	amkey, ccm := getCCM()

	run(ccm)

	res := CCSearchResult{}
	ccm.Search(test.NewContext(), amkey, &res, 25, 45)
	if res.Complete != true {
		t.Fatalf("Expected result to be complete, but it was not")
	}

	if res.Start[0].Ts != 20 {
		t.Fatalf("Expected result to start at 20, but had %d", res.Start[0].Ts)
	}

	if res.Start[len(res.Start)-1].Ts != 40 {
		t.Fatalf("Expected result to start at 40, but had %d", res.Start[len(res.Start)-1].Ts)
	}
}

// BenchmarkAddAsc benchmarks adding ascending timestamp chunks individually
func BenchmarkAddAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	for _, chunk := range chunks {
		ccm.Add(prev, chunk)
		prev = chunk.Ts
	}
}

// BenchmarkAddDesc1 benchmarks adding chunks that are all descending
func BenchmarkAddDesc1(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 1; i >= 0; i-- {
		ccm.Add(0, chunks[i])
	}
}

// BenchmarkAddDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending
func BenchmarkAddDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.Add(0, chunks[i])
		ccm.Add(0, chunks[i+1])
		ccm.Add(0, chunks[i+2])
		ccm.Add(0, chunks[i+3])
	}
}

// BenchmarkAddDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending
func BenchmarkAddDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		for offset := 0; offset < 64; offset += 1 {
			ccm.Add(0, chunks[i+offset])
		}
	}

}

// BenchmarkAddRangeAsc benchmarks adding a contiguous range at once
func BenchmarkAddRangeAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	ccm.AddRange(prev, chunks)
}

// BenchmarkAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.AddRange(0, chunks[i:i+4])
	}
}

// BenchmarkAddRangeDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		ccm.AddRange(0, chunks[i:i+64])
	}
}

func TestCorruptionCase1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[3:6])
		ccm.AddRange(0, chunks[0:4])
		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	})
}

func getRandomNumber(min, max int) int {
	return rand.Intn(max-min) + min
}

// getRandomRange returns a range start-end so that
// end >= start and both numbers drawn from [min, max)
func getRandomRange(min, max int) (int, int) {
	number1 := getRandomNumber(min, max)
	number2 := getRandomNumber(min, max)
	if number1 > number2 {
		return number2, number1
	} else {
		return number1, number2
	}
}

func TestCorruptionCase2(t *testing.T) {
	rand.Seed(time.Now().Unix())
	_, ccm := getCCM()
	iterations := 100000

	// 100 chunks, first t0=t10, last is t0=1000
	chunks := generateChunks(t, 10, 100, 10)

	var opAdd, opAddRange, opDel, opDelRange int
	var adds, dels int

	for i := 0; i < iterations; i++ {
		// 0 = Add
		// 1 = AddRange
		// 2 = Del
		// 3 = Del range (via multi del cals)
		action := getRandomNumber(0, 4)
		switch action {
		case 0:
			chunk := getRandomNumber(0, 100)
			t.Logf("adding chunk %d", chunk)
			ccm.Add(0, chunks[chunk])
			opAdd++
			adds++
		case 1:
			from, to := getRandomRange(0, 100)
			t.Logf("adding range %d-%d", from, to)
			ccm.AddRange(0, chunks[from:to])
			adds += (to - from)
			opAddRange++
		case 2:
			chunk := getRandomNumber(0, 100)
			t.Logf("deleting chunk %d", chunk)
			ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			opDel++
			dels++
		case 3:
			from, to := getRandomRange(0, 100)
			t.Logf("deleting range %d-%d", from, to)
			for chunk := from; chunk < to; chunk++ {
				ccm.Del(chunks[chunk].Ts) // note: chunk may not exist
			}
			opDelRange++
			dels += (to - from)
		}

		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	}

	fmt.Printf("operations: add %d - addRange %d - del %d - delRange %d\n", opAdd, opAddRange, opDel, opDelRange)
	fmt.Printf("total chunk adds %d - total chunk deletes %d\n", adds, dels)
}

// verifyCcm verifies the integrity of a CCacheMetric
// it assumes that all itergens are span-aware
func verifyCcm(ccm *CCacheMetric) error {
	var chunk *CCacheChunk
	var ok bool

	if len(ccm.chunks) != len(ccm.keys) {
		return errors.New("Length of ccm.chunks does not match ccm.keys")
	}

	if !sort.IsSorted(accnt.Uint32Asc(ccm.keys)) {
		return errors.New("keys are not sorted")
	}

	for i, ts := range ccm.keys {
		if chunk, ok = ccm.chunks[ts]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ts)
		}

		if i == 0 {
			if chunk.Prev != 0 {
				return errors.New("First chunk has Prev != 0")
			}
		} else {
			if chunk.Prev == 0 {
				if ccm.chunks[ccm.keys[i-1]].Ts == chunk.Ts-chunk.Itgen.Span {
					return fmt.Errorf("Chunk of ts %d has Prev == 0, but the previous chunk is present", ts)
				}
			} else {
				if ccm.chunks[ccm.keys[i-1]].Ts != chunk.Prev {
					return fmt.Errorf("Chunk of ts %d has Prev set to wrong ts %d but should be %d", ts, chunk.Prev, ccm.chunks[ccm.keys[i-1]].Ts)
				}
			}
		}

		if i == len(ccm.keys)-1 {
			if chunk.Next != 0 {
				return fmt.Errorf("Next of last chunk should be 0, but it's %d", chunk.Next)
			}

			// all checks completed
			break
		}

		var nextChunk *CCacheChunk
		if nextChunk, ok = ccm.chunks[ccm.keys[i+1]]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ccm.keys[i+1])
		}

		if chunk.Next == 0 {
			if chunk.Ts+chunk.Itgen.Span == nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to 0, but the next chunk is present", ts)
			}
		} else {
			if chunk.Next != nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to %d, but it should be %d", ts, chunk.Next, nextChunk.Ts)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package query

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"

	"github.com/Jeffail/gabs/v2"
)

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

// ValueType represents a discrete value type supported by Bloblang queries.
type ValueType string

// ValueType variants.
var (
	ValueString  ValueType = "string"
	ValueBytes   ValueType = "bytes"
	ValueNumber  ValueType = "number"
	ValueBool    ValueType = "bool"
	ValueArray   ValueType = "array"
	ValueObject  ValueType = "object"
	ValueNull    ValueType = "null"
	ValueDelete  ValueType = "delete"
	ValueNothing ValueType = "nothing"
	ValueUnknown ValueType = "unknown"
)

// ITypeOf returns the type of a boxed value as a discrete ValueType. If the
// type of the value is unknown then ValueUnknown is returned.
func ITypeOf(i interface{}) ValueType {
	switch i.(type) {
	case string:
		return ValueString
	case []byte:
		return ValueBytes
	case int64, uint64, float64, json.Number:
		return ValueNumber
	case bool:
		return ValueBool
	case []interface{}:
		return ValueArray
	case map[string]interface{}:
		return ValueObject
	case Delete:
		return ValueDelete
	case Nothing:
		return ValueNothing
	case nil:
		return ValueNull
	}
	return ValueUnknown
}

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

// Delete is a special type that serializes to `null` when forced but indicates
// a target should be deleted.
type Delete *struct{}

// Nothing is a special type that serializes to `null` when forced but indicates
// a query should be disregarded (and not mapped).
type Nothing *struct{}

// IGetNumber takes a boxed value and attempts to extract a number (float64)
// from it.
func IGetNumber(v interface{}) (float64, error) {
	switch t := v.(type) {
	case int:
		return float64(t), nil
	case int64:
		return float64(t), nil
	case uint64:
		return float64(t), nil
	case float64:
		return t, nil
	case json.Number:
		return t.Float64()
	}
	return 0, NewTypeError(v, ValueNumber)
}

// IGetInt takes a boxed value and attempts to extract an integer (int64) from
// it.
func IGetInt(v interface{}) (int64, error) {
	switch t := v.(type) {
	case int:
		return int64(t), nil
	case int64:
		return t, nil
	case uint64:
		return int64(t), nil
	case float64:
		return int64(t), nil
	case json.Number:
		return t.Int64()
	}
	return 0, NewTypeError(v, ValueNumber)
}

// IGetBool takes a boxed value and attempts to extract a boolean from it.
func IGetBool(v interface{}) (bool, error) {
	switch t := v.(type) {
	case bool:
		return t, nil
	case int:
		return t != 0, nil
	case int64:
		return t != 0, nil
	case uint64:
		return t != 0, nil
	case float64:
		return t != 0, nil
	case json.Number:
		return t.String() != "0", nil
	}
	return false, NewTypeError(v, ValueBool)
}

// IGetString takes a boxed value and attempts to return a string value. Returns
// an error if the value is not a string or byte slice.
func IGetString(v interface{}) (string, error) {
	switch t := v.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	}
	return "", NewTypeError(v, ValueString)
}

// IIsNull returns whether a bloblang type is null, this includes Delete and
// Nothing types.
func IIsNull(i interface{}) bool {
	if i == nil {
		return true
	}
	switch i.(type) {
	case Delete, Nothing:
		return true
	}
	return false
}

// ISanitize takes a boxed value of any type and attempts to convert it into one
// of the following types: string, []byte, int64, uint64, float64, bool,
// []interface{}, map[string]interface{}, Delete, Nothing.
func ISanitize(i interface{}) interface{} {
	switch t := i.(type) {
	case string, []byte, int64, uint64, float64, bool, []interface{}, map[string]interface{}, Delete, Nothing:
		return i
	case json.RawMessage:
		return []byte(t)
	case json.Number:
		if i, err := t.Int64(); err == nil {
			return int64(i)
		}
		if f, err := t.Float64(); err == nil {
			return f
		}
		return t.String()
	case int:
		return int64(t)
	case int32:
		return int64(t)
	case uint32:
		return uint64(t)
	case uint:
		return uint64(t)
	case float32:
		return float64(t)
	}
	// Do NOT support unknown types (for now).
	return nil
}

// IToBytes takes a boxed value of any type and attempts to convert it into a
// byte slice.
func IToBytes(i interface{}) []byte {
	switch t := i.(type) {
	case string:
		return []byte(t)
	case []byte:
		return t
	case json.Number:
		return []byte(t.String())
	case int64, uint64, float64:
		return []byte(fmt.Sprintf("%v", t)) // TODO
	case bool:
		if t {
			return []byte("true")
		}
		return []byte("false")
	case nil:
		return []byte(`null`)
	}
	// Last resort
	return gabs.Wrap(i).Bytes()
}

// IToString takes a boxed value of any type and attempts to convert it into a
// string.
func IToString(i interface{}) string {
	switch t := i.(type) {
	case string:
		return t
	case []byte:
		return string(t)
	case int64, uint64, float64:
		return fmt.Sprintf("%v", t) // TODO
	case json.Number:
		return t.String()
	case bool:
		if t {
			return "true"
		}
		return "false"
	case nil:
		return `null`
	}
	// Last resort
	return gabs.Wrap(i).String()
}

// IToNumber takes a boxed value and attempts to extract a number (float64)
// from it or parse one.
func IToNumber(v interface{}) (float64, error) {
	switch t := v.(type) {
	case int:
		return float64(t), nil
	case int64:
		return float64(t), nil
	case uint64:
		return float64(t), nil
	case float64:
		return t, nil
	case json.Number:
		return t.Float64()
	case []byte:
		return strconv.ParseFloat(string(t), 64)
	case string:
		return strconv.ParseFloat(t, 64)
	}
	return 0, NewTypeError(v, ValueNumber)
}

const maxUint = ^uint64(0)
const maxInt = uint64(maxUint >> 1)

// IToInt takes a boxed value and attempts to extract a number (int64) from it
// or parse one.
func IToInt(v interface{}) (int64, error) {
	switch t := v.(type) {
	case int:
		return int64(t), nil
	case int64:
		return t, nil
	case uint64:
		if t > maxInt {
			return 0, errors.New("unsigned integer value is too large to be cast as a signed integer")
		}
		return int64(t), nil
	case float64:
		return int64(t), nil
	case json.Number:
		return t.Int64()
	case []byte:
		return strconv.ParseInt(string(t), 10, 64)
	case string:
		return strconv.ParseInt(t, 10, 64)
	}
	return 0, NewTypeError(v, ValueNumber)
}

// IToBool takes a boxed value and attempts to extract a boolean from it or
// parse it into a bool.
func IToBool(v interface{}) (bool, error) {
	switch t := v.(type) {
	case bool:
		return t, nil
	case int:
		return t != 0, nil
	case int64:
		return t != 0, nil
	case uint64:
		return t != 0, nil
	case float64:
		return t != 0, nil
	case json.Number:
		return t.String() != "0", nil
	case []byte:
		if v, err := strconv.ParseBool(string(t)); err == nil {
			return v, nil
		}
	case string:
		if v, err := strconv.ParseBool(t); err == nil {
			return v, nil
		}
	}
	return false, NewTypeError(v, ValueBool)
}

// IClone performs a deep copy of a generic value.
func IClone(root interface{}) interface{} {
	switch t := root.(type) {
	case map[string]interface{}:
		newMap := make(map[string]interface{}, len(t))
		for k, v := range t {
			newMap[k] = IClone(v)
		}
		return newMap
	case []interface{}:
		newSlice := make([]interface{}, len(t))
		for i, v := range t {
			newSlice[i] = IClone(v)
		}
		return newSlice
	}
	return root
}

//------------------------------------------------------------------------------
<commit_msg>Include int in typeof<commit_after>package query

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"

	"github.com/Jeffail/gabs/v2"
)

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

// ValueType represents a discrete value type supported by Bloblang queries.
type ValueType string

// ValueType variants.
var (
	ValueString  ValueType = "string"
	ValueBytes   ValueType = "bytes"
	ValueNumber  ValueType = "number"
	ValueBool    ValueType = "bool"
	ValueArray   ValueType = "array"
	ValueObject  ValueType = "object"
	ValueNull    ValueType = "null"
	ValueDelete  ValueType = "delete"
	ValueNothing ValueType = "nothing"
	ValueUnknown ValueType = "unknown"
)

// ITypeOf returns the type of a boxed value as a discrete ValueType. If the
// type of the value is unknown then ValueUnknown is returned.
func ITypeOf(i interface{}) ValueType {
	switch i.(type) {
	case string:
		return ValueString
	case []byte:
		return ValueBytes
	case int, int64, uint64, float64, json.Number:
		return ValueNumber
	case bool:
		return ValueBool
	case []interface{}:
		return ValueArray
	case map[string]interface{}:
		return ValueObject
	case Delete:
		return ValueDelete
	case Nothing:
		return ValueNothing
	case nil:
		return ValueNull
	}
	return ValueUnknown
}

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

// Delete is a special type that serializes to `null` when forced but indicates
// a target should be deleted.
type Delete *struct{}

// Nothing is a special type that serializes to `null` when forced but indicates
// a query should be disregarded (and not mapped).
type Nothing *struct{}

// IGetNumber takes a boxed value and attempts to extract a number (float64)
// from it.
func IGetNumber(v interface{}) (float64, error) {
	switch t := v.(type) {
	case int:
		return float64(t), nil
	case int64:
		return float64(t), nil
	case uint64:
		return float64(t), nil
	case float64:
		return t, nil
	case json.Number:
		return t.Float64()
	}
	return 0, NewTypeError(v, ValueNumber)
}

// IGetInt takes a boxed value and attempts to extract an integer (int64) from
// it.
func IGetInt(v interface{}) (int64, error) {
	switch t := v.(type) {
	case int:
		return int64(t), nil
	case int64:
		return t, nil
	case uint64:
		return int64(t), nil
	case float64:
		return int64(t), nil
	case json.Number:
		return t.Int64()
	}
	return 0, NewTypeError(v, ValueNumber)
}

// IGetBool takes a boxed value and attempts to extract a boolean from it.
func IGetBool(v interface{}) (bool, error) {
	switch t := v.(type) {
	case bool:
		return t, nil
	case int:
		return t != 0, nil
	case int64:
		return t != 0, nil
	case uint64:
		return t != 0, nil
	case float64:
		return t != 0, nil
	case json.Number:
		return t.String() != "0", nil
	}
	return false, NewTypeError(v, ValueBool)
}

// IGetString takes a boxed value and attempts to return a string value. Returns
// an error if the value is not a string or byte slice.
func IGetString(v interface{}) (string, error) {
	switch t := v.(type) {
	case string:
		return t, nil
	case []byte:
		return string(t), nil
	}
	return "", NewTypeError(v, ValueString)
}

// IIsNull returns whether a bloblang type is null, this includes Delete and
// Nothing types.
func IIsNull(i interface{}) bool {
	if i == nil {
		return true
	}
	switch i.(type) {
	case Delete, Nothing:
		return true
	}
	return false
}

// ISanitize takes a boxed value of any type and attempts to convert it into one
// of the following types: string, []byte, int64, uint64, float64, bool,
// []interface{}, map[string]interface{}, Delete, Nothing.
func ISanitize(i interface{}) interface{} {
	switch t := i.(type) {
	case string, []byte, int64, uint64, float64, bool, []interface{}, map[string]interface{}, Delete, Nothing:
		return i
	case json.RawMessage:
		return []byte(t)
	case json.Number:
		if i, err := t.Int64(); err == nil {
			return int64(i)
		}
		if f, err := t.Float64(); err == nil {
			return f
		}
		return t.String()
	case int:
		return int64(t)
	case int32:
		return int64(t)
	case uint32:
		return uint64(t)
	case uint:
		return uint64(t)
	case float32:
		return float64(t)
	}
	// Do NOT support unknown types (for now).
	return nil
}

// IToBytes takes a boxed value of any type and attempts to convert it into a
// byte slice.
func IToBytes(i interface{}) []byte {
	switch t := i.(type) {
	case string:
		return []byte(t)
	case []byte:
		return t
	case json.Number:
		return []byte(t.String())
	case int64, uint64, float64:
		return []byte(fmt.Sprintf("%v", t)) // TODO
	case bool:
		if t {
			return []byte("true")
		}
		return []byte("false")
	case nil:
		return []byte(`null`)
	}
	// Last resort
	return gabs.Wrap(i).Bytes()
}

// IToString takes a boxed value of any type and attempts to convert it into a
// string.
func IToString(i interface{}) string {
	switch t := i.(type) {
	case string:
		return t
	case []byte:
		return string(t)
	case int64, uint64, float64:
		return fmt.Sprintf("%v", t) // TODO
	case json.Number:
		return t.String()
	case bool:
		if t {
			return "true"
		}
		return "false"
	case nil:
		return `null`
	}
	// Last resort
	return gabs.Wrap(i).String()
}

// IToNumber takes a boxed value and attempts to extract a number (float64)
// from it or parse one.
func IToNumber(v interface{}) (float64, error) {
	switch t := v.(type) {
	case int:
		return float64(t), nil
	case int64:
		return float64(t), nil
	case uint64:
		return float64(t), nil
	case float64:
		return t, nil
	case json.Number:
		return t.Float64()
	case []byte:
		return strconv.ParseFloat(string(t), 64)
	case string:
		return strconv.ParseFloat(t, 64)
	}
	return 0, NewTypeError(v, ValueNumber)
}

const maxUint = ^uint64(0)
const maxInt = uint64(maxUint >> 1)

// IToInt takes a boxed value and attempts to extract a number (int64) from it
// or parse one.
func IToInt(v interface{}) (int64, error) {
	switch t := v.(type) {
	case int:
		return int64(t), nil
	case int64:
		return t, nil
	case uint64:
		if t > maxInt {
			return 0, errors.New("unsigned integer value is too large to be cast as a signed integer")
		}
		return int64(t), nil
	case float64:
		return int64(t), nil
	case json.Number:
		return t.Int64()
	case []byte:
		return strconv.ParseInt(string(t), 10, 64)
	case string:
		return strconv.ParseInt(t, 10, 64)
	}
	return 0, NewTypeError(v, ValueNumber)
}

// IToBool takes a boxed value and attempts to extract a boolean from it or
// parse it into a bool.
func IToBool(v interface{}) (bool, error) {
	switch t := v.(type) {
	case bool:
		return t, nil
	case int:
		return t != 0, nil
	case int64:
		return t != 0, nil
	case uint64:
		return t != 0, nil
	case float64:
		return t != 0, nil
	case json.Number:
		return t.String() != "0", nil
	case []byte:
		if v, err := strconv.ParseBool(string(t)); err == nil {
			return v, nil
		}
	case string:
		if v, err := strconv.ParseBool(t); err == nil {
			return v, nil
		}
	}
	return false, NewTypeError(v, ValueBool)
}

// IClone performs a deep copy of a generic value.
func IClone(root interface{}) interface{} {
	switch t := root.(type) {
	case map[string]interface{}:
		newMap := make(map[string]interface{}, len(t))
		for k, v := range t {
			newMap[k] = IClone(v)
		}
		return newMap
	case []interface{}:
		newSlice := make([]interface{}, len(t))
		for i, v := range t {
			newSlice[i] = IClone(v)
		}
		return newSlice
	}
	return root
}

//------------------------------------------------------------------------------
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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 generator

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"cloud.google.com/go/internal/aliasfix"
	"cloud.google.com/go/internal/aliasgen"
	"cloud.google.com/go/internal/gapicgen/execv"
	"cloud.google.com/go/internal/gapicgen/execv/gocmd"
	"cloud.google.com/go/internal/gapicgen/git"
	"golang.org/x/sync/errgroup"
)

var goPkgOptRe = regexp.MustCompile(`(?m)^option go_package = (.*);`)

// denylist is a set of clients to NOT generate.
var denylist = map[string]bool{
	// Temporarily stop generation of removed protos. Will be manually cleaned
	// up with: https://github.com/googleapis/google-cloud-go/issues/4098
	"google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1alpha2": true,

	// Not properly configured:
	"google.golang.org/genproto/googleapis/cloud/ondemandscanning/v1beta1": true,
	"google.golang.org/genproto/googleapis/cloud/ondemandscanning/v1":      true,
}

// noGRPC is the set of APIs that do not need gRPC stubs.
var noGRPC = map[string]bool{
	"google.golang.org/genproto/googleapis/cloud/compute/v1": true,
}

// GenprotoGenerator is used to generate code for googleapis/go-genproto.
type GenprotoGenerator struct {
	genprotoDir     string
	googleapisDir   string
	protoSrcDir     string
	googleCloudDir  string
	gapicToGenerate string
	forceAll        bool
	genAlias        bool
}

// NewGenprotoGenerator creates a new GenprotoGenerator.
func NewGenprotoGenerator(c *Config) *GenprotoGenerator {
	return &GenprotoGenerator{
		genprotoDir:     c.GenprotoDir,
		googleapisDir:   c.GoogleapisDir,
		protoSrcDir:     filepath.Join(c.ProtoDir, "/src"),
		googleCloudDir:  c.GapicDir,
		gapicToGenerate: c.GapicToGenerate,
		forceAll:        c.ForceAll,
		genAlias:        c.GenAlias,
	}
}

var skipPrefixes = []string{
	"google.golang.org/genproto/googleapis/ads",
	"google.golang.org/genproto/googleapis/storage",
}

func hasPrefix(s string, prefixes []string) bool {
	for _, prefix := range prefixes {
		if strings.HasPrefix(s, prefix) {
			return true
		}
	}
	return false
}

// Regen regenerates the genproto repository.
// regenGenproto regenerates the genproto repository.
//
// regenGenproto recursively walks through each directory named by given
// arguments, looking for all .proto files. (Symlinks are not followed.) Any
// proto file without `go_package` option or whose option does not begin with
// the genproto prefix is ignored.
//
// If multiple roots contain files with the same name, eg "root1/path/to/file"
// and "root2/path/to/file", only the first file is processed; the rest are
// ignored.
//
// Protoc is executed on remaining files, one invocation per set of files
// declaring the same Go package.
func (g *GenprotoGenerator) Regen(ctx context.Context) error {
	log.Println("regenerating genproto")

	if g.genAlias {
		return g.generateAliases()
	}

	// Create space to put generated .pb.go's.
	c := execv.Command("mkdir", "-p", "generated")
	c.Dir = g.genprotoDir
	if err := c.Run(); err != nil {
		return err
	}

	// Get the last processed googleapis hash.
	lastHash, err := ioutil.ReadFile(filepath.Join(g.genprotoDir, "regen.txt"))
	if err != nil {
		return err
	}

	// TODO(noahdietz): In local mode, since it clones a shallow copy with 1 commit,
	// if the last regenerated hash is earlier than the top commit, the git diff-tree
	// command fails. This is is a bit of a rough edge. Using my local clone of
	// googleapis rectified the issue.
	pkgFiles, err := g.getUpdatedPackages(string(lastHash))
	if err != nil {
		return err
	}
	if len(pkgFiles) == 0 {
		return errors.New("couldn't find any pkgfiles")
	}

	log.Println("generating from protos")
	grp, _ := errgroup.WithContext(ctx)
	for pkg, fileNames := range pkgFiles {
		if !strings.HasPrefix(pkg, "google.golang.org/genproto") || denylist[pkg] || hasPrefix(pkg, skipPrefixes) {
			continue
		}
		grpc := !noGRPC[pkg]
		pk := pkg
		fn := fileNames

		if !isMigrated(pkg) {
			grp.Go(func() error {
				log.Println("running protoc on", pk)
				return g.protoc(fn, grpc)
			})
		} else {
			log.Printf("skipping, %q has been migrated", pkg)
		}
	}
	if err := grp.Wait(); err != nil {
		return err
	}

	if err := g.moveAndCleanupGeneratedSrc(); err != nil {
		return err
	}

	if err := gocmd.Vet(g.genprotoDir); err != nil {
		return err
	}

	if err := gocmd.Build(g.genprotoDir); err != nil {
		return err
	}

	return nil
}

// goPkg reports the import path declared in the given file's `go_package`
// option. If the option is missing, goPkg returns empty string.
func goPkg(fileName string) (string, error) {
	content, err := ioutil.ReadFile(fileName)
	if err != nil {
		return "", err
	}

	var pkgName string
	if match := goPkgOptRe.FindSubmatch(content); len(match) > 0 {
		pn, err := strconv.Unquote(string(match[1]))
		if err != nil {
			return "", err
		}
		pkgName = pn
	}
	if p := strings.IndexRune(pkgName, ';'); p > 0 {
		pkgName = pkgName[:p]
	}
	return pkgName, nil
}

// protoc executes the "protoc" command on files named in fileNames, and outputs
// to "<genprotoDir>/generated".
func (g *GenprotoGenerator) protoc(fileNames []string, grpc bool) error {
	stubs := fmt.Sprintf("--go_out=%s/generated", g.genprotoDir)
	if grpc {
		stubs = fmt.Sprintf("--go_out=plugins=grpc:%s/generated", g.genprotoDir)
	}
	args := []string{"--experimental_allow_proto3_optional", stubs, "-I", g.googleapisDir, "-I", g.protoSrcDir}
	args = append(args, fileNames...)
	c := execv.Command("protoc", args...)
	c.Dir = g.genprotoDir
	return c.Run()
}

// getUpdatedPackages parses all of the new commits to find what packages need
// to be regenerated.
func (g *GenprotoGenerator) getUpdatedPackages(googleapisHash string) (map[string][]string, error) {
	if g.forceAll {
		return g.getAllPackages()
	}
	files, err := git.UpdateFilesSinceHash(g.googleapisDir, googleapisHash)
	if err != nil {
		return nil, err
	}
	pkgFiles := make(map[string][]string)
	for _, v := range files {
		if !strings.HasSuffix(v, ".proto") {
			continue
		}
		if strings.HasSuffix(v, "compute_small.proto") {
			continue
		}
		path := filepath.Join(g.googleapisDir, v)
		pkg, err := goPkg(path)
		if err != nil {
			return nil, err
		}
		pkgFiles[pkg] = append(pkgFiles[pkg], path)
	}
	return pkgFiles, nil
}

func (g *GenprotoGenerator) getAllPackages() (map[string][]string, error) {
	seenFiles := make(map[string]bool)
	pkgFiles := make(map[string][]string)
	for _, root := range []string{g.googleapisDir} {
		walkFn := func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.Mode().IsRegular() || !strings.HasSuffix(path, ".proto") {
				return nil
			}

			switch rel, err := filepath.Rel(root, path); {
			case err != nil:
				return err
			case seenFiles[rel]:
				return nil
			default:
				seenFiles[rel] = true
			}

			pkg, err := goPkg(path)
			if err != nil {
				return err
			}
			pkgFiles[pkg] = append(pkgFiles[pkg], path)
			return nil
		}
		if err := filepath.Walk(root, walkFn); err != nil {
			return nil, err
		}
	}
	return pkgFiles, nil
}

// moveAndCleanupGeneratedSrc moves all generated src to their correct locations
// in the repository, because protoc puts it in a folder called `generated/“.
func (g *GenprotoGenerator) moveAndCleanupGeneratedSrc() error {
	log.Println("moving generated code")
	// The period at the end is analogous to * (copy everything in this dir).
	c := execv.Command("cp", "-R", filepath.Join(g.genprotoDir, "generated", "google.golang.org", "genproto", "googleapis"), g.genprotoDir)
	if err := c.Run(); err != nil {
		return err
	}

	c = execv.Command("rm", "-rf", "generated")
	c.Dir = g.genprotoDir
	if err := c.Run(); err != nil {
		return err
	}

	return nil
}

func (g *GenprotoGenerator) generateAliases() error {
	for genprotoImport, newPkg := range aliasfix.GenprotoPkgMigration {
		if !isMigrated(genprotoImport) || g.gapicToGenerate == "" {
			continue
		}
		// remove the stubs dir segment from path
		gapicImport := newPkg.ImportPath[:strings.LastIndex(newPkg.ImportPath, "/")]
		if !strings.Contains(g.gapicToGenerate, gapicImport) {
			continue
		}
		srdDir := filepath.Join(g.googleCloudDir, strings.TrimPrefix(newPkg.ImportPath, "cloud.google.com/go/"))
		destDir := filepath.Join(g.genprotoDir, "googleapis", strings.TrimPrefix(genprotoImport, "google.golang.org/genproto/googleapis/"))
		if err := aliasgen.Run(srdDir, destDir); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>fix(internal/gapicgen): add slashes to be more strict with matching (#6672)<commit_after>// Copyright 2019 Google LLC
//
// 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 generator

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"cloud.google.com/go/internal/aliasfix"
	"cloud.google.com/go/internal/aliasgen"
	"cloud.google.com/go/internal/gapicgen/execv"
	"cloud.google.com/go/internal/gapicgen/execv/gocmd"
	"cloud.google.com/go/internal/gapicgen/git"
	"golang.org/x/sync/errgroup"
)

var goPkgOptRe = regexp.MustCompile(`(?m)^option go_package = (.*);`)

// denylist is a set of clients to NOT generate.
var denylist = map[string]bool{
	// Temporarily stop generation of removed protos. Will be manually cleaned
	// up with: https://github.com/googleapis/google-cloud-go/issues/4098
	"google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1alpha2": true,

	// Not properly configured:
	"google.golang.org/genproto/googleapis/cloud/ondemandscanning/v1beta1": true,
	"google.golang.org/genproto/googleapis/cloud/ondemandscanning/v1":      true,
}

// noGRPC is the set of APIs that do not need gRPC stubs.
var noGRPC = map[string]bool{
	"google.golang.org/genproto/googleapis/cloud/compute/v1": true,
}

// GenprotoGenerator is used to generate code for googleapis/go-genproto.
type GenprotoGenerator struct {
	genprotoDir     string
	googleapisDir   string
	protoSrcDir     string
	googleCloudDir  string
	gapicToGenerate string
	forceAll        bool
	genAlias        bool
}

// NewGenprotoGenerator creates a new GenprotoGenerator.
func NewGenprotoGenerator(c *Config) *GenprotoGenerator {
	return &GenprotoGenerator{
		genprotoDir:     c.GenprotoDir,
		googleapisDir:   c.GoogleapisDir,
		protoSrcDir:     filepath.Join(c.ProtoDir, "/src"),
		googleCloudDir:  c.GapicDir,
		gapicToGenerate: c.GapicToGenerate,
		forceAll:        c.ForceAll,
		genAlias:        c.GenAlias,
	}
}

var skipPrefixes = []string{
	"google.golang.org/genproto/googleapis/ads/",
	"google.golang.org/genproto/googleapis/storage/",
}

func hasPrefix(s string, prefixes []string) bool {
	for _, prefix := range prefixes {
		if strings.HasPrefix(s, prefix) {
			return true
		}
	}
	return false
}

// Regen regenerates the genproto repository.
// regenGenproto regenerates the genproto repository.
//
// regenGenproto recursively walks through each directory named by given
// arguments, looking for all .proto files. (Symlinks are not followed.) Any
// proto file without `go_package` option or whose option does not begin with
// the genproto prefix is ignored.
//
// If multiple roots contain files with the same name, eg "root1/path/to/file"
// and "root2/path/to/file", only the first file is processed; the rest are
// ignored.
//
// Protoc is executed on remaining files, one invocation per set of files
// declaring the same Go package.
func (g *GenprotoGenerator) Regen(ctx context.Context) error {
	log.Println("regenerating genproto")

	if g.genAlias {
		return g.generateAliases()
	}

	// Create space to put generated .pb.go's.
	c := execv.Command("mkdir", "-p", "generated")
	c.Dir = g.genprotoDir
	if err := c.Run(); err != nil {
		return err
	}

	// Get the last processed googleapis hash.
	lastHash, err := ioutil.ReadFile(filepath.Join(g.genprotoDir, "regen.txt"))
	if err != nil {
		return err
	}

	// TODO(noahdietz): In local mode, since it clones a shallow copy with 1 commit,
	// if the last regenerated hash is earlier than the top commit, the git diff-tree
	// command fails. This is is a bit of a rough edge. Using my local clone of
	// googleapis rectified the issue.
	pkgFiles, err := g.getUpdatedPackages(string(lastHash))
	if err != nil {
		return err
	}
	if len(pkgFiles) == 0 {
		return errors.New("couldn't find any pkgfiles")
	}

	log.Println("generating from protos")
	grp, _ := errgroup.WithContext(ctx)
	for pkg, fileNames := range pkgFiles {
		if !strings.HasPrefix(pkg, "google.golang.org/genproto") || denylist[pkg] || hasPrefix(pkg, skipPrefixes) {
			continue
		}
		grpc := !noGRPC[pkg]
		pk := pkg
		fn := fileNames

		if !isMigrated(pkg) {
			grp.Go(func() error {
				log.Println("running protoc on", pk)
				return g.protoc(fn, grpc)
			})
		} else {
			log.Printf("skipping, %q has been migrated", pkg)
		}
	}
	if err := grp.Wait(); err != nil {
		return err
	}

	if err := g.moveAndCleanupGeneratedSrc(); err != nil {
		return err
	}

	if err := gocmd.Vet(g.genprotoDir); err != nil {
		return err
	}

	if err := gocmd.Build(g.genprotoDir); err != nil {
		return err
	}

	return nil
}

// goPkg reports the import path declared in the given file's `go_package`
// option. If the option is missing, goPkg returns empty string.
func goPkg(fileName string) (string, error) {
	content, err := ioutil.ReadFile(fileName)
	if err != nil {
		return "", err
	}

	var pkgName string
	if match := goPkgOptRe.FindSubmatch(content); len(match) > 0 {
		pn, err := strconv.Unquote(string(match[1]))
		if err != nil {
			return "", err
		}
		pkgName = pn
	}
	if p := strings.IndexRune(pkgName, ';'); p > 0 {
		pkgName = pkgName[:p]
	}
	return pkgName, nil
}

// protoc executes the "protoc" command on files named in fileNames, and outputs
// to "<genprotoDir>/generated".
func (g *GenprotoGenerator) protoc(fileNames []string, grpc bool) error {
	stubs := fmt.Sprintf("--go_out=%s/generated", g.genprotoDir)
	if grpc {
		stubs = fmt.Sprintf("--go_out=plugins=grpc:%s/generated", g.genprotoDir)
	}
	args := []string{"--experimental_allow_proto3_optional", stubs, "-I", g.googleapisDir, "-I", g.protoSrcDir}
	args = append(args, fileNames...)
	c := execv.Command("protoc", args...)
	c.Dir = g.genprotoDir
	return c.Run()
}

// getUpdatedPackages parses all of the new commits to find what packages need
// to be regenerated.
func (g *GenprotoGenerator) getUpdatedPackages(googleapisHash string) (map[string][]string, error) {
	if g.forceAll {
		return g.getAllPackages()
	}
	files, err := git.UpdateFilesSinceHash(g.googleapisDir, googleapisHash)
	if err != nil {
		return nil, err
	}
	pkgFiles := make(map[string][]string)
	for _, v := range files {
		if !strings.HasSuffix(v, ".proto") {
			continue
		}
		if strings.HasSuffix(v, "compute_small.proto") {
			continue
		}
		path := filepath.Join(g.googleapisDir, v)
		pkg, err := goPkg(path)
		if err != nil {
			return nil, err
		}
		pkgFiles[pkg] = append(pkgFiles[pkg], path)
	}
	return pkgFiles, nil
}

func (g *GenprotoGenerator) getAllPackages() (map[string][]string, error) {
	seenFiles := make(map[string]bool)
	pkgFiles := make(map[string][]string)
	for _, root := range []string{g.googleapisDir} {
		walkFn := func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.Mode().IsRegular() || !strings.HasSuffix(path, ".proto") {
				return nil
			}

			switch rel, err := filepath.Rel(root, path); {
			case err != nil:
				return err
			case seenFiles[rel]:
				return nil
			default:
				seenFiles[rel] = true
			}

			pkg, err := goPkg(path)
			if err != nil {
				return err
			}
			pkgFiles[pkg] = append(pkgFiles[pkg], path)
			return nil
		}
		if err := filepath.Walk(root, walkFn); err != nil {
			return nil, err
		}
	}
	return pkgFiles, nil
}

// moveAndCleanupGeneratedSrc moves all generated src to their correct locations
// in the repository, because protoc puts it in a folder called `generated/“.
func (g *GenprotoGenerator) moveAndCleanupGeneratedSrc() error {
	log.Println("moving generated code")
	// The period at the end is analogous to * (copy everything in this dir).
	c := execv.Command("cp", "-R", filepath.Join(g.genprotoDir, "generated", "google.golang.org", "genproto", "googleapis"), g.genprotoDir)
	if err := c.Run(); err != nil {
		return err
	}

	c = execv.Command("rm", "-rf", "generated")
	c.Dir = g.genprotoDir
	if err := c.Run(); err != nil {
		return err
	}

	return nil
}

func (g *GenprotoGenerator) generateAliases() error {
	for genprotoImport, newPkg := range aliasfix.GenprotoPkgMigration {
		if !isMigrated(genprotoImport) || g.gapicToGenerate == "" {
			continue
		}
		// remove the stubs dir segment from path
		gapicImport := newPkg.ImportPath[:strings.LastIndex(newPkg.ImportPath, "/")]
		if !strings.Contains(g.gapicToGenerate, gapicImport) {
			continue
		}
		srdDir := filepath.Join(g.googleCloudDir, strings.TrimPrefix(newPkg.ImportPath, "cloud.google.com/go/"))
		destDir := filepath.Join(g.genprotoDir, "googleapis", strings.TrimPrefix(genprotoImport, "google.golang.org/genproto/googleapis/"))
		if err := aliasgen.Run(srdDir, destDir); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package interceptor

import (
	"encoding/json"
	"log"
	"net"
	"strings"
	"sync"

	"github.com/datawire/ambassador/pkg/supervisor"

	"github.com/datawire/ambassador/internal/pkg/nat"
	rt "github.com/datawire/ambassador/internal/pkg/route"
)

type Interceptor struct {
	translator *nat.Translator
	tables     map[string]rt.Table
	tablesLock sync.RWMutex

	domains     map[string]rt.Route
	domainsLock sync.RWMutex

	search     []string
	searchLock sync.RWMutex

	work chan func(*supervisor.Process) error
}

func NewInterceptor(name string) *Interceptor {
	ret := &Interceptor{
		tables:     make(map[string]rt.Table),
		translator: nat.NewTranslator(name),
		domains:    make(map[string]rt.Route),
		search:     []string{""},
		work:       make(chan func(*supervisor.Process) error),
	}
	ret.tablesLock.Lock() // leave it locked until .Start() unlocks it
	return ret
}

func (i *Interceptor) Work(p *supervisor.Process) error {
	i.translator.Enable(p)
	i.tablesLock.Unlock()

	p.Ready()

	for {
		select {
		case <-p.Shutdown():
			i.tablesLock.Lock()
			i.translator.Disable(p)
			// leave it locked
			return nil
		case f := <-i.work:
			err := f(p)
			if err != nil {
				return err
			}
		}
	}
}

// Resolve looks up the given query in the (FIXME: somewhere), trying
// all the suffixes in the search path, and returns a Route on success
// or nil on failure. This implementation does not count the number of
// dots in the query.
func (i *Interceptor) Resolve(query string) *rt.Route {
	if !strings.HasSuffix(query, ".") {
		query += "."
	}

	const prefix = "teleproxy"
	const suffix = ".cachebust.telepresence.io." // must end with .
	const replacement = "teleproxy."             // must end with .
	if strings.HasPrefix(query, prefix) && strings.HasSuffix(query, suffix) {
		query = replacement
	}

	i.searchLock.RLock()
	defer i.searchLock.RUnlock()
	i.domainsLock.RLock()
	defer i.domainsLock.RUnlock()

	for _, suffix := range i.search {
		name := query + suffix
		value, ok := i.domains[strings.ToLower(name)]
		if ok {
			return &value
		}
	}
	return nil
}

func (i *Interceptor) Destination(conn *net.TCPConn) (string, error) {
	_, host, err := i.translator.GetOriginalDst(conn)
	return host, err
}

func (i *Interceptor) Render(table string) string {
	var obj interface{}

	if table == "" {
		var tables []rt.Table
		i.tablesLock.RLock()
		for _, t := range i.tables {
			tables = append(tables, t)
		}
		i.tablesLock.RUnlock()
		obj = tables
	} else {
		var ok bool
		i.tablesLock.RLock()
		obj, ok = i.tables[table]
		i.tablesLock.RUnlock()
		if !ok {
			return ""
		}
	}

	bytes, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		return err.Error()
	} else {
		return string(bytes)
	}
}

func (i *Interceptor) Delete(table string) bool {
	result := make(chan bool)
	i.work <- func(p *supervisor.Process) error {
		i.tablesLock.Lock()
		defer i.tablesLock.Unlock()
		i.domainsLock.Lock()
		defer i.domainsLock.Unlock()

		var names []string
		if table == "" {
			for name := range i.tables {
				names = append(names, name)
			}
		} else if _, ok := i.tables[table]; ok {
			names = []string{table}
		} else {
			result <- false
			return nil
		}

		for _, name := range names {
			if name != "bootstrap" {
				err := i.update(p, rt.Table{Name: name})
				if err != nil {
					return err
				}
			}
		}

		result <- true
		return nil
	}

	return <-result
}

func (i *Interceptor) Update(table rt.Table) {
	result := make(chan struct{})
	i.work <- func(p *supervisor.Process) error {
		defer close(result)
		return i.update(p, table)
	}
	<-result
}

func (i *Interceptor) update(p *supervisor.Process, table rt.Table) error {
	// Make a copy of the current table
	i.tablesLock.Lock()
	oldTable, ok := i.tables[table.Name]
	oldRoutes := make(map[string]rt.Route)
	if ok {
		for _, route := range oldTable.Routes {
			oldRoutes[route.Name] = route
		}
	}
	i.tablesLock.Unlock()

	// Operate on the copy of the current table and the new table
	for _, newRoute := range table.Routes {
		oldRoute, oldRouteOk := oldRoutes[newRoute.Name]
		// A nil Route (when oldRouteOk != true) will compare
		// inequal to any valid new Route.
		if newRoute != oldRoute {
			log.Printf("INT: new route %v:%v domain %v does not match old route", newRoute.Ip, newRoute.Port, newRoute.Domain())
			log.Printf("INT: old route exists for %v:%v and domain %v", oldRoute.Ip, oldRoute.Port, oldRoute.Domain())
			// We're updating a route. Make sure DNS waits until the new answer
			// is ready, i.e. don't serve the old answer.
			i.domainsLock.Lock()

			// delete the old version
			if oldRouteOk {
				switch newRoute.Proto {
				case "tcp":
					i.translator.ClearTCP(p, oldRoute.Ip, oldRoute.Port)
				case "udp":
					i.translator.ClearUDP(p, oldRoute.Ip, oldRoute.Port)
				default:
					log.Printf("INT: unrecognized protocol: %v", newRoute)
				}
			}
			// and add the new version
			if newRoute.Target != "" {
				switch newRoute.Proto {
				case "tcp":
					i.translator.ForwardTCP(p, newRoute.Ip, newRoute.Port, newRoute.Target)
				case "udp":
					i.translator.ForwardUDP(p, newRoute.Ip, newRoute.Port, newRoute.Target)
				default:
					log.Printf("INT: unrecognized protocol: %v", newRoute)
				}
			}

			if newRoute.Name != "" {
				log.Printf("INT: STORE %v->%v", newRoute.Domain(), newRoute)
				i.domains[newRoute.Domain()] = newRoute
			}

			i.domainsLock.Unlock()
		}

		// remove the route from our map of old routes so we
		// don't end up deleting it below
		delete(oldRoutes, newRoute.Name)
	}

	// Clear out stale routes and DNS names
	i.domainsLock.Lock()
	for _, route := range oldRoutes {
		log.Printf("INT: CLEAR %v->%v", route.Domain(), route)
		delete(i.domains, route.Domain())

		switch route.Proto {
		case "tcp":
			i.translator.ClearTCP(p, route.Ip, route.Port)
		case "udp":
			i.translator.ClearUDP(p, route.Ip, route.Port)
		default:
			log.Printf("INT: unrecognized protocol: %v", route)
		}
	}
	i.domainsLock.Unlock()

	// Update the externally-visible table
	i.tablesLock.Lock()
	if table.Routes == nil || len(table.Routes) == 0 {
		delete(i.tables, table.Name)
	} else {
		i.tables[table.Name] = table
	}
	i.tablesLock.Unlock()

	return nil
}

// SetSearchPath updates the DNS search path used by the resolver
func (i *Interceptor) SetSearchPath(paths []string) {
	i.searchLock.Lock()
	defer i.searchLock.Unlock()

	i.search = paths
}

// GetSearchPath retrieves the current search path
func (i *Interceptor) GetSearchPath() []string {
	i.searchLock.RLock()
	defer i.searchLock.RUnlock()

	return i.search
}
<commit_msg>(from AES) Revert "Add debug statements"<commit_after>package interceptor

import (
	"encoding/json"
	"log"
	"net"
	"strings"
	"sync"

	"github.com/datawire/ambassador/pkg/supervisor"

	"github.com/datawire/ambassador/internal/pkg/nat"
	rt "github.com/datawire/ambassador/internal/pkg/route"
)

type Interceptor struct {
	translator *nat.Translator
	tables     map[string]rt.Table
	tablesLock sync.RWMutex

	domains     map[string]rt.Route
	domainsLock sync.RWMutex

	search     []string
	searchLock sync.RWMutex

	work chan func(*supervisor.Process) error
}

func NewInterceptor(name string) *Interceptor {
	ret := &Interceptor{
		tables:     make(map[string]rt.Table),
		translator: nat.NewTranslator(name),
		domains:    make(map[string]rt.Route),
		search:     []string{""},
		work:       make(chan func(*supervisor.Process) error),
	}
	ret.tablesLock.Lock() // leave it locked until .Start() unlocks it
	return ret
}

func (i *Interceptor) Work(p *supervisor.Process) error {
	i.translator.Enable(p)
	i.tablesLock.Unlock()

	p.Ready()

	for {
		select {
		case <-p.Shutdown():
			i.tablesLock.Lock()
			i.translator.Disable(p)
			// leave it locked
			return nil
		case f := <-i.work:
			err := f(p)
			if err != nil {
				return err
			}
		}
	}
}

// Resolve looks up the given query in the (FIXME: somewhere), trying
// all the suffixes in the search path, and returns a Route on success
// or nil on failure. This implementation does not count the number of
// dots in the query.
func (i *Interceptor) Resolve(query string) *rt.Route {
	if !strings.HasSuffix(query, ".") {
		query += "."
	}

	const prefix = "teleproxy"
	const suffix = ".cachebust.telepresence.io." // must end with .
	const replacement = "teleproxy."             // must end with .
	if strings.HasPrefix(query, prefix) && strings.HasSuffix(query, suffix) {
		query = replacement
	}

	i.searchLock.RLock()
	defer i.searchLock.RUnlock()
	i.domainsLock.RLock()
	defer i.domainsLock.RUnlock()

	for _, suffix := range i.search {
		name := query + suffix
		value, ok := i.domains[strings.ToLower(name)]
		if ok {
			return &value
		}
	}
	return nil
}

func (i *Interceptor) Destination(conn *net.TCPConn) (string, error) {
	_, host, err := i.translator.GetOriginalDst(conn)
	return host, err
}

func (i *Interceptor) Render(table string) string {
	var obj interface{}

	if table == "" {
		var tables []rt.Table
		i.tablesLock.RLock()
		for _, t := range i.tables {
			tables = append(tables, t)
		}
		i.tablesLock.RUnlock()
		obj = tables
	} else {
		var ok bool
		i.tablesLock.RLock()
		obj, ok = i.tables[table]
		i.tablesLock.RUnlock()
		if !ok {
			return ""
		}
	}

	bytes, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		return err.Error()
	} else {
		return string(bytes)
	}
}

func (i *Interceptor) Delete(table string) bool {
	result := make(chan bool)
	i.work <- func(p *supervisor.Process) error {
		i.tablesLock.Lock()
		defer i.tablesLock.Unlock()
		i.domainsLock.Lock()
		defer i.domainsLock.Unlock()

		var names []string
		if table == "" {
			for name := range i.tables {
				names = append(names, name)
			}
		} else if _, ok := i.tables[table]; ok {
			names = []string{table}
		} else {
			result <- false
			return nil
		}

		for _, name := range names {
			if name != "bootstrap" {
				err := i.update(p, rt.Table{Name: name})
				if err != nil {
					return err
				}
			}
		}

		result <- true
		return nil
	}

	return <-result
}

func (i *Interceptor) Update(table rt.Table) {
	result := make(chan struct{})
	i.work <- func(p *supervisor.Process) error {
		defer close(result)
		return i.update(p, table)
	}
	<-result
}

func (i *Interceptor) update(p *supervisor.Process, table rt.Table) error {
	// Make a copy of the current table
	i.tablesLock.Lock()
	oldTable, ok := i.tables[table.Name]
	oldRoutes := make(map[string]rt.Route)
	if ok {
		for _, route := range oldTable.Routes {
			oldRoutes[route.Name] = route
		}
	}
	i.tablesLock.Unlock()

	// Operate on the copy of the current table and the new table
	for _, newRoute := range table.Routes {
		oldRoute, oldRouteOk := oldRoutes[newRoute.Name]
		// A nil Route (when oldRouteOk != true) will compare
		// inequal to any valid new Route.
		if newRoute != oldRoute {
			// We're updating a route. Make sure DNS waits until the new answer
			// is ready, i.e. don't serve the old answer.
			i.domainsLock.Lock()

			// delete the old version
			if oldRouteOk {
				switch newRoute.Proto {
				case "tcp":
					i.translator.ClearTCP(p, oldRoute.Ip, oldRoute.Port)
				case "udp":
					i.translator.ClearUDP(p, oldRoute.Ip, oldRoute.Port)
				default:
					log.Printf("INT: unrecognized protocol: %v", newRoute)
				}
			}
			// and add the new version
			if newRoute.Target != "" {
				switch newRoute.Proto {
				case "tcp":
					i.translator.ForwardTCP(p, newRoute.Ip, newRoute.Port, newRoute.Target)
				case "udp":
					i.translator.ForwardUDP(p, newRoute.Ip, newRoute.Port, newRoute.Target)
				default:
					log.Printf("INT: unrecognized protocol: %v", newRoute)
				}
			}

			if newRoute.Name != "" {
				log.Printf("INT: STORE %v->%v", newRoute.Domain(), newRoute)
				i.domains[newRoute.Domain()] = newRoute
			}

			i.domainsLock.Unlock()
		}

		// remove the route from our map of old routes so we
		// don't end up deleting it below
		delete(oldRoutes, newRoute.Name)
	}

	// Clear out stale routes and DNS names
	i.domainsLock.Lock()
	for _, route := range oldRoutes {
		log.Printf("INT: CLEAR %v->%v", route.Domain(), route)
		delete(i.domains, route.Domain())

		switch route.Proto {
		case "tcp":
			i.translator.ClearTCP(p, route.Ip, route.Port)
		case "udp":
			i.translator.ClearUDP(p, route.Ip, route.Port)
		default:
			log.Printf("INT: unrecognized protocol: %v", route)
		}
	}
	i.domainsLock.Unlock()

	// Update the externally-visible table
	i.tablesLock.Lock()
	if table.Routes == nil || len(table.Routes) == 0 {
		delete(i.tables, table.Name)
	} else {
		i.tables[table.Name] = table
	}
	i.tablesLock.Unlock()

	return nil
}

// SetSearchPath updates the DNS search path used by the resolver
func (i *Interceptor) SetSearchPath(paths []string) {
	i.searchLock.Lock()
	defer i.searchLock.Unlock()

	i.search = paths
}

// GetSearchPath retrieves the current search path
func (i *Interceptor) GetSearchPath() []string {
	i.searchLock.RLock()
	defer i.searchLock.RUnlock()

	return i.search
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsKinesisStream() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsKinesisStreamCreate,
		Read:   resourceAwsKinesisStreamRead,
		Update: resourceAwsKinesisStreamUpdate,
		Delete: resourceAwsKinesisStreamDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"shard_count": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"retention_period": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  24,
				ValidateFunc: func(v interface{}, k string) (ws []string, errors []error) {
					value := v.(int)
					if value < 24 || value > 168 {
						errors = append(errors, fmt.Errorf(
							"%q must be between 24 and 168 hours", k))
					}
					return
				},
			},

			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"tags": tagsSchema(),
		},
	}
}

func resourceAwsKinesisStreamCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn
	sn := d.Get("name").(string)
	createOpts := &kinesis.CreateStreamInput{
		ShardCount: aws.Int64(int64(d.Get("shard_count").(int))),
		StreamName: aws.String(sn),
	}

	_, err := conn.CreateStream(createOpts)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			return fmt.Errorf("[WARN] Error creating Kinesis Stream: \"%s\", code: \"%s\"", awsErr.Message(), awsErr.Code())
		}
		return err
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"CREATING"},
		Target:     []string{"ACTIVE"},
		Refresh:    streamStateRefreshFunc(conn, sn),
		Timeout:    5 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}

	streamRaw, err := stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf(
			"Error waiting for Kinesis Stream (%s) to become active: %s",
			sn, err)
	}

	s := streamRaw.(kinesisStreamState)
	d.SetId(s.arn)
	d.Set("arn", s.arn)
	d.Set("shard_count", s.shardCount)

	return resourceAwsKinesisStreamUpdate(d, meta)
}

func resourceAwsKinesisStreamUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn

	d.Partial(true)
	if err := setTagsKinesis(conn, d); err != nil {
		return err
	}

	d.SetPartial("tags")
	d.Partial(false)

	if err := setKinesisRetentionPeriod(conn, d); err != nil {
		return err
	}

	return resourceAwsKinesisStreamRead(d, meta)
}

func resourceAwsKinesisStreamRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn
	sn := d.Get("name").(string)

	state, err := readKinesisStreamState(conn, sn)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				d.SetId("")
				return nil
			}
			return fmt.Errorf("[WARN] Error reading Kinesis Stream: \"%s\", code: \"%s\"", awsErr.Message(), awsErr.Code())
		}
		return err

	}
	d.Set("arn", state.arn)
	d.Set("shard_count", state.shardCount)
	d.Set("retention_period", state.retentionPeriod)

	// set tags
	describeTagsOpts := &kinesis.ListTagsForStreamInput{
		StreamName: aws.String(sn),
	}
	tagsResp, err := conn.ListTagsForStream(describeTagsOpts)
	if err != nil {
		log.Printf("[DEBUG] Error retrieving tags for Stream: %s. %s", sn, err)
	} else {
		d.Set("tags", tagsToMapKinesis(tagsResp.Tags))
	}

	return nil
}

func resourceAwsKinesisStreamDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn
	sn := d.Get("name").(string)
	_, err := conn.DeleteStream(&kinesis.DeleteStreamInput{
		StreamName: aws.String(sn),
	})

	if err != nil {
		return err
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"DELETING"},
		Target:     []string{"DESTROYED"},
		Refresh:    streamStateRefreshFunc(conn, sn),
		Timeout:    5 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}

	_, err = stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf(
			"Error waiting for Stream (%s) to be destroyed: %s",
			sn, err)
	}

	d.SetId("")
	return nil
}

func setKinesisRetentionPeriod(conn *kinesis.Kinesis, d *schema.ResourceData) error {
	sn := d.Get("name").(string)

	oraw, nraw := d.GetChange("retention_period")
	o := oraw.(int)
	n := nraw.(int)

	if n == 0 {
		log.Printf("[DEBUG] Kinesis Stream (%q) Retention Period Not Changed", sn)
		return nil
	}

	if n > o {
		log.Printf("[DEBUG] Increasing %s Stream Retention Period to %d", sn, n)
		_, err := conn.IncreaseStreamRetentionPeriod(&kinesis.IncreaseStreamRetentionPeriodInput{
			StreamName:           aws.String(sn),
			RetentionPeriodHours: aws.Int64(int64(n)),
		})
		if err != nil {
			return err
		}

	} else {
		log.Printf("[DEBUG] Decreasing %s Stream Retention Period to %d", sn, n)
		_, err := conn.DecreaseStreamRetentionPeriod(&kinesis.DecreaseStreamRetentionPeriodInput{
			StreamName:           aws.String(sn),
			RetentionPeriodHours: aws.Int64(int64(n)),
		})
		if err != nil {
			return err
		}
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"UPDATING"},
		Target:     []string{"ACTIVE"},
		Refresh:    streamStateRefreshFunc(conn, sn),
		Timeout:    5 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}

	_, err := stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf(
			"Error waiting for Kinesis Stream (%s) to become active: %s",
			sn, err)
	}

	return nil
}

type kinesisStreamState struct {
	arn             string
	status          string
	shardCount      int
	retentionPeriod int64
}

func readKinesisStreamState(conn *kinesis.Kinesis, sn string) (kinesisStreamState, error) {
	describeOpts := &kinesis.DescribeStreamInput{
		StreamName: aws.String(sn),
	}

	var state kinesisStreamState
	err := conn.DescribeStreamPages(describeOpts, func(page *kinesis.DescribeStreamOutput, last bool) (shouldContinue bool) {
		state.arn = aws.StringValue(page.StreamDescription.StreamARN)
		state.status = aws.StringValue(page.StreamDescription.StreamStatus)
		state.shardCount += len(page.StreamDescription.Shards)
		state.retentionPeriod = aws.Int64Value(page.StreamDescription.RetentionPeriodHours)
		return !last
	})
	return state, err
}

func streamStateRefreshFunc(conn *kinesis.Kinesis, sn string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		state, err := readKinesisStreamState(conn, sn)
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "ResourceNotFoundException" {
					return 42, "DESTROYED", nil
				}
				return nil, awsErr.Code(), err
			}
			return nil, "failed", err
		}

		return state, state.status, nil
	}
}
<commit_msg>Report the correct number of shards<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsKinesisStream() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsKinesisStreamCreate,
		Read:   resourceAwsKinesisStreamRead,
		Update: resourceAwsKinesisStreamUpdate,
		Delete: resourceAwsKinesisStreamDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"shard_count": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"retention_period": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  24,
				ValidateFunc: func(v interface{}, k string) (ws []string, errors []error) {
					value := v.(int)
					if value < 24 || value > 168 {
						errors = append(errors, fmt.Errorf(
							"%q must be between 24 and 168 hours", k))
					}
					return
				},
			},

			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"tags": tagsSchema(),
		},
	}
}

func resourceAwsKinesisStreamCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn
	sn := d.Get("name").(string)
	createOpts := &kinesis.CreateStreamInput{
		ShardCount: aws.Int64(int64(d.Get("shard_count").(int))),
		StreamName: aws.String(sn),
	}

	_, err := conn.CreateStream(createOpts)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			return fmt.Errorf("[WARN] Error creating Kinesis Stream: \"%s\", code: \"%s\"", awsErr.Message(), awsErr.Code())
		}
		return err
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"CREATING"},
		Target:     []string{"ACTIVE"},
		Refresh:    streamStateRefreshFunc(conn, sn),
		Timeout:    5 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}

	streamRaw, err := stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf(
			"Error waiting for Kinesis Stream (%s) to become active: %s",
			sn, err)
	}

	s := streamRaw.(kinesisStreamState)
	d.SetId(s.arn)
	d.Set("arn", s.arn)
	d.Set("shard_count", s.shardCount)

	return resourceAwsKinesisStreamUpdate(d, meta)
}

func resourceAwsKinesisStreamUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn

	d.Partial(true)
	if err := setTagsKinesis(conn, d); err != nil {
		return err
	}

	d.SetPartial("tags")
	d.Partial(false)

	if err := setKinesisRetentionPeriod(conn, d); err != nil {
		return err
	}

	return resourceAwsKinesisStreamRead(d, meta)
}

func resourceAwsKinesisStreamRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn
	sn := d.Get("name").(string)

	state, err := readKinesisStreamState(conn, sn)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				d.SetId("")
				return nil
			}
			return fmt.Errorf("[WARN] Error reading Kinesis Stream: \"%s\", code: \"%s\"", awsErr.Message(), awsErr.Code())
		}
		return err

	}
	d.Set("arn", state.arn)
	d.Set("shard_count", state.shardCount)
	d.Set("retention_period", state.retentionPeriod)

	// set tags
	describeTagsOpts := &kinesis.ListTagsForStreamInput{
		StreamName: aws.String(sn),
	}
	tagsResp, err := conn.ListTagsForStream(describeTagsOpts)
	if err != nil {
		log.Printf("[DEBUG] Error retrieving tags for Stream: %s. %s", sn, err)
	} else {
		d.Set("tags", tagsToMapKinesis(tagsResp.Tags))
	}

	return nil
}

func resourceAwsKinesisStreamDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kinesisconn
	sn := d.Get("name").(string)
	_, err := conn.DeleteStream(&kinesis.DeleteStreamInput{
		StreamName: aws.String(sn),
	})

	if err != nil {
		return err
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"DELETING"},
		Target:     []string{"DESTROYED"},
		Refresh:    streamStateRefreshFunc(conn, sn),
		Timeout:    5 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}

	_, err = stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf(
			"Error waiting for Stream (%s) to be destroyed: %s",
			sn, err)
	}

	d.SetId("")
	return nil
}

func setKinesisRetentionPeriod(conn *kinesis.Kinesis, d *schema.ResourceData) error {
	sn := d.Get("name").(string)

	oraw, nraw := d.GetChange("retention_period")
	o := oraw.(int)
	n := nraw.(int)

	if n == 0 {
		log.Printf("[DEBUG] Kinesis Stream (%q) Retention Period Not Changed", sn)
		return nil
	}

	if n > o {
		log.Printf("[DEBUG] Increasing %s Stream Retention Period to %d", sn, n)
		_, err := conn.IncreaseStreamRetentionPeriod(&kinesis.IncreaseStreamRetentionPeriodInput{
			StreamName:           aws.String(sn),
			RetentionPeriodHours: aws.Int64(int64(n)),
		})
		if err != nil {
			return err
		}

	} else {
		log.Printf("[DEBUG] Decreasing %s Stream Retention Period to %d", sn, n)
		_, err := conn.DecreaseStreamRetentionPeriod(&kinesis.DecreaseStreamRetentionPeriodInput{
			StreamName:           aws.String(sn),
			RetentionPeriodHours: aws.Int64(int64(n)),
		})
		if err != nil {
			return err
		}
	}

	stateConf := &resource.StateChangeConf{
		Pending:    []string{"UPDATING"},
		Target:     []string{"ACTIVE"},
		Refresh:    streamStateRefreshFunc(conn, sn),
		Timeout:    5 * time.Minute,
		Delay:      10 * time.Second,
		MinTimeout: 3 * time.Second,
	}

	_, err := stateConf.WaitForState()
	if err != nil {
		return fmt.Errorf(
			"Error waiting for Kinesis Stream (%s) to become active: %s",
			sn, err)
	}

	return nil
}

type kinesisStreamState struct {
	arn             string
	status          string
	shardCount      int
	retentionPeriod int64
}

func readKinesisStreamState(conn *kinesis.Kinesis, sn string) (kinesisStreamState, error) {
	describeOpts := &kinesis.DescribeStreamInput{
		StreamName: aws.String(sn),
	}

	var state kinesisStreamState
	err := conn.DescribeStreamPages(describeOpts, func(page *kinesis.DescribeStreamOutput, last bool) (shouldContinue bool) {
		state.arn = aws.StringValue(page.StreamDescription.StreamARN)
		state.status = aws.StringValue(page.StreamDescription.StreamStatus)
		state.shardCount += len(openShards(page.StreamDescription.Shards))
		state.retentionPeriod = aws.Int64Value(page.StreamDescription.RetentionPeriodHours)
		return !last
	})
	return state, err
}

func streamStateRefreshFunc(conn *kinesis.Kinesis, sn string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		state, err := readKinesisStreamState(conn, sn)
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "ResourceNotFoundException" {
					return 42, "DESTROYED", nil
				}
				return nil, awsErr.Code(), err
			}
			return nil, "failed", err
		}

		return state, state.status, nil
	}
}

// See http://docs.aws.amazon.com/kinesis/latest/dev/kinesis-using-sdk-java-resharding-merge.html
func openShards(shards []*kinesis.Shard) []*kinesis.Shard {
	var open []*kinesis.Shard
	for _, s := range shards {
		if s.SequenceNumberRange.EndingSequenceNumber == nil {
			open = append(open, s)
		}
	}

	return open
}
<|endoftext|>