<|start_filename|>src/query/storage/promremote/storage_test.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package promremote

import (
	"context"
	"io"
	"math/rand"
	"net/http"
	"testing"
	"time"

	"github.com/m3db/m3/src/query/models"
	"github.com/m3db/m3/src/query/storage"
	"github.com/m3db/m3/src/query/storage/m3/storagemetadata"
	"github.com/m3db/m3/src/query/storage/promremote/promremotetest"
	"github.com/m3db/m3/src/query/ts"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/tallytest"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/prometheus/prometheus/prompb"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
	"go.uber.org/zap"
)

var (
	logger, _ = zap.NewDevelopment()
	scope     = tally.NewTestScope("test_scope", map[string]string{})
)

func TestWrite(t *testing.T) {
	fakeProm := promremotetest.NewServer(t)
	defer fakeProm.Close()

	promStorage, err := NewStorage(Options{
		endpoints: []EndpointOptions{{name: "testEndpoint", address: fakeProm.WriteAddr()}},
		scope:     scope,
		logger:    logger,
	})
	require.NoError(t, err)
	defer closeWithCheck(t, promStorage)

	now := xtime.Now()
	wq, err := storage.NewWriteQuery(storage.WriteQueryOptions{
		Tags: models.Tags{
			Opts: models.NewTagOptions(),
			Tags: []models.Tag{{
				Name:  []byte("test_tag_name"),
				Value: []byte("test_tag_value"),
			}},
		},
		Datapoints: ts.Datapoints{{
			Timestamp: now,
			Value:     42,
		}},
		Unit: xtime.Millisecond,
	})
	require.NoError(t, err)
	err = promStorage.Write(context.TODO(), wq)
	require.NoError(t, err)

	promWrite := fakeProm.GetLastWriteRequest()

	expectedLabel := prompb.Label{
		Name:  "test_tag_name",
		Value: "test_tag_value",
	}
	expectedSample := prompb.Sample{
		Value:     42,
		Timestamp: now.ToNormalizedTime(time.Millisecond),
	}
	require.Len(t, promWrite.Timeseries, 1)
	require.Len(t, promWrite.Timeseries[0].Labels, 1)
	require.Len(t, promWrite.Timeseries[0].Samples, 1)
	assert.Equal(t, expectedLabel, promWrite.Timeseries[0].Labels[0])
	assert.Equal(t, expectedSample, promWrite.Timeseries[0].Samples[0])

	tallytest.AssertCounterValue(
		t, 1, scope.Snapshot(), "test_scope.prom_remote_storage.writeSingle.success",
		map[string]string{"endpoint_name": "testEndpoint"},
	)
	tallytest.AssertCounterValue(
		t, 0, scope.Snapshot(), "test_scope.prom_remote_storage.writeSingle.errors",
		map[string]string{"endpoint_name": "testEndpoint"},
	)
}

func TestWriteBasedOnRetention(t *testing.T) {
	promShortRetention := promremotetest.NewServer(t)
	defer promShortRetention.Close()
	promMediumRetention := promremotetest.NewServer(t)
	defer promMediumRetention.Close()
	promLongRetention := promremotetest.NewServer(t)
	defer promLongRetention.Close()
	promLongRetention2 := promremotetest.NewServer(t)
	defer promLongRetention2.Close()
	reset := func() {
		promShortRetention.Reset()
		promMediumRetention.Reset()
		promLongRetention.Reset()
		promLongRetention2.Reset()
	}

	mediumRetentionAttr := storagemetadata.Attributes{
		MetricsType: storagemetadata.AggregatedMetricsType,
		Retention:   720 * time.Hour,
		Resolution:  5 * time.Minute,
	}
	shortRetentionAttr := storagemetadata.Attributes{
		MetricsType: storagemetadata.AggregatedMetricsType,
		Retention:   120 * time.Hour,
		Resolution:  15 * time.Second,
	}
	longRetentionAttr := storagemetadata.Attributes{
		Resolution: 10 * time.Minute,
		Retention:  8760 * time.Hour,
	}
	promStorage, err := NewStorage(Options{
		endpoints: []EndpointOptions{
			{
				address:    promShortRetention.WriteAddr(),
				attributes: shortRetentionAttr,
			},
			{
				address:    promMediumRetention.WriteAddr(),
				attributes: mediumRetentionAttr,
			},
			{
				address:    promLongRetention.WriteAddr(),
				attributes: longRetentionAttr,
			},
			{
				address:    promLongRetention2.WriteAddr(),
				attributes: longRetentionAttr,
			},
		},
		scope:  scope,
		logger: logger,
	})
	require.NoError(t, err)
	defer closeWithCheck(t, promStorage)

	t.Run("send short retention write", func(t *testing.T) {
		reset()
		err := writeTestMetric(t, promStorage, shortRetentionAttr)
		require.NoError(t, err)
		assert.NotNil(t, promShortRetention.GetLastWriteRequest())
		assert.Nil(t, promMediumRetention.GetLastWriteRequest())
		assert.Nil(t, promLongRetention.GetLastWriteRequest())
	})

	t.Run("send medium retention write", func(t *testing.T) {
		reset()
		err := writeTestMetric(t, promStorage, mediumRetentionAttr)
		require.NoError(t, err)
		assert.Nil(t, promShortRetention.GetLastWriteRequest())
		assert.NotNil(t, promMediumRetention.GetLastWriteRequest())
		assert.Nil(t, promLongRetention.GetLastWriteRequest())
	})

	t.Run("send write to multiple instances configured with same retention", func(t *testing.T) {
		reset()
		err := writeTestMetric(t, promStorage, longRetentionAttr)
		require.NoError(t, err)
		assert.Nil(t, promShortRetention.GetLastWriteRequest())
		assert.Nil(t, promMediumRetention.GetLastWriteRequest())
		assert.NotNil(t, promLongRetention.GetLastWriteRequest())
		assert.NotNil(t, promLongRetention2.GetLastWriteRequest())
	})

	t.Run("send unconfigured retention write", func(t *testing.T) {
		reset()
		err := writeTestMetric(t, promStorage, storagemetadata.Attributes{
			Resolution: mediumRetentionAttr.Resolution + 1,
			Retention:  mediumRetentionAttr.Retention,
		})
		require.Error(t, err)
		err = writeTestMetric(t, promStorage, storagemetadata.Attributes{
			Resolution: mediumRetentionAttr.Resolution,
			Retention:  mediumRetentionAttr.Retention + 1,
		})
		require.Error(t, err)
		assert.Contains(t, err.Error(), "write did not match any of known endpoints")
		assert.Nil(t, promShortRetention.GetLastWriteRequest())
		assert.Nil(t, promMediumRetention.GetLastWriteRequest())
		assert.Nil(t, promLongRetention.GetLastWriteRequest())
		const droppedWrites = "test_scope.prom_remote_storage.dropped_writes"
		tallytest.AssertCounterValue(t, 2, scope.Snapshot(), droppedWrites, map[string]string{})
	})

	t.Run("error should not prevent sending to other instances", func(t *testing.T) {
		reset()
		promLongRetention.SetError("test err", http.StatusInternalServerError)
		err := writeTestMetric(t, promStorage, longRetentionAttr)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "test err")
		assert.NotNil(t, promLongRetention2.GetLastWriteRequest())
	})
}

func TestErrorHandling(t *testing.T) {
	svr := promremotetest.NewServer(t)
	defer svr.Close()

	attr := storagemetadata.Attributes{
		MetricsType: storagemetadata.AggregatedMetricsType,
		Retention:   720 * time.Hour,
		Resolution:  5 * time.Minute,
	}
	promStorage, err := NewStorage(Options{
		endpoints: []EndpointOptions{{address: svr.WriteAddr(), attributes: attr}},
		scope:     scope,
		logger:    logger,
	})
	require.NoError(t, err)
	defer closeWithCheck(t, promStorage)

	t.Run("wrap non 5xx errors as invalid params error", func(t *testing.T) {
		svr.Reset()
		svr.SetError("test err", http.StatusForbidden)
		err := writeTestMetric(t, promStorage, attr)
		require.Error(t, err)
		assert.True(t, xerrors.IsInvalidParams(err))
	})

	t.Run("429 should not be wrapped as invalid params", func(t *testing.T) {
		svr.Reset()
		svr.SetError("test err", http.StatusTooManyRequests)
		err := writeTestMetric(t, promStorage, attr)
		require.Error(t, err)
		assert.False(t, xerrors.IsInvalidParams(err))
	})
}

func closeWithCheck(t *testing.T, c io.Closer) {
	require.NoError(t, c.Close())
}

func writeTestMetric(t *testing.T, s storage.Storage, attr storagemetadata.Attributes) error {
	//nolint: gosec
	datapoint := ts.Datapoint{Value: rand.Float64(), Timestamp: xtime.Now()}
	wq, err := storage.NewWriteQuery(storage.WriteQueryOptions{
		Tags: models.Tags{
			Opts: models.NewTagOptions(),
			Tags: []models.Tag{{
				Name:  []byte("test_tag_name"),
				Value: []byte("test_tag_value"),
			}},
		},
		Datapoints: ts.Datapoints{datapoint},
		Unit:       xtime.Millisecond,
		Attributes: attr,
	})
	require.NoError(t, err)
	return s.Write(context.TODO(), wq)
}


<|start_filename|>src/cmd/tools/annotation_checker/main/main.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"sort"
	"strconv"

	"github.com/m3db/m3/src/dbnode/encoding"
	"github.com/m3db/m3/src/dbnode/encoding/m3tsz"
	"github.com/m3db/m3/src/dbnode/persist"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/dbnode/x/xio"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/ident"
	"github.com/m3db/m3/src/x/pool"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/pborman/getopt"
	"go.uber.org/zap"
)

const (
	snapshotType = "snapshot"
	flushType    = "flush"

	allShards = -1
)

func main() {
	var (
		optPathPrefix  = getopt.StringLong("path-prefix", 'p', "", "Path prefix [e.g. /var/lib/m3db]")
		fileSetTypeArg = getopt.StringLong("fileset-type", 't', flushType, fmt.Sprintf("%s|%s", flushType, snapshotType))
		optNamespace   = getopt.StringLong("namespace", 'n', "default", "Namespace [e.g. metrics]")
		optShard       = getopt.IntLong("shard", 's', allShards,
			fmt.Sprintf("Shard number, or %v for all shards in the directory", allShards))
		optBlockstart = getopt.Int64Long("block-start", 'b', 0, "Block Start Time [in nsec]")
		volume        = getopt.Int64Long("volume", 'v', 0, "Volume number")

		annotationRewrittenFilter = getopt.BoolLong("annotation-rewritten", 'R', "Filters metrics with annotation rewrites")
	)
	getopt.Parse()

	rawLogger, err := zap.NewDevelopment()
	if err != nil {
		log.Fatalf("unable to create logger: %+v", err)
	}
	log := rawLogger.Sugar()

	if *optPathPrefix == "" ||
		*optNamespace == "" ||
		*optShard < allShards ||
		*optBlockstart <= 0 ||
		*volume < 0 ||
		(*fileSetTypeArg != snapshotType && *fileSetTypeArg != flushType) {
		getopt.Usage()
		os.Exit(1)
	}

	var fileSetType persist.FileSetType
	switch *fileSetTypeArg {
	case flushType:
		fileSetType = persist.FileSetFlushType
	case snapshotType:
		fileSetType = persist.FileSetSnapshotType
	default:
		log.Fatalf("unknown fileset type: %s", *fileSetTypeArg)
	}

	shards := []uint32{uint32(*optShard)}
	if *optShard == allShards {
		shards, err = getShards(*optPathPrefix, fileSetType, *optNamespace)
		if err != nil {
			log.Fatalf("failed extracting the list of shards: %v", err)
		}
	}

	// Not using bytes pool with streaming reads/writes to avoid the fixed memory overhead.
	var bytesPool pool.CheckedBytesPool
	encodingOpts := encoding.NewOptions().SetBytesPool(bytesPool)

	fsOpts := fs.NewOptions().SetFilePathPrefix(*optPathPrefix)
	reader, err := fs.NewReader(bytesPool, fsOpts)
	if err != nil {
		log.Fatalf("could not create new reader: %v", err)
	}

	metricsMap := make(map[string]struct{})

	for _, shard := range shards {
		openOpts := fs.DataReaderOpenOptions{
			Identifier: fs.FileSetFileIdentifier{
				Namespace:   ident.StringID(*optNamespace),
				Shard:       shard,
				BlockStart:  xtime.UnixNano(*optBlockstart),
				VolumeIndex: int(*volume),
			},
			FileSetType:      fileSetType,
			StreamingEnabled: true,
		}

		err = reader.Open(openOpts)
		if err != nil {
			log.Fatalf("unable to open reader for shard %v: %v", shard, err)
		}

		for {
			entry, err := reader.StreamingRead()
			if xerrors.Is(err, io.EOF) {
				break
			}
			if err != nil {
				log.Fatalf("err reading metadata: %v", err)
			}

			noInitialAnnotation, annotationRewritten, err := checkAnnotations(entry.Data, encodingOpts)
			if err != nil {
				log.Fatalf("failed checking annotations: %v", err)
			}

			if (!*annotationRewrittenFilter && noInitialAnnotation) || (*annotationRewrittenFilter && annotationRewritten) {
				metricsMap[entry.ID.String()] = struct{}{}
			}
		}

		if err := reader.Close(); err != nil {
			log.Fatalf("unable to close reader: %v", err)
		}
	}

	metrics := make([]string, 0)
	for m := range metricsMap {
		metrics = append(metrics, m)
	}
	sort.Strings(metrics)
	for _, metric := range metrics {
		fmt.Println(metric) // nolint: forbidigo
	}
}

func checkAnnotations(data []byte, encodingOpts encoding.Options) (bool, bool, error) {
	iter := m3tsz.NewReaderIterator(xio.NewBytesReader64(data), true, encodingOpts)
	defer iter.Close()

	var (
		previousAnnotationBase64 *string
		noInitialAnnotation      = true
		annotationRewritten      = false

		firstDatapoint = true
	)

	for iter.Next() {
		_, _, annotation := iter.Current()
		if len(annotation) > 0 {
			if firstDatapoint {
				noInitialAnnotation = false
			}
			annotationBase64 := base64.StdEncoding.EncodeToString(annotation)
			if previousAnnotationBase64 != nil && *previousAnnotationBase64 != annotationBase64 {
				annotationRewritten = true
			}
			previousAnnotationBase64 = &annotationBase64
		}
		firstDatapoint = false
	}
	if err := iter.Err(); err != nil {
		return false, false, fmt.Errorf("unable to iterate original data: %w", err)
	}

	return noInitialAnnotation, annotationRewritten, nil
}

func getShards(pathPrefix string, fileSetType persist.FileSetType, namespace string) ([]uint32, error) {
	nsID := ident.StringID(namespace)
	path := fs.NamespaceDataDirPath(pathPrefix, nsID)
	if fileSetType == persist.FileSetSnapshotType {
		path = fs.NamespaceSnapshotsDirPath(pathPrefix, nsID)
	}

	files, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, fmt.Errorf("failed reading namespace directory: %w", err)
	}

	shards := make([]uint32, 0)
	for _, f := range files {
		if !f.IsDir() {
			continue
		}
		i, err := strconv.Atoi(f.Name())
		if err != nil {
			return nil, fmt.Errorf("failed extracting shard number: %w", err)
		}
		if i < 0 {
			return nil, fmt.Errorf("negative shard number %v", i)
		}
		shards = append(shards, uint32(i))
	}

	return shards, nil
}


<|start_filename|>src/dbnode/encoding/annotation_holder_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package encoding

import (
	"bytes"
	"testing"

	"github.com/m3db/m3/src/dbnode/ts"

	"github.com/stretchr/testify/assert"
)

var annotationHolderTests = []struct {
	name       string
	annotation ts.Annotation
}{
	{
		name: "nil",
	},
	{
		name:       "empty",
		annotation: []byte{},
	},
	{
		name:       "short",
		annotation: []byte{1, 2, 3},
	},
	{
		name:       "OptimizedAnnotationLen - 1",
		annotation: repeat(1, ts.OptimizedAnnotationLen-1),
	},
	{
		name:       "OptimizedAnnotationLen",
		annotation: repeat(2, ts.OptimizedAnnotationLen),
	},
	{
		name:       "OptimizedAnnotationLen + 1",
		annotation: repeat(3, ts.OptimizedAnnotationLen+1),
	},
	{
		name:       "OptimizedAnnotationLen * 2",
		annotation: repeat(4, ts.OptimizedAnnotationLen*2),
	},
}

func TestAnnotationHolder(t *testing.T) {
	holder := annotationHolder{}
	assert.Nil(t, holder.get())

	for _, tt := range annotationHolderTests {
		t.Run(tt.name, func(t *testing.T) {
			holder.set(tt.annotation)
			if len(tt.annotation) == 0 {
				assert.Empty(t, holder.get())
			} else {
				assert.Equal(t, tt.annotation, holder.get())
			}
		})
	}
}

func TestAnnotationHolderWithReset(t *testing.T) {
	holder := annotationHolder{}
	holder.reset()
	assert.Nil(t, holder.get())

	for _, tt := range annotationHolderTests {
		t.Run(tt.name, func(t *testing.T) {
			holder.set(tt.annotation)
			if len(tt.annotation) == 0 {
				assert.Empty(t, holder.get())
			} else {
				assert.Equal(t, tt.annotation, holder.get())
			}
			holder.reset()
			assert.Nil(t, holder.get())
		})
	}
}

func repeat(b byte, n int) ts.Annotation {
	return bytes.Repeat([]byte{b}, n)
}


<|start_filename|>src/x/debug/test/test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// Package debugtest contains test helper methods
package debugtest

import (
	"testing"

	clusterclient "github.com/m3db/m3/src/cluster/client"
	"github.com/m3db/m3/src/cluster/generated/proto/placementpb"
	"github.com/m3db/m3/src/cluster/kv"
	clusterplacement "github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/placementhandler"
	"github.com/m3db/m3/src/cluster/services"
	"github.com/m3db/m3/src/query/api/v1/handler/namespace"
	"github.com/m3db/m3/src/x/instrument"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/require"
)

// NewTestHandlerOptsAndClient returns a new test handler.
func NewTestHandlerOptsAndClient(
	t *testing.T,
) (
	placementhandler.HandlerOptions,
	*kv.MockStore,
	*clusterclient.MockClient,
) {
	ctrl := gomock.NewController(t)

	placementProto := &placementpb.Placement{
		Instances: map[string]*placementpb.Instance{
			"host1": {
				Id:             "host1",
				IsolationGroup: "rack1",
				Zone:           "test",
				Weight:         1,
				Endpoint:       "http://host1:1234",
				Hostname:       "host1",
				Port:           1234,
			},
			"host2": {
				Id:             "host2",
				IsolationGroup: "rack1",
				Zone:           "test",
				Weight:         1,
				Endpoint:       "http://host2:1234",
				Hostname:       "host2",
				Port:           1234,
			},
		},
	}

	mockPlacement := clusterplacement.NewMockPlacement(ctrl)
	mockPlacement.EXPECT().Proto().Return(placementProto, nil).AnyTimes()
	mockPlacement.EXPECT().Version().Return(0).AnyTimes()

	return NewTestHandlerOptsAndClientWithPlacement(t, ctrl, mockPlacement)
}

// NewTestHandlerOptsAndClientWithPlacement returns a new test handler with supplied placement.
func NewTestHandlerOptsAndClientWithPlacement(
	t *testing.T,
	ctrl *gomock.Controller,
	placement clusterplacement.Placement,
) (
	placementhandler.HandlerOptions,
	*kv.MockStore,
	*clusterclient.MockClient,
) {
	mockClient := clusterclient.NewMockClient(ctrl)
	require.NotNil(t, mockClient)

	mockKV := kv.NewMockStore(ctrl)
	require.NotNil(t, mockKV)

	mockServices := services.NewMockServices(ctrl)
	require.NotNil(t, mockServices)

	mockPlacementService := clusterplacement.NewMockService(ctrl)
	require.NotNil(t, mockPlacementService)

	mockClient.EXPECT().Services(gomock.Not(nil)).Return(mockServices, nil).AnyTimes()
	mockServices.EXPECT().PlacementService(gomock.Not(nil), gomock.Not(nil)).Return(mockPlacementService, nil).AnyTimes()
	mockPlacementService.EXPECT().Placement().Return(placement, nil).AnyTimes()

	mockClient.EXPECT().KV().Return(mockKV, nil).AnyTimes()
	mockKV.EXPECT().Get(namespace.M3DBNodeNamespacesKey).Return(nil, kv.ErrNotFound).AnyTimes()

	handlerOpts, err := placementhandler.NewHandlerOptions(
		mockClient, clusterplacement.Configuration{}, nil, instrument.NewOptions())
	require.NoError(t, err)

	return handlerOpts, mockKV, mockClient
}


<|start_filename|>src/msg/producer/config/writer.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package config

import (
	"errors"
	"time"

	"github.com/m3db/m3/src/cluster/client"
	"github.com/m3db/m3/src/cluster/kv"
	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/services"
	"github.com/m3db/m3/src/msg/producer/writer"
	"github.com/m3db/m3/src/msg/protocol/proto"
	"github.com/m3db/m3/src/msg/topic"
	"github.com/m3db/m3/src/x/instrument"
	xio "github.com/m3db/m3/src/x/io"
	"github.com/m3db/m3/src/x/pool"
	"github.com/m3db/m3/src/x/retry"

	"github.com/uber-go/tally"
)

// ConnectionConfiguration configs the connection options.
type ConnectionConfiguration struct {
	NumConnections  *int                 `yaml:"numConnections"`
	DialTimeout     *time.Duration       `yaml:"dialTimeout"`
	WriteTimeout    *time.Duration       `yaml:"writeTimeout"`
	KeepAlivePeriod *time.Duration       `yaml:"keepAlivePeriod"`
	ResetDelay      *time.Duration       `yaml:"resetDelay"`
	Retry           *retry.Configuration `yaml:"retry"`
	FlushInterval   *time.Duration       `yaml:"flushInterval"`
	WriteBufferSize *int                 `yaml:"writeBufferSize"`
	ReadBufferSize  *int                 `yaml:"readBufferSize"`
	// ContextDialer specifies a custom dialer to use when creating TCP connections to the consumer.
	// See writer.ConnectionOptions.ContextDialer for details.
	ContextDialer writer.ContextDialerFn `yaml:"-"` // not serializable
}

// NewOptions creates connection options.
func (c *ConnectionConfiguration) NewOptions(iOpts instrument.Options) writer.ConnectionOptions {
	opts := writer.NewConnectionOptions()
	if c.NumConnections != nil {
		opts = opts.SetNumConnections(*c.NumConnections)
	}
	if c.ContextDialer != nil {
		opts = opts.SetContextDialer(c.ContextDialer)
	}
	if c.DialTimeout != nil {
		opts = opts.SetDialTimeout(*c.DialTimeout)
	}
	if c.WriteTimeout != nil {
		opts = opts.SetWriteTimeout(*c.WriteTimeout)
	}
	if c.KeepAlivePeriod != nil {
		opts = opts.SetKeepAlivePeriod(*c.KeepAlivePeriod)
	}
	if c.ResetDelay != nil {
		opts = opts.SetResetDelay(*c.ResetDelay)
	}
	if c.Retry != nil {
		opts = opts.SetRetryOptions(c.Retry.NewOptions(iOpts.MetricsScope()))
	}
	if c.FlushInterval != nil {
		opts = opts.SetFlushInterval(*c.FlushInterval)
	}
	if c.WriteBufferSize != nil {
		opts = opts.SetWriteBufferSize(*c.WriteBufferSize)
	}
	if c.ReadBufferSize != nil {
		opts = opts.SetReadBufferSize(*c.ReadBufferSize)
	}
	return opts.SetInstrumentOptions(iOpts)
}

// WriterConfiguration configs the writer options.
type WriterConfiguration struct {
	TopicName                 string                         `yaml:"topicName" validate:"nonzero"`
	TopicServiceOverride      kv.OverrideConfiguration       `yaml:"topicServiceOverride"`
	TopicWatchInitTimeout     *time.Duration                 `yaml:"topicWatchInitTimeout"`
	PlacementOptions          placement.Configuration        `yaml:"placement"`
	PlacementServiceOverride  services.OverrideConfiguration `yaml:"placementServiceOverride"`
	PlacementWatchInitTimeout *time.Duration                 `yaml:"placementWatchInitTimeout"`
	// MessagePool configuration is deprecated for producer side.
	MessagePool                       *pool.ObjectPoolConfiguration `yaml:"messagePool"`
	MessageQueueNewWritesScanInterval *time.Duration                `yaml:"messageQueueNewWritesScanInterval"`
	MessageQueueFullScanInterval      *time.Duration                `yaml:"messageQueueFullScanInterval"`
	MessageQueueScanBatchSize         *int                          `yaml:"messageQueueScanBatchSize"`
	InitialAckMapSize                 *int                          `yaml:"initialAckMapSize"`
	CloseCheckInterval                *time.Duration                `yaml:"closeCheckInterval"`
	AckErrorRetry                     *retry.Configuration          `yaml:"ackErrorRetry"`
	Encoder                           *proto.Configuration          `yaml:"encoder"`
	Decoder                           *proto.Configuration          `yaml:"decoder"`
	Connection                        *ConnectionConfiguration      `yaml:"connection"`

	// StaticMessageRetry configs a static message retry policy.
	StaticMessageRetry *StaticMessageRetryConfiguration `yaml:"staticMessageRetry"`
	// MessageRetry configs a algorithmic retry policy.
	// Only one of the retry configuration should be used.
	MessageRetry *retry.Configuration `yaml:"messageRetry"`

	// IgnoreCutoffCutover allows producing writes ignoring cutoff/cutover timestamp.
	// Must be in sync with AggregatorConfiguration.WritesIgnoreCutoffCutover.
	IgnoreCutoffCutover bool `yaml:"ignoreCutoffCutover"`
	// WithoutConsumerScope drops the consumer tag from the metrics. For large m3msg deployments the consumer tag can
	// add a lot of cardinality to the metrics.
	WithoutConsumerScope bool `yaml:"withoutConsumerScope"`
}

// StaticMessageRetryConfiguration configs the static message retry policy.
// When messageRetry config exists, messageRetry will override the static config.
type StaticMessageRetryConfiguration struct {
	Backoff []time.Duration `yaml:"backoff"`
}

// NewOptions creates writer options.
func (c *WriterConfiguration) NewOptions(
	cs client.Client,
	iOpts instrument.Options,
	rwOptions xio.Options,
) (writer.Options, error) {
	opts := writer.NewOptions().
		SetTopicName(c.TopicName).
		SetPlacementOptions(c.PlacementOptions.NewOptions()).
		SetInstrumentOptions(iOpts).
		SetWithoutConsumerScope(c.WithoutConsumerScope)

	kvOpts, err := c.TopicServiceOverride.NewOverrideOptions()
	if err != nil {
		return nil, err
	}

	topicServiceOpts := topic.NewServiceOptions().
		SetConfigService(cs).
		SetKVOverrideOptions(kvOpts)
	ts, err := topic.NewService(topicServiceOpts)
	if err != nil {
		return nil, err
	}

	opts = opts.SetTopicService(ts)

	if c.TopicWatchInitTimeout != nil {
		opts = opts.SetTopicWatchInitTimeout(*c.TopicWatchInitTimeout)
	}
	sd, err := cs.Services(c.PlacementServiceOverride.NewOptions())
	if err != nil {
		return nil, err
	}

	opts = opts.SetServiceDiscovery(sd)

	if c.PlacementWatchInitTimeout != nil {
		opts = opts.SetPlacementWatchInitTimeout(*c.PlacementWatchInitTimeout)
	}
	opts, err = c.setRetryOptions(opts, iOpts)
	if err != nil {
		return nil, err
	}
	if c.MessageQueueNewWritesScanInterval != nil {
		opts = opts.SetMessageQueueNewWritesScanInterval(*c.MessageQueueNewWritesScanInterval)
	}
	if c.MessageQueueFullScanInterval != nil {
		opts = opts.SetMessageQueueFullScanInterval(*c.MessageQueueFullScanInterval)
	}
	if c.MessageQueueScanBatchSize != nil {
		opts = opts.SetMessageQueueScanBatchSize(*c.MessageQueueScanBatchSize)
	}
	if c.InitialAckMapSize != nil {
		opts = opts.SetInitialAckMapSize(*c.InitialAckMapSize)
	}
	if c.CloseCheckInterval != nil {
		opts = opts.SetCloseCheckInterval(*c.CloseCheckInterval)
	}
	if c.AckErrorRetry != nil {
		opts = opts.SetAckErrorRetryOptions(c.AckErrorRetry.NewOptions(tally.NoopScope))
	}
	if c.Encoder != nil {
		opts = opts.SetEncoderOptions(c.Encoder.NewOptions(iOpts))
	}
	if c.Decoder != nil {
		opts = opts.SetDecoderOptions(c.Decoder.NewOptions(iOpts))
	}
	if c.Connection != nil {
		opts = opts.SetConnectionOptions(c.Connection.NewOptions(iOpts))
	}

	opts = opts.SetIgnoreCutoffCutover(c.IgnoreCutoffCutover)

	opts = opts.SetDecoderOptions(opts.DecoderOptions().SetRWOptions(rwOptions))
	return opts, nil
}

func (c *WriterConfiguration) setRetryOptions(
	opts writer.Options,
	iOpts instrument.Options,
) (writer.Options, error) {
	if c.StaticMessageRetry != nil && c.MessageRetry != nil {
		return nil, errors.New("invalid writer config with both static and algorithmic retry config set")
	}
	if c.MessageRetry != nil {
		return opts.SetMessageRetryNanosFn(
			writer.NextRetryNanosFn(c.MessageRetry.NewOptions(iOpts.MetricsScope())),
		), nil
	}
	if c.StaticMessageRetry != nil {
		fn, err := writer.StaticRetryNanosFn(c.StaticMessageRetry.Backoff)
		if err != nil {
			return nil, err
		}
		return opts.SetMessageRetryNanosFn(fn), nil
	}
	return opts, nil
}


<|start_filename|>src/integration/resources/resources.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package resources

import (
	"errors"
	"fmt"
	"net"
	"strconv"
	"time"

	"github.com/m3db/m3/src/cluster/generated/proto/placementpb"
	"github.com/m3db/m3/src/cluster/placementhandler/handleroptions"
	"github.com/m3db/m3/src/dbnode/generated/proto/namespace"
	"github.com/m3db/m3/src/msg/generated/proto/topicpb"
	"github.com/m3db/m3/src/query/generated/proto/admin"
	"github.com/m3db/m3/src/x/instrument"

	protobuftypes "github.com/gogo/protobuf/types"
	"go.uber.org/zap"
)

const (
	retention = "6h"

	// AggName is the name of the aggregated namespace.
	AggName = "aggregated"
	// UnaggName is the name of the unaggregated namespace.
	UnaggName = "default"
	// ColdWriteNsName is the name for cold write namespace.
	ColdWriteNsName = "coldWritesRepairAndNoIndex"
	// AggregatorInputTopic is the m3msg topic name for coordinator->aggregator traffic.
	AggregatorInputTopic = "aggregator_ingest"
	// AggregatorOutputTopic is the m3msg topic name for aggregator->coordinator traffic.
	AggregatorOutputTopic = "aggregated_metrics"
)

// SetupCluster setups m3 cluster on provided docker containers.
func SetupCluster(
	cluster M3Resources,
	opts ClusterOptions,
) error { // nolint: gocyclo
	coordinator := cluster.Coordinator()
	iOpts := instrument.NewOptions()
	logger := iOpts.Logger().With(zap.String("source", "harness"))
	hosts := make([]*admin.Host, 0, len(cluster.Nodes()))
	ids := make([]string, 0, len(cluster.Nodes()))
	for i, n := range cluster.Nodes() {
		h, err := n.HostDetails(9000)
		if err != nil {
			logger.Error("could not get host details", zap.Error(err))
			return err
		}
		if opts.DBNode != nil && opts.DBNode.NumIsolationGroups > 0 {
			h.IsolationGroup = fmt.Sprintf("isogroup-%d", int32(i)%opts.DBNode.NumIsolationGroups)
		}

		hosts = append(hosts, h)
		ids = append(ids, h.GetId())
	}

	replicationFactor := int32(1)
	numShards := int32(4)
	if opts.DBNode != nil {
		if opts.DBNode.RF > 0 {
			replicationFactor = opts.DBNode.RF
		}
		if opts.DBNode.NumShards > 0 {
			numShards = opts.DBNode.NumShards
		}
	}

	var (
		unaggDatabase = admin.DatabaseCreateRequest{
			Type:              "cluster",
			NamespaceName:     UnaggName,
			RetentionTime:     retention,
			NumShards:         numShards,
			ReplicationFactor: replicationFactor,
			Hosts:             hosts,
		}

		aggNamespace = admin.NamespaceAddRequest{
			Name: AggName,
			Options: &namespace.NamespaceOptions{
				BootstrapEnabled:  true,
				FlushEnabled:      true,
				WritesToCommitLog: true,
				CleanupEnabled:    true,
				SnapshotEnabled:   true,
				IndexOptions: &namespace.IndexOptions{
					Enabled:        true,
					BlockSizeNanos: int64(30 * time.Minute),
				},
				RetentionOptions: &namespace.RetentionOptions{
					RetentionPeriodNanos:                     int64(6 * time.Hour),
					BlockSizeNanos:                           int64(30 * time.Minute),
					BufferFutureNanos:                        int64(2 * time.Minute),
					BufferPastNanos:                          int64(10 * time.Minute),
					BlockDataExpiry:                          true,
					BlockDataExpiryAfterNotAccessPeriodNanos: int64(time.Minute * 5),
				},
				AggregationOptions: &namespace.AggregationOptions{
					Aggregations: []*namespace.Aggregation{
						{
							Aggregated: true,
							Attributes: &namespace.AggregatedAttributes{
								ResolutionNanos: int64(5 * time.Second),
							},
						},
					},
				},
			},
		}

		coldWriteNamespace = admin.NamespaceAddRequest{
			Name: ColdWriteNsName,
			Options: &namespace.NamespaceOptions{
				BootstrapEnabled:      true,
				FlushEnabled:          true,
				WritesToCommitLog:     true,
				CleanupEnabled:        true,
				SnapshotEnabled:       true,
				RepairEnabled:         true,
				ColdWritesEnabled:     true,
				CacheBlocksOnRetrieve: &protobuftypes.BoolValue{Value: true},
				RetentionOptions: &namespace.RetentionOptions{
					RetentionPeriodNanos:                     int64(4 * time.Hour),
					BlockSizeNanos:                           int64(time.Hour),
					BufferFutureNanos:                        int64(time.Minute * 10),
					BufferPastNanos:                          int64(time.Minute * 10),
					BlockDataExpiry:                          true,
					BlockDataExpiryAfterNotAccessPeriodNanos: int64(time.Minute * 5),
				},
			},
		}
	)

	logger.Info("waiting for coordinator")
	if err := coordinator.WaitForNamespace(""); err != nil {
		return err
	}

	logger.Info("creating database", zap.Any("request", unaggDatabase))
	if _, err := coordinator.CreateDatabase(unaggDatabase); err != nil {
		return err
	}

	logger.Info("waiting for placements", zap.Strings("placement ids", ids))
	if err := coordinator.WaitForInstances(ids); err != nil {
		return err
	}

	logger.Info("waiting for namespace", zap.String("name", UnaggName))
	if err := coordinator.WaitForNamespace(UnaggName); err != nil {
		return err
	}

	logger.Info("creating namespace", zap.Any("request", aggNamespace))
	if _, err := coordinator.AddNamespace(aggNamespace); err != nil {
		return err
	}

	logger.Info("waiting for namespace", zap.String("name", AggName))
	if err := coordinator.WaitForNamespace(AggName); err != nil {
		return err
	}

	logger.Info("creating namespace", zap.Any("request", coldWriteNamespace))
	if _, err := coordinator.AddNamespace(coldWriteNamespace); err != nil {
		return err
	}

	logger.Info("waiting for namespace", zap.String("name", ColdWriteNsName))
	if err := coordinator.WaitForNamespace(ColdWriteNsName); err != nil {
		return err
	}

	logger.Info("waiting for healthy")
	if err := cluster.Nodes().WaitForHealthy(); err != nil {
		return err
	}

	logger.Info("waiting for shards ready")
	if err := coordinator.WaitForShardsReady(); err != nil {
		return err
	}

	logger.Info("waiting for cluster to be ready")
	if err := coordinator.WaitForClusterReady(); err != nil {
		return err
	}

	if opts.Aggregator != nil {
		aggregators := cluster.Aggregators()
		if len(aggregators) == 0 {
			return errors.New("no aggregators have been initiazted")
		}

		if err := aggregators.WaitForHealthy(); err != nil {
			return err
		}
	}

	logger.Info("all healthy")

	return nil
}

// SetupPlacement configures the placement for the provided coordinators
// and aggregators.
func SetupPlacement(
	coordAPI Coordinator,
	coordHost InstanceInfo,
	aggs Aggregators,
	opts AggregatorClusterOptions,
) error {
	// Setup aggregator placement.
	var env, zone string
	instances := make([]*placementpb.Instance, 0, len(aggs))
	for i, agg := range aggs {
		info, err := agg.HostDetails()
		if err != nil {
			return err
		}

		if i == 0 {
			env = info.Env
			zone = info.Zone
		}

		instance := &placementpb.Instance{
			Id:             info.ID,
			IsolationGroup: fmt.Sprintf("isogroup-%02d", i%int(opts.NumIsolationGroups)),
			Zone:           info.Zone,
			Weight:         1,
			Endpoint:       net.JoinHostPort(info.M3msgAddress, strconv.Itoa(int(info.M3msgPort))),
			Hostname:       info.ID,
			Port:           info.M3msgPort,
		}

		instances = append(instances, instance)
	}

	aggPlacementRequestOptions := PlacementRequestOptions{
		Service: ServiceTypeM3Aggregator,
		Zone:    zone,
		Env:     env,
	}

	_, err := coordAPI.InitPlacement(
		aggPlacementRequestOptions,
		admin.PlacementInitRequest{
			NumShards:         opts.NumShards,
			ReplicationFactor: opts.RF,
			Instances:         instances,
			OptionOverride: &placementpb.Options{
				SkipPortMirroring: &protobuftypes.BoolValue{Value: true},
			},
		},
	)
	if err != nil {
		return fmt.Errorf("failed to init aggregator placement: %w", err)
	}

	// Setup coordinator placement.
	coordPlacementRequestOptions := PlacementRequestOptions{
		Service: ServiceTypeM3Coordinator,
		Zone:    coordHost.Zone,
		Env:     coordHost.Env,
	}
	_, err = coordAPI.InitPlacement(
		coordPlacementRequestOptions,
		admin.PlacementInitRequest{
			Instances: []*placementpb.Instance{
				{
					Id:       coordHost.ID,
					Zone:     coordHost.Zone,
					Endpoint: net.JoinHostPort(coordHost.M3msgAddress, strconv.Itoa(int(coordHost.M3msgPort))),
					Hostname: coordHost.ID,
					Port:     coordHost.M3msgPort,
				},
			},
		},
	)
	if err != nil {
		return fmt.Errorf("failed to init coordinator placement: %w", err)
	}

	return nil
}

// SetupM3MsgTopics sets up the m3msg topics for the provided coordinator
// and aggregator.
func SetupM3MsgTopics(
	coord Coordinator,
	aggInstanceInfo InstanceInfo,
	opts ClusterOptions,
) error {
	aggInputTopicOpts := M3msgTopicOptions{
		Zone:      aggInstanceInfo.Zone,
		Env:       aggInstanceInfo.Env,
		TopicName: AggregatorInputTopic,
	}

	aggOutputTopicOpts := M3msgTopicOptions{
		Zone:      aggInstanceInfo.Zone,
		Env:       aggInstanceInfo.Env,
		TopicName: AggregatorOutputTopic,
	}

	_, err := coord.InitM3msgTopic(
		aggInputTopicOpts,
		admin.TopicInitRequest{NumberOfShards: uint32(opts.Aggregator.NumShards)},
	)
	if err != nil {
		return fmt.Errorf("failed to init aggregator input m3msg topic: %w", err)
	}

	_, err = coord.AddM3msgTopicConsumer(aggInputTopicOpts, admin.TopicAddRequest{
		ConsumerService: &topicpb.ConsumerService{
			ServiceId: &topicpb.ServiceID{
				Name:        handleroptions.M3AggregatorServiceName,
				Environment: aggInputTopicOpts.Env,
				Zone:        aggInputTopicOpts.Zone,
			},
			ConsumptionType: topicpb.ConsumptionType_REPLICATED,
			MessageTtlNanos: 600000000000, // 10 mins
		},
	})
	if err != nil {
		return fmt.Errorf("failed to add aggregator input m3msg topic consumer: %w", err)
	}

	_, err = coord.InitM3msgTopic(
		aggOutputTopicOpts,
		admin.TopicInitRequest{NumberOfShards: uint32(opts.DBNode.NumShards)},
	)
	if err != nil {
		return fmt.Errorf("failed to init aggregator output m3msg topic: %w", err)
	}

	_, err = coord.AddM3msgTopicConsumer(aggOutputTopicOpts, admin.TopicAddRequest{
		ConsumerService: &topicpb.ConsumerService{
			ServiceId: &topicpb.ServiceID{
				Name:        handleroptions.M3CoordinatorServiceName,
				Environment: aggInputTopicOpts.Env,
				Zone:        aggInputTopicOpts.Zone,
			},
			ConsumptionType: topicpb.ConsumptionType_SHARED,
			MessageTtlNanos: 600000000000, // 10 mins
		},
	})
	if err != nil {
		return fmt.Errorf("failed to add aggregator output m3msg topic consumer: %w", err)
	}

	return nil
}


<|start_filename|>src/cmd/tools/read_commitlog/main/main.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"encoding/base64"
	"fmt"
	"log"
	"os"
	"sort"
	"time"

	"github.com/m3db/m3/src/x/ident"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/pborman/getopt"
	"go.uber.org/zap"
)

func main() {
	var (
		path         = getopt.StringLong("path", 'p', "", "file path [e.g. /var/lib/m3db/commitlogs/commitlog-0-161023.db]")
		idFilter     = getopt.StringLong("id-filter", 'f', "", "ID Contains Filter (optional)")
		idSizeFilter = getopt.IntLong("id-size-filter", 's', 0, "ID Size (bytes) Filter (optional)")
		mode         = getopt.StringLong("mode", 'm', "", "Action [print,summary]. Defaults to 'print'")
		top          = getopt.IntLong("top", 't', 0, "Print out only top N IDs")
	)
	getopt.Parse()

	rawLogger, err := zap.NewDevelopment()
	if err != nil {
		log.Fatalf("unable to create logger: %+v", err)
	}
	logger := rawLogger.Sugar()

	if *path == "" {
		getopt.Usage()
		os.Exit(1)
	}

	reader, err := newFilteringReader(*path, idFilter, idSizeFilter)
	if err != nil {
		logger.Fatalf("unable to open reader: %v", err)
	}
	defer reader.Close()

	switch *mode {
	case "summary":
		err = printSummary(reader, top)
	default:
		err = printMetrics(reader)
	}
	if err != nil {
		logger.Fatalf("error while reading commitlog: %v", err)
	}
}

func printMetrics(reader *filteringReader) error {
	var (
		entryCount          uint32
		annotationSizeTotal uint64
		start               = time.Now()
	)

	for {
		entry, found, err := reader.Read()
		if err != nil {
			return err
		}
		if !found {
			break
		}

		series := entry.Series
		fmt.Printf("{id: %s, dp: %+v, ns: %s, shard: %d", // nolint: forbidigo
			series.ID, entry.Datapoint, entry.Series.Namespace, entry.Series.Shard)
		if len(entry.Annotation) > 0 {
			fmt.Printf(", annotation: %s", // nolint: forbidigo
				base64.StdEncoding.EncodeToString(entry.Annotation))
			annotationSizeTotal += uint64(len(entry.Annotation))
		}
		fmt.Println("}") // nolint: forbidigo

		entryCount++
	}

	runTime := time.Since(start)

	fmt.Printf("\nRunning time: %s\n", runTime)                          // nolint: forbidigo
	fmt.Printf("%d entries read\n", entryCount)                          // nolint: forbidigo
	fmt.Printf("Total annotation size: %d bytes\n", annotationSizeTotal) // nolint: forbidigo
	return nil
}

func printSummary(reader *filteringReader, top *int) error {
	var (
		entryCount        uint32
		start             = time.Now()
		datapointCount    = map[ident.ID]uint32{}
		totalIDSize       uint64
		earliestDatapoint xtime.UnixNano
		oldestDatapoint   xtime.UnixNano
	)

	for {
		entry, found, err := reader.Read()
		if err != nil {
			return err
		}
		if !found {
			break
		}
		dp := entry.Datapoint

		if earliestDatapoint == 0 || earliestDatapoint > dp.TimestampNanos {
			earliestDatapoint = dp.TimestampNanos
		}
		if oldestDatapoint == 0 || oldestDatapoint < dp.TimestampNanos {
			oldestDatapoint = dp.TimestampNanos
		}

		datapointCount[entry.Series.ID]++

		entryCount++
	}

	runTime := time.Since(start)

	fmt.Printf("\nRunning time: %s\n", runTime)                                              // nolint: forbidigo
	fmt.Printf("%d entries read\n", entryCount)                                              // nolint: forbidigo
	fmt.Printf("time range [%s:%s]\n", earliestDatapoint.String(), oldestDatapoint.String()) // nolint: forbidigo

	datapointCountArr := idPairs{}
	sizeArr := idPairs{}
	for ID, count := range datapointCount {
		IDSize := len(ID.Bytes())
		totalIDSize += uint64(IDSize)
		datapointCountArr = append(datapointCountArr, idPair{ID: ID, Value: count})
		sizeArr = append(sizeArr, idPair{ID: ID, Value: uint32(IDSize)})
	}

	sort.Sort(sort.Reverse(datapointCountArr))
	sort.Sort(sort.Reverse(sizeArr))

	fmt.Printf("total ID size: %d bytes\n", totalIDSize)                  // nolint: forbidigo
	fmt.Printf("total distinct number of IDs %d \n", len(datapointCount)) // nolint: forbidigo

	limit := len(datapointCountArr)
	if *top > 0 && *top < limit {
		limit = *top
	}
	fmt.Printf("ID datapoint counts: \n") // nolint: forbidigo
	for i := 0; i < limit; i++ {
		pair := datapointCountArr[i]
		fmt.Printf("%-10d %s\n", pair.Value, pair.ID.String()) // nolint: forbidigo
	}

	fmt.Printf("ID sizes(bytes): \n") // nolint: forbidigo
	for i := 0; i < limit; i++ {
		pair := sizeArr[i]
		fmt.Printf("%-10d %s\n", pair.Value, pair.ID.String()) // nolint: forbidigo
	}

	return nil
}

type idPair struct {
	ID    ident.ID
	Value uint32
}

type idPairs []idPair

func (p idPairs) Len() int           { return len(p) }
func (p idPairs) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p idPairs) Less(i, j int) bool { return p[i].Value < p[j].Value }


<|start_filename|>src/dbnode/storage/bootstrap_instrumentation_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package storage

import (
	"errors"
	"strings"
	"testing"

	"github.com/m3db/m3/src/dbnode/storage/bootstrap"
	xerrors "github.com/m3db/m3/src/x/errors"

	"github.com/stretchr/testify/assert"
	"github.com/uber-go/tally"
)

func TestBootstrapRetryMetricReason(t *testing.T) {
	tests := []struct {
		name           string
		err            error
		expectedReason string
	}{
		{
			name:           "any error",
			err:            xerrors.NewInvalidParamsError(errors.New("some error")),
			expectedReason: "other",
		},
		{
			name:           "obsolete ranges error",
			err:            xerrors.NewInvalidParamsError(bootstrap.ErrFileSetSnapshotTypeRangeAdvanced),
			expectedReason: "obsolete-ranges",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			testScope := tally.NewTestScope("testScope", map[string]string{})
			options := NewOptions()
			options = options.SetInstrumentOptions(options.InstrumentOptions().SetMetricsScope(testScope))
			instr := newBootstrapInstrumentation(options)

			instr.bootstrapFailed(1, tt.err)

			retryMetrics := getBootstrapRetryMetrics(testScope)
			assert.NotEmpty(t, retryMetrics)
			assert.Equal(t, 1, retryMetrics[tt.expectedReason], "metric with reason:%v should be set to 1")
			for reason, value := range retryMetrics {
				if reason != tt.expectedReason {
					assert.Equal(t, 0, value, "metric with reason:%v should stay 0")
				}
			}
		})
	}
}

func getBootstrapRetryMetrics(testScope tally.TestScope) map[string]int {
	const (
		metricName = "bootstrap-retries"
		reasonTag  = "reason"
	)
	valuesByReason := make(map[string]int)
	for _, counter := range testScope.Snapshot().Counters() {
		if strings.Contains(counter.Name(), metricName) {
			reason := ""
			if r, ok := counter.Tags()[reasonTag]; ok {
				reason = r
			}
			valuesByReason[reason] = int(counter.Value())
		}
	}
	return valuesByReason
}


<|start_filename|>src/integration/resources/docker/prometheus.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package docker

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

	"github.com/m3db/m3/src/integration/resources"
	"github.com/m3db/m3/src/x/instrument"

	"github.com/ory/dockertest/v3"
	"github.com/prometheus/common/model"
)

// Prometheus is a docker-backed instantiation of Prometheus.
type Prometheus struct {
	pool      *dockertest.Pool
	pathToCfg string
	iOpts     instrument.Options

	resource *Resource
}

// PrometheusOptions contains the options for
// spinning up docker container running Prometheus
type PrometheusOptions struct {
	// Pool is the connection to the docker API
	Pool *dockertest.Pool
	// PathToCfg contains the path to the prometheus.yml configuration
	// file to be used on startup.
	PathToCfg string
	// InstrumentOptions are the instrument.Options to use when
	// creating the resource.
	InstrumentOptions instrument.Options
}

// NewPrometheus creates a new docker-backed Prometheus
// that implements the resources.ExternalResources interface.
func NewPrometheus(opts PrometheusOptions) resources.ExternalResources {
	if opts.InstrumentOptions == nil {
		opts.InstrumentOptions = instrument.NewOptions()
	}
	return &Prometheus{
		pool:      opts.Pool,
		pathToCfg: opts.PathToCfg,
		iOpts:     opts.InstrumentOptions,
	}
}

// Setup is a method that setups up the prometheus instance.
func (p *Prometheus) Setup() error {
	if p.resource != nil {
		return errors.New("prometheus already setup. must close resource " +
			"before attempting to setup again")
	}

	if err := SetupNetwork(p.pool); err != nil {
		return err
	}

	res, err := NewDockerResource(p.pool, ResourceOptions{
		ContainerName: "prometheus",
		Image: Image{
			Name: "prom/prometheus",
			Tag:  "latest",
		},
		PortList: []int{9090},
		Mounts: []string{
			fmt.Sprintf("%s:/etc/prometheus/prometheus.yml", p.pathToCfg),
		},
		InstrumentOpts: p.iOpts,
	})
	if err != nil {
		return err
	}

	p.resource = res

	return p.waitForHealthy()
}

func (p *Prometheus) waitForHealthy() error {
	return resources.Retry(func() error {
		req, err := http.NewRequestWithContext(
			context.Background(),
			http.MethodGet,
			"http://0.0.0.0:9090/-/ready",
			nil,
		)
		if err != nil {
			return err
		}

		client := http.Client{}
		res, _ := client.Do(req)
		if res != nil {
			_ = res.Body.Close()

			if res.StatusCode == http.StatusOK {
				return nil
			}
		}

		return errors.New("prometheus not ready")
	})
}

// PrometheusQueryRequest contains the parameters for making a query request.
type PrometheusQueryRequest struct {
	// Query is the prometheus query to execute
	Query string
	// Time is the time to execute the query at
	Time time.Time
}

// String converts the query request into a string suitable for use
// in the url params or request body
func (p *PrometheusQueryRequest) String() string {
	str := fmt.Sprintf("query=%v", p.Query)

	if !p.Time.IsZero() {
		str += fmt.Sprintf("&time=%v", p.Time.Unix())
	}

	return str
}

// Query executes a query request against the prometheus resource.
func (p *Prometheus) Query(req PrometheusQueryRequest) (model.Vector, error) {
	if p.resource.Closed() {
		return nil, errClosed
	}

	r, err := http.NewRequestWithContext(
		context.Background(),
		http.MethodGet,
		fmt.Sprintf("http://0.0.0.0:9090/api/v1/query?%s", req.String()),
		nil,
	)
	if err != nil {
		return nil, err
	}

	client := http.Client{}
	res, err := client.Do(r)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("non-200 status code received. "+
			"status=%v responseBody=%v", res.StatusCode, string(body))
	}

	var parsedResp jsonInstantQueryResponse
	if err := json.Unmarshal(body, &parsedResp); err != nil {
		return nil, err
	}

	return parsedResp.Data.Result, nil
}

type jsonInstantQueryResponse struct {
	Status string
	Data   vectorResult
}

type vectorResult struct {
	ResultType model.ValueType
	Result     model.Vector
}

// Close cleans up the prometheus instance.
func (p *Prometheus) Close() error {
	if p.resource.Closed() {
		return errClosed
	}

	if err := p.resource.Close(); err != nil {
		return err
	}

	p.resource = nil

	return nil
}


<|start_filename|>src/dbnode/network/server/tchannelthrift/cluster/service_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cluster

import (
	"fmt"
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/client"
	"github.com/m3db/m3/src/dbnode/encoding"
	"github.com/m3db/m3/src/dbnode/generated/thrift/rpc"
	"github.com/m3db/m3/src/dbnode/storage/index"
	"github.com/m3db/m3/src/dbnode/topology"
	xtest "github.com/m3db/m3/src/x/test"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel-go/thrift"
)

func TestSessionOpts(t *testing.T) {
	clientOpts := client.NewOptions().
		SetIterationOptions(index.IterationOptions{
			IterateEqualTimestampStrategy: encoding.IterateLowestValue,
		}).
		SetReadConsistencyLevel(topology.ReadConsistencyLevelMajority)
	rcAll := rpc.ReadConsistency_ALL
	csHighest := rpc.EqualTimestampStrategy_HIGHEST_VALUE
	cases := []struct {
		name     string
		reqOpts  *rpc.ClusterQueryOptions
		sessOpts sessionOpts
	}{
		{
			name: "nil opts",
			sessOpts: sessionOpts{
				readConsistency:        topology.ReadConsistencyLevelMajority,
				equalTimestampStrategy: encoding.IterateLowestValue,
			},
		},
		{
			name:    "not set",
			reqOpts: &rpc.ClusterQueryOptions{},
			sessOpts: sessionOpts{
				readConsistency:        topology.ReadConsistencyLevelMajority,
				equalTimestampStrategy: encoding.IterateLowestValue,
			},
		},
		{
			name: "only read consistency",
			reqOpts: &rpc.ClusterQueryOptions{
				ReadConsistency: &rcAll,
			},
			sessOpts: sessionOpts{
				readConsistency:        topology.ReadConsistencyLevelAll,
				equalTimestampStrategy: encoding.IterateLowestValue,
				consistencyOverride:    true,
			},
		},
		{
			name: "only conflict strategy",
			reqOpts: &rpc.ClusterQueryOptions{
				ConflictResolutionStrategy: &csHighest,
			},
			sessOpts: sessionOpts{
				readConsistency:        topology.ReadConsistencyLevelMajority,
				equalTimestampStrategy: encoding.IterateHighestValue,
			},
		},
		{
			name: "both set",
			reqOpts: &rpc.ClusterQueryOptions{
				ReadConsistency:            &rcAll,
				ConflictResolutionStrategy: &csHighest,
			},
			sessOpts: sessionOpts{
				readConsistency:        topology.ReadConsistencyLevelAll,
				equalTimestampStrategy: encoding.IterateHighestValue,
				consistencyOverride:    true,
			},
		},
	}

	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			ctrl := xtest.NewController(t)
			defer ctrl.Finish()

			c := client.NewMockClient(ctrl)
			c.EXPECT().Options().Return(clientOpts).AnyTimes()
			sess := client.NewMockSession(ctrl)
			iters := encoding.NewMockSeriesIterators(ctrl)
			sess.EXPECT().FetchTagged(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(
				iters, client.FetchResponseMetadata{}, nil)
			iters.EXPECT().Len().Return(0)
			iters.EXPECT().Iters().Return(nil)
			iters.EXPECT().Close()

			ctx, cancel := thrift.NewContext(time.Second)
			defer cancel()
			req := &rpc.QueryRequest{
				ClusterOptions: tc.reqOpts,
			}
			c.EXPECT().NewSessionWithOptions(tc.sessOpts).Return(sess, nil)
			s := NewService(c).(*service)
			res, err := s.Query(ctx, req)
			require.NoError(t, err)
			require.NotNil(t, res)
		})
	}
}

func (s sessionOpts) Matches(x interface{}) bool {
	if c, ok := x.(client.Options); ok {
		if s.equalTimestampStrategy != c.IterationOptions().IterateEqualTimestampStrategy {
			return false
		}
		if s.readConsistency != c.ReadConsistencyLevel() {
			return false
		}
		if s.consistencyOverride && c.RuntimeOptionsManager() != nil {
			return false
		}
		return true
	}
	return false
}

func (s sessionOpts) String() string {
	return fmt.Sprintf("%v %v", s.readConsistency, s.equalTimestampStrategy)
}

var _ gomock.Matcher = &sessionOpts{}


<|start_filename|>src/aggregator/aggregator/handler/writer/protobuf.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package writer

import (
	"errors"

	"github.com/m3db/m3/src/aggregator/sharding"
	"github.com/m3db/m3/src/metrics/encoding/protobuf"
	"github.com/m3db/m3/src/metrics/metric/aggregated"
	"github.com/m3db/m3/src/metrics/policy"
	"github.com/m3db/m3/src/msg/producer"

	"github.com/uber-go/tally"
)

var (
	errWriterClosed = errors.New("writer is closed")
)

type protobufWriterMetrics struct {
	writerClosed tally.Counter
	encodeErrors tally.Counter
	routeErrors  tally.Counter
}

func newProtobufWriterMetrics(scope tally.Scope) protobufWriterMetrics {
	encodeScope := scope.SubScope("encode")
	routeScope := scope.SubScope("route")
	return protobufWriterMetrics{
		writerClosed: scope.Counter("writer-closed"),
		encodeErrors: encodeScope.Counter("errors"),
		routeErrors:  routeScope.Counter("errors"),
	}
}

// protobufWriter encodes data and routes them to the backend.
// protobufWriter is not thread safe.
type protobufWriter struct {
	encoder   *protobuf.AggregatedEncoder
	p         producer.Producer
	numShards uint32

	m       aggregated.MetricWithStoragePolicy
	metrics protobufWriterMetrics

	shardFn sharding.ShardFn
	closed  bool
}

// NewProtobufWriter creates a writer that encodes metric in protobuf.
func NewProtobufWriter(
	producer producer.Producer,
	shardFn sharding.ShardFn,
	opts Options,
) Writer {
	instrumentOpts := opts.InstrumentOptions()
	w := &protobufWriter{
		encoder:   protobuf.NewAggregatedEncoder(opts.BytesPool()),
		p:         producer,
		numShards: producer.NumShards(),
		closed:    false,
		metrics:   newProtobufWriterMetrics(instrumentOpts.MetricsScope()),
		shardFn:   shardFn,
	}
	return w
}

func (w *protobufWriter) Write(mp aggregated.ChunkedMetricWithStoragePolicy) error {
	if w.closed {
		w.metrics.writerClosed.Inc(1)
		return errWriterClosed
	}

	m, shard := w.prepare(mp)
	if err := w.encoder.Encode(m); err != nil {
		w.metrics.encodeErrors.Inc(1)
		return err
	}

	if err := w.p.Produce(newMessage(shard, mp.StoragePolicy, w.encoder.Buffer())); err != nil {
		w.metrics.routeErrors.Inc(1)
		return err
	}

	return nil
}

func (w *protobufWriter) prepare(mp aggregated.ChunkedMetricWithStoragePolicy) (aggregated.MetricWithStoragePolicy, uint32) {
	// TODO(cw) Chunked metric has no 'type' field, consider adding one.
	w.m.ID = w.m.ID[:0]
	w.m.ID = append(w.m.ID, mp.Prefix...)
	w.m.ID = append(w.m.ID, mp.Data...)
	w.m.ID = append(w.m.ID, mp.Suffix...)
	w.m.Metric.TimeNanos = mp.TimeNanos
	w.m.Metric.Value = mp.Value
	w.m.Annotation = mp.ChunkedMetric.Annotation
	w.m.StoragePolicy = mp.StoragePolicy
	shard := w.shardFn(w.m.ID, w.numShards)
	return w.m, shard
}

func (w *protobufWriter) Flush() error {
	return nil
}

func (w *protobufWriter) Close() error {
	if w.closed {
		w.metrics.writerClosed.Inc(1)
		return errWriterClosed
	}
	// Don't close the producer here, it maybe shared by other writers.
	w.closed = true
	return nil
}

type message struct {
	shard uint32
	sp    policy.StoragePolicy
	data  protobuf.Buffer
}

func newMessage(shard uint32, sp policy.StoragePolicy, data protobuf.Buffer) producer.Message {
	return message{shard: shard, sp: sp, data: data}
}

func (d message) Shard() uint32 {
	return d.shard
}

func (d message) Bytes() []byte {
	return d.data.Bytes()
}

func (d message) Size() int {
	// Use the cap of the underlying byte slice in the buffer instead of
	// the length of the byte encoded to avoid "memory leak", for example
	// when the underlying buffer is 2KB, and it only encoded 300B, if we
	// use 300 as the size, then a producer with a buffer of 3GB could be
	// actually buffering 20GB in total for the underlying buffers.
	return cap(d.data.Bytes())
}

func (d message) Finalize(producer.FinalizeReason) {
	d.data.Close()
}

type storagePolicyFilter struct {
	acceptedStoragePolicies []policy.StoragePolicy
}

// NewStoragePolicyFilter creates a new storage policy based filter.
func NewStoragePolicyFilter(acceptedStoragePolicies []policy.StoragePolicy) producer.FilterFunc {
	return storagePolicyFilter{acceptedStoragePolicies}.Filter
}

func (f storagePolicyFilter) Filter(m producer.Message) bool {
	msg, ok := m.(message)
	if !ok {
		return true
	}
	for i := 0; i < len(f.acceptedStoragePolicies); i++ {
		if f.acceptedStoragePolicies[i].Equivalent(msg.sp) {
			return true
		}
	}
	return false
}


<|start_filename|>src/query/storage/promremote/storage.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package promremote

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
	"time"

	"github.com/m3db/m3/src/query/block"
	"github.com/m3db/m3/src/query/storage"
	"github.com/m3db/m3/src/query/storage/m3/consolidators"
	"github.com/m3db/m3/src/query/storage/m3/storagemetadata"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/instrument"
	xhttp "github.com/m3db/m3/src/x/net/http"

	"github.com/pkg/errors"
	"github.com/uber-go/tally"
	"go.uber.org/zap"
)

const metricsScope = "prom_remote_storage"

var errorReadingBody = []byte("error reading body")

var errNoEndpoints = errors.New("write did not match any of known endpoints")

// NewStorage returns new Prometheus remote write compatible storage
func NewStorage(opts Options) (storage.Storage, error) {
	client := xhttp.NewHTTPClient(opts.httpOptions)
	scope := opts.scope.SubScope(metricsScope)
	s := &promStorage{
		opts:            opts,
		client:          client,
		endpointMetrics: initEndpointMetrics(opts.endpoints, scope),
		droppedWrites:   scope.Counter("dropped_writes"),
		logger:          opts.logger,
	}
	return s, nil
}

type promStorage struct {
	unimplementedPromStorageMethods
	opts            Options
	client          *http.Client
	endpointMetrics map[string]instrument.MethodMetrics
	droppedWrites   tally.Counter
	logger          *zap.Logger
}

func (p *promStorage) Write(ctx context.Context, query *storage.WriteQuery) error {
	encoded, err := convertAndEncodeWriteQuery(query)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	multiErr := xerrors.NewMultiError()
	var errLock sync.Mutex
	atLeastOneEndpointMatched := false
	for _, endpoint := range p.opts.endpoints {
		endpoint := endpoint
		if endpoint.attributes.Resolution != query.Attributes().Resolution ||
			endpoint.attributes.Retention != query.Attributes().Retention {
			continue
		}

		metrics := p.endpointMetrics[endpoint.name]
		atLeastOneEndpointMatched = true

		wg.Add(1)
		go func() {
			defer wg.Done()
			err := p.writeSingle(ctx, metrics, endpoint.address, bytes.NewReader(encoded))
			if err != nil {
				errLock.Lock()
				multiErr = multiErr.Add(err)
				errLock.Unlock()
				return
			}
		}()
	}

	wg.Wait()

	if !atLeastOneEndpointMatched {
		p.droppedWrites.Inc(1)
		multiErr = multiErr.Add(errNoEndpoints)
		p.logger.Warn(
			"write did not match any of known endpoints",
			zap.Duration("retention", query.Attributes().Retention),
			zap.Duration("resolution", query.Attributes().Resolution),
		)
	}
	return multiErr.FinalError()
}

func (p *promStorage) Type() storage.Type {
	return storage.TypeRemoteDC
}

func (p *promStorage) Close() error {
	p.client.CloseIdleConnections()
	return nil
}

func (p *promStorage) ErrorBehavior() storage.ErrorBehavior {
	return storage.BehaviorFail
}

func (p *promStorage) Name() string {
	return "prom-remote"
}

func (p *promStorage) writeSingle(
	ctx context.Context,
	metrics instrument.MethodMetrics,
	address string,
	encoded io.Reader,
) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, address, encoded)
	if err != nil {
		return err
	}
	req.Header.Set("content-encoding", "snappy")
	req.Header.Set(xhttp.HeaderContentType, xhttp.ContentTypeProtobuf)

	start := time.Now()
	resp, err := p.client.Do(req)
	methodDuration := time.Since(start)
	if err != nil {
		metrics.ReportError(methodDuration)
		return err
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode/100 != 2 {
		metrics.ReportError(methodDuration)
		response, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			p.logger.Error("error reading body", zap.Error(err))
			response = errorReadingBody
		}
		genericError := fmt.Errorf(
			"expected status code 2XX: actual=%v, address=%v, resp=%s",
			resp.StatusCode, address, response,
		)
		if resp.StatusCode < 500 && resp.StatusCode != http.StatusTooManyRequests {
			return xerrors.NewInvalidParamsError(genericError)
		}
		return genericError
	}
	metrics.ReportSuccess(methodDuration)
	return nil
}

func initEndpointMetrics(endpoints []EndpointOptions, scope tally.Scope) map[string]instrument.MethodMetrics {
	metrics := make(map[string]instrument.MethodMetrics, len(endpoints))
	for _, endpoint := range endpoints {
		endpointScope := scope.Tagged(map[string]string{"endpoint_name": endpoint.name})
		methodMetrics := instrument.NewMethodMetrics(endpointScope, "writeSingle", instrument.TimerOptions{
			Type:             instrument.HistogramTimerType,
			HistogramBuckets: tally.DefaultBuckets,
		})
		metrics[endpoint.name] = methodMetrics
	}
	return metrics
}

var _ storage.Storage = &promStorage{}

type unimplementedPromStorageMethods struct{}

func (p *unimplementedPromStorageMethods) FetchProm(
	_ context.Context,
	_ *storage.FetchQuery,
	_ *storage.FetchOptions,
) (storage.PromResult, error) {
	return storage.PromResult{}, unimplementedError("FetchProm")
}

func (p *unimplementedPromStorageMethods) FetchBlocks(
	_ context.Context,
	_ *storage.FetchQuery,
	_ *storage.FetchOptions,
) (block.Result, error) {
	return block.Result{}, unimplementedError("FetchBlocks")
}

func (p *unimplementedPromStorageMethods) FetchCompressed(
	_ context.Context,
	_ *storage.FetchQuery,
	_ *storage.FetchOptions,
) (consolidators.MultiFetchResult, error) {
	return nil, unimplementedError("FetchCompressed")
}

func (p *unimplementedPromStorageMethods) SearchSeries(
	_ context.Context,
	_ *storage.FetchQuery,
	_ *storage.FetchOptions,
) (*storage.SearchResults, error) {
	return nil, unimplementedError("SearchSeries")
}

func (p *unimplementedPromStorageMethods) CompleteTags(
	_ context.Context,
	_ *storage.CompleteTagsQuery,
	_ *storage.FetchOptions,
) (*consolidators.CompleteTagsResult, error) {
	return nil, unimplementedError("CompleteTags")
}

func (p *unimplementedPromStorageMethods) QueryStorageMetadataAttributes(
	_ context.Context,
	_, _ time.Time,
	_ *storage.FetchOptions,
) ([]storagemetadata.Attributes, error) {
	return nil, unimplementedError("QueryStorageMetadataAttributes")
}

func unimplementedError(name string) error {
	return fmt.Errorf("promStorage: %s method is not supported", name)
}


<|start_filename|>src/integration/resources/coordinator_client.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package resources

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/m3db/m3/src/cluster/generated/proto/placementpb"
	"github.com/m3db/m3/src/cluster/placementhandler"
	"github.com/m3db/m3/src/query/api/v1/handler/graphite"
	"github.com/m3db/m3/src/query/api/v1/handler/prometheus/native"
	"github.com/m3db/m3/src/query/api/v1/handler/topic"
	"github.com/m3db/m3/src/query/api/v1/options"
	"github.com/m3db/m3/src/query/api/v1/route"
	"github.com/m3db/m3/src/query/generated/proto/admin"
	"github.com/m3db/m3/src/query/generated/proto/prompb"
	"github.com/m3db/m3/src/x/headers"
	xhttp "github.com/m3db/m3/src/x/net/http"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/prometheus/common/model"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

var errUnknownServiceType = errors.New("unknown service type")

// RetryFunc is a function that retries the provided
// operation until successful.
type RetryFunc func(op func() error) error

// ZapMethod appends the method as a log field.
func ZapMethod(s string) zapcore.Field { return zap.String("method", s) }

// CoordinatorClient is a client use to invoke API calls
// on a coordinator
type CoordinatorClient struct {
	client    *http.Client
	httpPort  int
	logger    *zap.Logger
	retryFunc RetryFunc
}

// CoordinatorClientOptions are the options for the CoordinatorClient.
type CoordinatorClientOptions struct {
	Client    *http.Client
	HTTPPort  int
	Logger    *zap.Logger
	RetryFunc RetryFunc
}

// NewCoordinatorClient creates a new CoordinatorClient.
func NewCoordinatorClient(opts CoordinatorClientOptions) CoordinatorClient {
	return CoordinatorClient{
		client:    opts.Client,
		httpPort:  opts.HTTPPort,
		logger:    opts.Logger,
		retryFunc: opts.RetryFunc,
	}
}

func (c *CoordinatorClient) makeURL(resource string) string {
	return fmt.Sprintf("http://0.0.0.0:%d/%s", c.httpPort, strings.TrimPrefix(resource, "/"))
}

// GetNamespace gets namespaces.
func (c *CoordinatorClient) GetNamespace() (admin.NamespaceGetResponse, error) {
	url := c.makeURL("api/v1/services/m3db/namespace")
	logger := c.logger.With(
		ZapMethod("getNamespace"), zap.String("url", url))

	//nolint:noctx
	resp, err := c.client.Get(url)
	if err != nil {
		logger.Error("failed get", zap.Error(err))
		return admin.NamespaceGetResponse{}, err
	}

	var response admin.NamespaceGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		return admin.NamespaceGetResponse{}, err
	}

	return response, nil
}

// GetPlacement gets placements.
func (c *CoordinatorClient) GetPlacement(opts PlacementRequestOptions) (admin.PlacementGetResponse, error) {
	var handlerurl string
	switch opts.Service {
	case ServiceTypeM3DB:
		handlerurl = placementhandler.M3DBGetURL
	case ServiceTypeM3Aggregator:
		handlerurl = placementhandler.M3AggGetURL
	case ServiceTypeM3Coordinator:
		handlerurl = placementhandler.M3CoordinatorGetURL
	default:
		return admin.PlacementGetResponse{}, errUnknownServiceType
	}
	url := c.makeURL(handlerurl)
	logger := c.logger.With(
		ZapMethod("getPlacement"), zap.String("url", url))

	resp, err := c.makeRequest(logger, url, placementhandler.GetHTTPMethod, nil, placementOptsToMap(opts))
	if err != nil {
		logger.Error("failed get", zap.Error(err))
		return admin.PlacementGetResponse{}, err
	}

	var response admin.PlacementGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		return admin.PlacementGetResponse{}, err
	}

	return response, nil
}

// InitPlacement initializes placements.
func (c *CoordinatorClient) InitPlacement(
	opts PlacementRequestOptions,
	initRequest admin.PlacementInitRequest,
) (admin.PlacementGetResponse, error) {
	var handlerurl string
	switch opts.Service {
	case ServiceTypeM3DB:
		handlerurl = placementhandler.M3DBInitURL
	case ServiceTypeM3Aggregator:
		handlerurl = placementhandler.M3AggInitURL
	case ServiceTypeM3Coordinator:
		handlerurl = placementhandler.M3CoordinatorInitURL
	default:
		return admin.PlacementGetResponse{}, errUnknownServiceType
	}
	url := c.makeURL(handlerurl)
	logger := c.logger.With(
		ZapMethod("initPlacement"), zap.String("url", url))

	resp, err := c.makeRequest(logger, url, placementhandler.InitHTTPMethod, &initRequest, placementOptsToMap(opts))
	if err != nil {
		logger.Error("failed init", zap.Error(err))
		return admin.PlacementGetResponse{}, err
	}

	var response admin.PlacementGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		return admin.PlacementGetResponse{}, err
	}

	return response, nil
}

// DeleteAllPlacements deletes all placements for the specified service.
func (c *CoordinatorClient) DeleteAllPlacements(opts PlacementRequestOptions) error {
	var handlerurl string
	switch opts.Service {
	case ServiceTypeM3DB:
		handlerurl = placementhandler.M3DBDeleteAllURL
	case ServiceTypeM3Aggregator:
		handlerurl = placementhandler.M3AggDeleteAllURL
	case ServiceTypeM3Coordinator:
		handlerurl = placementhandler.M3CoordinatorDeleteAllURL
	default:
		return errUnknownServiceType
	}
	url := c.makeURL(handlerurl)
	logger := c.logger.With(
		ZapMethod("deleteAllPlacements"), zap.String("url", url))

	resp, err := c.makeRequest(
		logger, url, placementhandler.DeleteAllHTTPMethod, nil, placementOptsToMap(opts),
	)
	if err != nil {
		logger.Error("failed to delete all placements", zap.Error(err))
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode/100 != 2 {
		logger.Error("status code not 2xx",
			zap.Int("status code", resp.StatusCode),
			zap.String("status", resp.Status))
		return fmt.Errorf("status code %d", resp.StatusCode)
	}

	logger.Info("placements deleted")

	return nil
}

// WaitForNamespace blocks until the given namespace is enabled.
// NB: if the name string is empty, this will instead
// check for a successful response.
func (c *CoordinatorClient) WaitForNamespace(name string) error {
	logger := c.logger.With(ZapMethod("waitForNamespace"))
	return c.retryFunc(func() error {
		ns, err := c.GetNamespace()
		if err != nil {
			return err
		}

		// If no name passed in, instad just check for success.
		if len(name) == 0 {
			return nil
		}

		nss := ns.GetRegistry().GetNamespaces()
		_, found := nss[name]
		if !found {
			err := fmt.Errorf("no namespace with name %s", name)
			logger.Error("could not get namespace", zap.Error(err))
			return err
		}

		logger.Info("namespace ready", zap.String("namespace", name))
		return nil
	})
}

// WaitForInstances blocks until the given instance is available.
func (c *CoordinatorClient) WaitForInstances(
	ids []string,
) error {
	logger := c.logger.With(ZapMethod("waitForPlacement"))
	return c.retryFunc(func() error {
		placement, err := c.GetPlacement(PlacementRequestOptions{Service: ServiceTypeM3DB})
		if err != nil {
			logger.Error("retrying get placement", zap.Error(err))
			return err
		}

		logger.Info("got placement", zap.Any("placement", placement))
		instances := placement.GetPlacement().GetInstances()
		for _, id := range ids {
			placement, found := instances[id]
			if !found {
				err = fmt.Errorf("no instance with id %s", id)
				logger.Error("could not get instance", zap.Error(err))
				return err
			}

			if pID := placement.GetId(); pID != id {
				err = fmt.Errorf("id mismatch: instance(%s) != placement(%s)", id, pID)
				logger.Error("could not get instance", zap.Error(err))
				return err
			}
		}

		logger.Info("instances ready")
		return nil
	})
}

// WaitForShardsReady waits until all shards gets ready.
func (c *CoordinatorClient) WaitForShardsReady() error {
	logger := c.logger.With(ZapMethod("waitForShards"))
	return c.retryFunc(func() error {
		placement, err := c.GetPlacement(PlacementRequestOptions{Service: ServiceTypeM3DB})
		if err != nil {
			logger.Error("retrying get placement", zap.Error(err))
			return err
		}

		for _, instance := range placement.Placement.Instances {
			for _, shard := range instance.Shards {
				if shard.State == placementpb.ShardState_INITIALIZING {
					err = fmt.Errorf("at least shard %d of dbnode %s still initializing", shard.Id, instance.Id)
					logger.Error("shards still are initializing", zap.Error(err))
					return err
				}
			}
		}
		return nil
	})
}

// WaitForClusterReady waits until the cluster is ready to receive reads and writes.
func (c *CoordinatorClient) WaitForClusterReady() error {
	var (
		url    = c.makeURL("ready")
		logger = c.logger.With(ZapMethod("waitForClusterReady"), zap.String("url", url))
	)
	return c.retryFunc(func() error {
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
		if err != nil {
			logger.Error("failed to create request", zap.Error(err))
			return err
		}

		resp, err := c.client.Do(req)
		if err != nil {
			logger.Error("failed checking cluster readiness", zap.Error(err))
			return err
		}
		defer resp.Body.Close()

		if resp.StatusCode != 200 {
			err = errors.New("non-200 status code received")

			body, rerr := ioutil.ReadAll(resp.Body)
			if rerr != nil {
				logger.Warn("failed parse response body", zap.Error(rerr))
				body = []byte("")
			}

			logger.Error("failed to check cluster readiness", zap.Error(err),
				zap.String("responseBody", string(body)),
			)
			return err
		}

		logger.Info("cluster ready to receive reads and writes")

		return nil
	})
}

// CreateDatabase creates a database.
func (c *CoordinatorClient) CreateDatabase(
	addRequest admin.DatabaseCreateRequest,
) (admin.DatabaseCreateResponse, error) {
	url := c.makeURL("api/v1/database/create")
	logger := c.logger.With(
		ZapMethod("createDatabase"), zap.String("url", url),
		zap.String("request", addRequest.String()))

	resp, err := c.makeRequest(logger, url, http.MethodPost, &addRequest, nil)
	if err != nil {
		logger.Error("failed post", zap.Error(err))
		return admin.DatabaseCreateResponse{}, err
	}

	var response admin.DatabaseCreateResponse
	if err := toResponse(resp, &response, logger); err != nil {
		logger.Error("failed response", zap.Error(err))
		return admin.DatabaseCreateResponse{}, err
	}

	if err = c.setNamespaceReady(addRequest.NamespaceName); err != nil {
		logger.Error("failed to set namespace to ready state",
			zap.Error(err),
			zap.String("namespace", addRequest.NamespaceName),
		)
		return response, err
	}

	logger.Info("created database")
	return response, nil
}

// AddNamespace adds a namespace.
func (c *CoordinatorClient) AddNamespace(
	addRequest admin.NamespaceAddRequest,
) (admin.NamespaceGetResponse, error) {
	url := c.makeURL("api/v1/services/m3db/namespace")
	logger := c.logger.With(
		ZapMethod("addNamespace"), zap.String("url", url),
		zap.String("request", addRequest.String()))

	resp, err := c.makeRequest(logger, url, http.MethodPost, &addRequest, nil)
	if err != nil {
		logger.Error("failed post", zap.Error(err))
		return admin.NamespaceGetResponse{}, err
	}

	var response admin.NamespaceGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		return admin.NamespaceGetResponse{}, err
	}

	if err = c.setNamespaceReady(addRequest.Name); err != nil {
		logger.Error("failed to set namespace to ready state", zap.Error(err), zap.String("namespace", addRequest.Name))
		return response, err
	}

	return response, nil
}

// UpdateNamespace updates the namespace.
func (c *CoordinatorClient) UpdateNamespace(
	req admin.NamespaceUpdateRequest,
) (admin.NamespaceGetResponse, error) {
	url := c.makeURL("api/v1/services/m3db/namespace")
	logger := c.logger.With(
		ZapMethod("updateNamespace"), zap.String("url", url),
		zap.String("request", req.String()))

	resp, err := c.makeRequest(logger, url, http.MethodPut, &req, nil)
	if err != nil {
		logger.Error("failed to update namespace", zap.Error(err))
		return admin.NamespaceGetResponse{}, err
	}

	var response admin.NamespaceGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		return admin.NamespaceGetResponse{}, err
	}

	return response, nil
}

func (c *CoordinatorClient) setNamespaceReady(name string) error {
	url := c.makeURL("api/v1/services/m3db/namespace/ready")
	logger := c.logger.With(
		ZapMethod("setNamespaceReady"), zap.String("url", url),
		zap.String("namespace", name))

	_, err := c.makeRequest(logger, url, http.MethodPost, // nolint: bodyclose
		&admin.NamespaceReadyRequest{
			Name:  name,
			Force: true,
		}, nil)
	return err
}

// DeleteNamespace removes the namespace.
func (c *CoordinatorClient) DeleteNamespace(namespaceID string) error {
	url := c.makeURL("api/v1/services/m3db/namespace/" + namespaceID)
	logger := c.logger.With(ZapMethod("deleteNamespace"), zap.String("url", url))

	if _, err := c.makeRequest(logger, url, http.MethodDelete, nil, nil); err != nil { // nolint: bodyclose
		logger.Error("failed to delete namespace", zap.Error(err))
		return err
	}
	return nil
}

//nolint:dupl
// InitM3msgTopic initializes an m3msg topic
func (c *CoordinatorClient) InitM3msgTopic(
	topicOpts M3msgTopicOptions,
	initRequest admin.TopicInitRequest,
) (admin.TopicGetResponse, error) {
	url := c.makeURL(topic.InitURL)
	logger := c.logger.With(
		ZapMethod("initM3msgTopic"),
		zap.String("url", url),
		zap.String("request", initRequest.String()),
		zap.String("topic", fmt.Sprintf("%v", topicOpts)))

	resp, err := c.makeRequest(logger, url, topic.InitHTTPMethod, &initRequest, m3msgTopicOptionsToMap(topicOpts))
	if err != nil {
		logger.Error("failed post", zap.Error(err))
		return admin.TopicGetResponse{}, err
	}

	var response admin.TopicGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		logger.Error("failed response", zap.Error(err))
		return admin.TopicGetResponse{}, err
	}

	logger.Info("topic initialized")
	return response, nil
}

// GetM3msgTopic fetches an m3msg topic
func (c *CoordinatorClient) GetM3msgTopic(
	topicOpts M3msgTopicOptions,
) (admin.TopicGetResponse, error) {
	url := c.makeURL(topic.GetURL)
	logger := c.logger.With(
		ZapMethod("getM3msgTopic"), zap.String("url", url),
		zap.String("topic", fmt.Sprintf("%v", topicOpts)))

	resp, err := c.makeRequest(logger, url, topic.GetHTTPMethod, nil, m3msgTopicOptionsToMap(topicOpts))
	if err != nil {
		logger.Error("failed get", zap.Error(err))
		return admin.TopicGetResponse{}, err
	}

	var response admin.TopicGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		logger.Error("failed response", zap.Error(err))
		return admin.TopicGetResponse{}, err
	}

	logger.Info("topic get")
	return response, nil
}

//nolint:dupl
// AddM3msgTopicConsumer adds a consumer service to an m3msg topic
func (c *CoordinatorClient) AddM3msgTopicConsumer(
	topicOpts M3msgTopicOptions,
	addRequest admin.TopicAddRequest,
) (admin.TopicGetResponse, error) {
	url := c.makeURL(topic.AddURL)
	logger := c.logger.With(
		ZapMethod("addM3msgTopicConsumer"),
		zap.String("url", url),
		zap.String("request", addRequest.String()),
		zap.String("topic", fmt.Sprintf("%v", topicOpts)))

	resp, err := c.makeRequest(logger, url, topic.AddHTTPMethod, &addRequest, m3msgTopicOptionsToMap(topicOpts))
	if err != nil {
		logger.Error("failed post", zap.Error(err))
		return admin.TopicGetResponse{}, err
	}

	var response admin.TopicGetResponse
	if err := toResponse(resp, &response, logger); err != nil {
		logger.Error("failed response", zap.Error(err))
		return admin.TopicGetResponse{}, err
	}

	logger.Info("topic consumer added")
	return response, nil
}

func placementOptsToMap(opts PlacementRequestOptions) map[string]string {
	return map[string]string{
		headers.HeaderClusterEnvironmentName: opts.Env,
		headers.HeaderClusterZoneName:        opts.Zone,
	}
}

func m3msgTopicOptionsToMap(opts M3msgTopicOptions) map[string]string {
	return map[string]string{
		headers.HeaderClusterEnvironmentName: opts.Env,
		headers.HeaderClusterZoneName:        opts.Zone,
		topic.HeaderTopicName:                opts.TopicName,
	}
}

// WriteCarbon writes a carbon metric datapoint at a given time.
func (c *CoordinatorClient) WriteCarbon(
	url string, metric string, v float64, t time.Time,
) error {
	logger := c.logger.With(
		ZapMethod("writeCarbon"), zap.String("url", url),
		zap.String("at time", time.Now().String()),
		zap.String("at ts", t.String()))

	con, err := net.Dial("tcp", url)
	if err != nil {
		logger.Error("could not dial", zap.Error(err))
		return err
	}

	write := fmt.Sprintf("%s %f %d", metric, v, t.Unix())
	logger.Info("writing", zap.String("metric", write))
	n, err := con.Write([]byte(write))
	if err != nil {
		logger.Error("could not write", zap.Error(err))
	}

	if n != len(write) {
		err := fmt.Errorf("wrote %d, wanted %d", n, len(write))
		logger.Error("write failure", zap.Error(err))
		return err
	}

	logger.Info("write success", zap.Int("bytes written", n))
	return con.Close()
}

// WriteProm writes a prometheus metric. Takes tags/labels as a map for convenience.
func (c *CoordinatorClient) WriteProm(
	name string,
	tags map[string]string,
	samples []prompb.Sample,
	headers Headers,
) error {
	labels := make([]prompb.Label, 0, len(tags)+1)

	labels = append(labels, prompb.Label{
		Name:  []byte(model.MetricNameLabel),
		Value: []byte(name),
	})

	for tag, value := range tags {
		labels = append(labels, prompb.Label{
			Name:  []byte(tag),
			Value: []byte(value),
		})
	}

	writeRequest := prompb.WriteRequest{
		Timeseries: []prompb.TimeSeries{
			{
				Labels:  labels,
				Samples: samples,
			},
		},
	}

	return c.WritePromWithRequest(writeRequest, headers)
}

// WritePromWithRequest executes a prometheus write request. Allows you to
// provide the request directly which is useful for batch metric requests.
func (c *CoordinatorClient) WritePromWithRequest(writeRequest prompb.WriteRequest, headers Headers) error {
	url := c.makeURL("api/v1/prom/remote/write")

	logger := c.logger.With(
		ZapMethod("writeProm"), zap.String("url", url),
		zap.String("request", writeRequest.String()))

	body, err := proto.Marshal(&writeRequest)
	if err != nil {
		logger.Error("failed marshaling request message", zap.Error(err))
		return err
	}
	data := bytes.NewBuffer(snappy.Encode(nil, body))

	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, url, data)
	if err != nil {
		logger.Error("failed constructing request", zap.Error(err))
		return err
	}
	for key, vals := range headers {
		for _, val := range vals {
			req.Header.Add(key, val)
		}
	}
	req.Header.Add(xhttp.HeaderContentType, xhttp.ContentTypeProtobuf)

	resp, err := c.client.Do(req)
	if err != nil {
		logger.Error("failed making a request", zap.Error(err))
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		logger.Error("status code not 2xx",
			zap.Int("status code", resp.StatusCode),
			zap.String("status", resp.Status))
		return fmt.Errorf("status code %d", resp.StatusCode)
	}

	return nil
}

func (c *CoordinatorClient) makeRequest(
	logger *zap.Logger,
	url string,
	method string,
	body proto.Message,
	header map[string]string,
) (*http.Response, error) {
	data := bytes.NewBuffer(nil)
	if body != nil {
		if err := (&jsonpb.Marshaler{}).Marshal(data, body); err != nil {
			logger.Error("failed to marshal", zap.Error(err))

			return nil, fmt.Errorf("failed to marshal: %w", err)
		}
	}

	req, err := http.NewRequestWithContext(context.Background(), method, url, data)
	if err != nil {
		logger.Error("failed to construct request", zap.Error(err))

		return nil, fmt.Errorf("failed to construct request: %w", err)
	}

	req.Header.Add(xhttp.HeaderContentType, xhttp.ContentTypeJSON)
	for k, v := range header {
		req.Header.Add(k, v)
	}

	return c.client.Do(req)
}

// ApplyKVUpdate applies a KV update.
func (c *CoordinatorClient) ApplyKVUpdate(update string) error {
	url := c.makeURL("api/v1/kvstore")

	logger := c.logger.With(
		ZapMethod("ApplyKVUpdate"), zap.String("url", url),
		zap.String("update", update))

	data := bytes.NewBuffer([]byte(update))
	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, url, data)
	if err != nil {
		logger.Error("failed to construct request", zap.Error(err))
		return fmt.Errorf("failed to construct request: %w", err)
	}

	req.Header.Add(xhttp.HeaderContentType, xhttp.ContentTypeJSON)

	resp, err := c.client.Do(req)
	if err != nil {
		logger.Error("failed to apply request", zap.Error(err))
		return fmt.Errorf("failed to apply request: %w", err)
	}

	bs, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logger.Error("failed to read body", zap.Error(err))
		return fmt.Errorf("failed to read body: %w", err)
	}

	logger.Info("applied KV update", zap.ByteString("response", bs))
	_ = resp.Body.Close()
	return nil
}

func (c *CoordinatorClient) query(
	verifier ResponseVerifier, query string, headers map[string][]string,
) error {
	url := c.makeURL(query)
	logger := c.logger.With(
		ZapMethod("query"), zap.String("url", url), zap.Any("headers", headers))
	logger.Info("running")
	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
	if err != nil {
		return err
	}

	if headers != nil {
		req.Header = headers
	}

	resp, err := c.client.Do(req)
	if err != nil {
		logger.Error("failed get", zap.Error(err))
		return err
	}

	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)

	return verifier(resp.StatusCode, resp.Header, string(b), err)
}

// InstantQuery runs an instant query with provided headers
func (c *CoordinatorClient) InstantQuery(req QueryRequest, headers Headers) (model.Vector, error) {
	return c.instantQuery(req, route.QueryURL, headers)
}

// InstantQueryWithEngine runs an instant query with provided headers and the specified
// query engine.
func (c *CoordinatorClient) InstantQueryWithEngine(
	req QueryRequest,
	engine options.QueryEngine,
	headers Headers,
) (model.Vector, error) {
	if engine == options.M3QueryEngine {
		return c.instantQuery(req, native.M3QueryReadInstantURL, headers)
	} else if engine == options.PrometheusEngine {
		return c.instantQuery(req, native.PrometheusReadInstantURL, headers)
	}
	return nil, fmt.Errorf("unknown query engine: %s", engine)
}

func (c *CoordinatorClient) instantQuery(
	req QueryRequest,
	queryRoute string,
	headers Headers,
) (model.Vector, error) {
	queryStr := fmt.Sprintf("%s?query=%s", queryRoute, req.Query)
	if req.Time != nil {
		queryStr = fmt.Sprintf("%s&time=%d", queryStr, req.Time.Unix())
	}

	resp, err := c.runQuery(queryStr, headers)
	if err != nil {
		return nil, err
	}

	var parsedResp jsonInstantQueryResponse
	if err := json.Unmarshal([]byte(resp), &parsedResp); err != nil {
		return nil, err
	}

	return parsedResp.Data.Result, nil
}

type jsonInstantQueryResponse struct {
	Status string
	Data   vectorResult
}

type vectorResult struct {
	ResultType model.ValueType
	Result     model.Vector
}

// RangeQuery runs a range query with provided headers
func (c *CoordinatorClient) RangeQuery(
	req RangeQueryRequest,
	headers Headers,
) (model.Matrix, error) {
	return c.rangeQuery(req, route.QueryRangeURL, headers)
}

// RangeQueryWithEngine runs a range query with provided headers and the specified
// query engine.
func (c *CoordinatorClient) RangeQueryWithEngine(
	req RangeQueryRequest,
	engine options.QueryEngine,
	headers Headers,
) (model.Matrix, error) {
	if engine == options.M3QueryEngine {
		return c.rangeQuery(req, native.M3QueryReadURL, headers)
	} else if engine == options.PrometheusEngine {
		return c.rangeQuery(req, native.PrometheusReadURL, headers)
	}
	return nil, fmt.Errorf("unknown query engine: %s", engine)
}

func (c *CoordinatorClient) rangeQuery(
	req RangeQueryRequest,
	queryRoute string,
	headers Headers,
) (model.Matrix, error) {
	if req.Step == 0 {
		req.Step = 15 * time.Second // default step is 15 seconds.
	}
	queryStr := fmt.Sprintf(
		"%s?query=%s&start=%d&end=%d&step=%f",
		queryRoute, req.Query,
		req.Start.Unix(),
		req.End.Unix(),
		req.Step.Seconds(),
	)

	resp, err := c.runQuery(queryStr, headers)
	if err != nil {
		return nil, err
	}

	var parsedResp jsonRangeQueryResponse
	if err := json.Unmarshal([]byte(resp), &parsedResp); err != nil {
		return nil, err
	}

	return parsedResp.Data.Result, nil
}

// LabelNames return matching label names based on the request.
func (c *CoordinatorClient) LabelNames(
	req LabelNamesRequest,
	headers Headers,
) (model.LabelNames, error) {
	urlPathAndQuery := fmt.Sprintf("%s?%s", route.LabelNamesURL, req.String())
	resp, err := c.runQuery(urlPathAndQuery, headers)
	if err != nil {
		return nil, err
	}

	var parsedResp labelResponse
	if err := json.Unmarshal([]byte(resp), &parsedResp); err != nil {
		return nil, err
	}

	labelNames := make(model.LabelNames, 0, len(parsedResp.Data))
	for _, label := range parsedResp.Data {
		labelNames = append(labelNames, model.LabelName(label))
	}

	return labelNames, nil
}

// LabelValues return matching label values based on the request.
func (c *CoordinatorClient) LabelValues(
	req LabelValuesRequest,
	headers Headers,
) (model.LabelValues, error) {
	urlPathAndQuery := fmt.Sprintf("%s?%s",
		path.Join(route.Prefix, "label", req.LabelName, "values"),
		req.String())
	resp, err := c.runQuery(urlPathAndQuery, headers)
	if err != nil {
		return nil, err
	}

	var parsedResp labelResponse
	if err := json.Unmarshal([]byte(resp), &parsedResp); err != nil {
		return nil, err
	}

	labelValues := make(model.LabelValues, 0, len(parsedResp.Data))
	for _, label := range parsedResp.Data {
		labelValues = append(labelValues, model.LabelValue(label))
	}

	return labelValues, nil
}

// Series returns matching series based on the request.
func (c *CoordinatorClient) Series(
	req SeriesRequest,
	headers Headers,
) ([]model.Metric, error) {
	urlPathAndQuery := fmt.Sprintf("%s?%s", route.SeriesMatchURL, req.String())
	resp, err := c.runQuery(urlPathAndQuery, headers)
	if err != nil {
		return nil, err
	}

	var parsedResp seriesResponse
	if err := json.Unmarshal([]byte(resp), &parsedResp); err != nil {
		return nil, err
	}

	series := make([]model.Metric, 0, len(parsedResp.Data))
	for _, labels := range parsedResp.Data {
		labelSet := make(model.LabelSet)
		for name, val := range labels {
			labelSet[model.LabelName(name)] = model.LabelValue(val)
		}
		series = append(series, model.Metric(labelSet))
	}

	return series, nil
}

type jsonRangeQueryResponse struct {
	Status string
	Data   matrixResult
}

type matrixResult struct {
	ResultType model.ValueType
	Result     model.Matrix
}

type labelResponse struct {
	Status string
	Data   []string
}

type seriesResponse struct {
	Status string
	Data   []map[string]string
}

func (c *CoordinatorClient) runQuery(
	query string, headers map[string][]string,
) (string, error) {
	url := c.makeURL(query)
	logger := c.logger.With(
		ZapMethod("query"), zap.String("url", url), zap.Any("headers", headers))
	logger.Info("running")
	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
	if err != nil {
		return "", err
	}

	if headers != nil {
		req.Header = headers
	}

	resp, err := c.client.Do(req)
	if err != nil {
		logger.Error("failed get", zap.Error(err))
		return "", err
	}

	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)

	if status := resp.StatusCode; status != http.StatusOK {
		return "", fmt.Errorf("query response status not OK, received %v. error=%v",
			status, string(b))
	}

	if contentType, ok := resp.Header["Content-Type"]; !ok {
		return "", fmt.Errorf("missing Content-Type header")
	} else if len(contentType) != 1 || contentType[0] != "application/json" { //nolint:goconst
		return "", fmt.Errorf("expected json content type, got %v", contentType)
	}

	return string(b), err
}

// RunQuery runs the given query with a given verification function.
func (c *CoordinatorClient) RunQuery(
	verifier ResponseVerifier, query string, headers map[string][]string,
) error {
	logger := c.logger.With(ZapMethod("runQuery"),
		zap.String("query", query))
	err := c.retryFunc(func() error {
		err := c.query(verifier, query, headers)
		if err != nil {
			logger.Info("retrying", zap.Error(err))
		}

		return err
	})
	if err != nil {
		logger.Error("failed run", zap.Error(err))
	}

	return err
}

func toResponse(
	resp *http.Response,
	response proto.Message,
	logger *zap.Logger,
) error {
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logger.Error("could not read body", zap.Error(err))
		return err
	}

	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		logger.Error("status code not 2xx",
			zap.Int("status code", resp.StatusCode),
			zap.String("status", resp.Status))
		return fmt.Errorf("status code %d", resp.StatusCode)
	}

	err = jsonpb.Unmarshal(bytes.NewReader(b), response)
	if err != nil {
		logger.Error("unable to unmarshal response",
			zap.Error(err),
			zap.Any("response", response))
		return err
	}

	return nil
}

// GraphiteQuery retrieves graphite raw data.
func (c *CoordinatorClient) GraphiteQuery(
	graphiteReq GraphiteQueryRequest,
) ([]Datapoint, error) {
	if graphiteReq.From.IsZero() {
		graphiteReq.From = time.Now().Add(-24 * time.Hour)
	}
	if graphiteReq.Until.IsZero() {
		graphiteReq.Until = time.Now()
	}

	queryStr := fmt.Sprintf(
		"%s?target=%s&from=%d&until=%d",
		graphite.ReadURL, graphiteReq.Target,
		graphiteReq.From.Unix(),
		graphiteReq.Until.Unix(),
	)

	url := c.makeURL(queryStr)
	logger := c.logger.With(
		ZapMethod("graphiteQuery"), zap.String("url", url))
	logger.Info("running")
	req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}

	resp, err := c.client.Do(req)
	if err != nil {
		logger.Error("failed get", zap.Error(err))
		return nil, err
	}

	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if status := resp.StatusCode; status != http.StatusOK {
		return nil, fmt.Errorf("query response status not OK, received %v %s", status, resp.Status)
	}

	var parsedResp jsonGraphiteQueryResponse
	if err := json.Unmarshal(b, &parsedResp); err != nil {
		return nil, err
	}

	if len(parsedResp) == 0 {
		return nil, nil
	}

	results := make([]Datapoint, 0, len(parsedResp[0].Datapoints))
	for _, dp := range parsedResp[0].Datapoints {
		if len(dp) != 2 {
			return nil, fmt.Errorf("failed to parse response: %s", string(b))
		}

		results = append(results, Datapoint{
			Value:     dp[0],
			Timestamp: int64(*dp[1]),
		})
	}

	return results, nil
}

type jsonGraphiteQueryResponse []series

type series struct {
	Target     string
	Datapoints []tuple
	StepSizeMs int
}

type tuple []*float64


<|start_filename|>src/dbnode/storage/entry.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package storage

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/m3db/m3/src/dbnode/storage/block"
	"github.com/m3db/m3/src/dbnode/storage/bootstrap"
	"github.com/m3db/m3/src/dbnode/storage/index"
	"github.com/m3db/m3/src/dbnode/storage/series"
	"github.com/m3db/m3/src/dbnode/ts/writes"
	"github.com/m3db/m3/src/m3ninx/doc"
	"github.com/m3db/m3/src/x/clock"
	"github.com/m3db/m3/src/x/context"
	m3errors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/ident"
	"github.com/m3db/m3/src/x/resource"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/uber-go/tally"
	xatomic "go.uber.org/atomic"
)

// IndexWriter accepts index inserts.
type IndexWriter interface {
	// WritePending indexes the provided pending entries.
	WritePending(
		pending []writes.PendingIndexInsert,
	) error

	// BlockStartForWriteTime returns the index block start
	// time for the given writeTime.
	BlockStartForWriteTime(
		writeTime xtime.UnixNano,
	) xtime.UnixNano
}

// EntryMetrics are metrics for an entry.
type EntryMetrics struct {
	gcNoReconcile           tally.Counter
	gcNeedsReconcile        tally.Counter
	gcSuccessShardClosed    tally.Counter
	gcSuccessEmpty          tally.Counter
	noGcNil                 tally.Counter
	noGcErr                 tally.Counter
	noGcHasReaders          tally.Counter
	noGcNotEmptySeries      tally.Counter
	duplicateNoReconcile    tally.Counter
	duplicateNeedsReconcile tally.Counter
}

// NewEntryMetrics builds an entry metrics.
func NewEntryMetrics(scope tally.Scope) *EntryMetrics {
	return &EntryMetrics{
		gcNoReconcile: scope.Tagged(map[string]string{
			"reconcile": "no_reconcile",
			"path":      "gc",
		}).Counter("count"),
		gcNeedsReconcile: scope.Tagged(map[string]string{
			"reconcile": "needs_reconcile",
			"path":      "gc",
		}).Counter("count"),
		gcSuccessShardClosed: scope.Tagged(map[string]string{
			"reason": "shard_closed",
			"path":   "gc",
		}).Counter("gc_count"),
		gcSuccessEmpty: scope.Tagged(map[string]string{
			"reason": "empty",
			"path":   "gc",
		}).Counter("gc_count"),
		noGcNil: scope.Tagged(map[string]string{
			"reason": "nil",
			"path":   "gc",
		}).Counter("no_gc_count"),
		noGcErr: scope.Tagged(map[string]string{
			"reason": "error",
			"path":   "gc",
		}).Counter("no_gc_count"),
		noGcHasReaders: scope.Tagged(map[string]string{
			"reason": "has_readers",
			"path":   "gc",
		}).Counter("no_gc_count"),
		noGcNotEmptySeries: scope.Tagged(map[string]string{
			"reason": "not_empty_series",
			"path":   "gc",
		}).Counter("no_gc_count"),

		duplicateNoReconcile: scope.Tagged(map[string]string{
			"reconcile": "no_reconcile",
			"path":      "duplicate",
		}).Counter("count"),
		duplicateNeedsReconcile: scope.Tagged(map[string]string{
			"reconcile": "needs_reconcile",
			"path":      "duplicate",
		}).Counter("count"),
	}
}

// Entry is the entry in the shard ident.ID -> series map. It has additional
// members to track lifecycle and minimize indexing overhead.
// NB: users are expected to use `NewEntry` to construct these objects.
type Entry struct {
	ID                       ident.ID
	Shard                    Shard
	Series                   series.DatabaseSeries
	Index                    uint64
	IndexGarbageCollected    *xatomic.Bool
	insertTime               *xatomic.Int64
	indexWriter              IndexWriter
	curReadWriters           int32
	reverseIndex             entryIndexState
	nowFn                    clock.NowFn
	metrics                  *EntryMetrics
	pendingIndexBatchSizeOne []writes.PendingIndexInsert
}

// ensure Entry satisfies the `doc.OnIndexSeries` interface.
var _ doc.OnIndexSeries = &Entry{}

// ensure Entry satisfies the `bootstrap.SeriesRef` interface.
var _ bootstrap.SeriesRef = &Entry{}

// ensure Entry satisfies the `bootstrap.SeriesRefResolver` interface.
var _ bootstrap.SeriesRefResolver = &Entry{}

// NewEntryOptions supplies options for a new entry.
type NewEntryOptions struct {
	Shard        Shard
	Series       series.DatabaseSeries
	Index        uint64
	IndexWriter  IndexWriter
	NowFn        clock.NowFn
	EntryMetrics *EntryMetrics
}

// NewEntry returns a new Entry.
func NewEntry(opts NewEntryOptions) *Entry {
	nowFn := time.Now
	if opts.NowFn != nil {
		nowFn = opts.NowFn
	}
	entry := &Entry{
		ID:                       opts.Series.ID(),
		Shard:                    opts.Shard,
		Series:                   opts.Series,
		Index:                    opts.Index,
		IndexGarbageCollected:    xatomic.NewBool(false),
		insertTime:               xatomic.NewInt64(0),
		indexWriter:              opts.IndexWriter,
		nowFn:                    nowFn,
		pendingIndexBatchSizeOne: make([]writes.PendingIndexInsert, 1),
		reverseIndex:             newEntryIndexState(),
		metrics:                  opts.EntryMetrics,
	}
	return entry
}

// StringID returns the index series ID, as a string.
func (entry *Entry) StringID() string {
	return entry.ID.String()
}

// ReaderWriterCount returns the current ref count on the Entry.
func (entry *Entry) ReaderWriterCount() int32 {
	return atomic.LoadInt32(&entry.curReadWriters)
}

// IncrementReaderWriterCount increments the ref count on the Entry.
func (entry *Entry) IncrementReaderWriterCount() {
	atomic.AddInt32(&entry.curReadWriters, 1)
}

// DecrementReaderWriterCount decrements the ref count on the Entry.
func (entry *Entry) DecrementReaderWriterCount() {
	atomic.AddInt32(&entry.curReadWriters, -1)
}

// IndexedBlockCount returns the count of indexed block states.
func (entry *Entry) IndexedBlockCount() int {
	entry.reverseIndex.RLock()
	count := len(entry.reverseIndex.states)
	entry.reverseIndex.RUnlock()
	return count
}

// IndexedForBlockStart returns a bool to indicate if the Entry has been successfully
// indexed for the given index blockStart.
func (entry *Entry) IndexedForBlockStart(indexBlockStart xtime.UnixNano) bool {
	entry.reverseIndex.RLock()
	isIndexed := entry.reverseIndex.indexedWithRLock(indexBlockStart)
	entry.reverseIndex.RUnlock()
	return isIndexed
}

// IndexedRange returns minimum and maximum blockStart values covered by index entry.
// The range is inclusive. Note that there may be uncovered gaps within the range.
// Returns (0, 0) for an empty range.
func (entry *Entry) IndexedRange() (xtime.UnixNano, xtime.UnixNano) {
	entry.reverseIndex.RLock()
	min, max := entry.reverseIndex.indexedRangeWithRLock()
	entry.reverseIndex.RUnlock()
	return min, max
}

// ReconciledOnIndexSeries attempts to retrieve the most recent index entry from the
// shard if the entry this method was called on was never inserted there. If there
// is an error during retrieval, simply returns the current entry. Additionally,
// returns a cleanup function to run once finished using the reconciled entry and
// a boolean value indicating whether the result came from reconciliation or not.
func (entry *Entry) ReconciledOnIndexSeries() (doc.OnIndexSeries, resource.SimpleCloser, bool) {
	if entry.insertTime.Load() > 0 {
		return entry, resource.SimpleCloserFn(func() {}), false
	}

	e, _, err := entry.Shard.TryRetrieveSeriesAndIncrementReaderWriterCount(entry.ID)
	if err != nil || e == nil {
		return entry, resource.SimpleCloserFn(func() {}), false
	}

	// NB: attempt to merge the index series here, to ensure the returned
	// reconciled series will have each index block marked from both this and the
	// reconciliated series.
	entry.mergeInto(e)

	return e, resource.SimpleCloserFn(func() {
		e.DecrementReaderWriterCount()
	}), true
}

// MergeEntryIndexBlockStates merges the given states into the current
// indexed entry.
func (entry *Entry) MergeEntryIndexBlockStates(states doc.EntryIndexBlockStates) {
	entry.reverseIndex.Lock()
	for t, state := range states {
		set := false
		if state.Success {
			set = true
			entry.reverseIndex.setSuccessWithWLock(t)
		} else {
			// NB: setSuccessWithWLock(t) will perform the logic to determine if
			// minIndexedT/maxIndexedT need to be updated; if this is not being called
			// these should be updated.
			if entry.reverseIndex.maxIndexedT < t {
				entry.reverseIndex.maxIndexedT = t
			}
			if entry.reverseIndex.minIndexedT > t {
				entry.reverseIndex.minIndexedT = t
			}
		}

		if state.Attempt {
			set = true
			entry.reverseIndex.setAttemptWithWLock(t, false)
		}

		if !set {
			// NB: if not set through the above methods, need to create an index block
			// state at the given timestamp.
			entry.reverseIndex.states[t] = doc.EntryIndexBlockState{}
		}
	}

	entry.reverseIndex.Unlock()
}

// NeedsIndexUpdate returns a bool to indicate if the Entry needs to be indexed
// for the provided blockStart. It only allows a single index attempt at a time
// for a single entry.
// NB(prateek): NeedsIndexUpdate is a CAS, i.e. when this method returns true, it
// also sets state on the entry to indicate that a write for the given blockStart
// is going to be sent to the index, and other go routines should not attempt the
// same write. Callers are expected to ensure they follow this guideline.
// Further, every call to NeedsIndexUpdate which returns true needs to have a corresponding
// OnIndexFinalize() call. This is required for correct lifecycle maintenance.
func (entry *Entry) NeedsIndexUpdate(indexBlockStartForWrite xtime.UnixNano) bool {
	// first we try the low-cost path: acquire a RLock and see if the given block start
	// has been marked successful or that we've attempted it.
	entry.reverseIndex.RLock()
	alreadyIndexedOrAttempted := entry.reverseIndex.indexedOrAttemptedWithRLock(indexBlockStartForWrite)
	entry.reverseIndex.RUnlock()
	if alreadyIndexedOrAttempted {
		// if so, the entry does not need to be indexed.
		return false
	}

	// now acquire a write lock and set that we're going to attempt to do this so we don't try
	// multiple times.
	entry.reverseIndex.Lock()
	// NB(prateek): not defer-ing here, need to avoid the the extra ~150ns to minimize contention.

	// but first, we have to ensure no one has done so since we released the read lock
	alreadyIndexedOrAttempted = entry.reverseIndex.indexedOrAttemptedWithRLock(indexBlockStartForWrite)
	if alreadyIndexedOrAttempted {
		entry.reverseIndex.Unlock()
		return false
	}

	entry.reverseIndex.setAttemptWithWLock(indexBlockStartForWrite, true)
	entry.reverseIndex.Unlock()
	return true
}

// OnIndexPrepare prepares the Entry to be handed off to the indexing sub-system.
// NB(prateek): we retain the ref count on the entry while the indexing is pending,
// the callback executed on the entry once the indexing is completed releases this
// reference.
func (entry *Entry) OnIndexPrepare(blockStartNanos xtime.UnixNano) {
	entry.reverseIndex.Lock()
	entry.reverseIndex.setAttemptWithWLock(blockStartNanos, true)
	entry.reverseIndex.Unlock()
	entry.IncrementReaderWriterCount()
}

// OnIndexSuccess marks the given block start as successfully indexed.
func (entry *Entry) OnIndexSuccess(blockStartNanos xtime.UnixNano) {
	entry.reverseIndex.Lock()
	entry.reverseIndex.setSuccessWithWLock(blockStartNanos)
	entry.reverseIndex.Unlock()
}

// OnIndexFinalize marks any attempt for the given block start as finished
// and decrements the entry ref count.
func (entry *Entry) OnIndexFinalize(blockStartNanos xtime.UnixNano) {
	entry.reverseIndex.Lock()
	entry.reverseIndex.setAttemptWithWLock(blockStartNanos, false)
	entry.reverseIndex.Unlock()
	// indicate the index has released held reference for provided write
	entry.DecrementReaderWriterCount()
}

// IfAlreadyIndexedMarkIndexSuccessAndFinalize marks the entry as successfully
// indexed if already indexed and returns true. Otherwise returns false.
func (entry *Entry) IfAlreadyIndexedMarkIndexSuccessAndFinalize(
	blockStart xtime.UnixNano,
) bool {
	successAlready := false
	entry.reverseIndex.Lock()
	for _, state := range entry.reverseIndex.states {
		if state.Success {
			successAlready = true
			break
		}
	}
	if successAlready {
		entry.reverseIndex.setSuccessWithWLock(blockStart)
		entry.reverseIndex.setAttemptWithWLock(blockStart, false)
	}
	entry.reverseIndex.Unlock()
	if successAlready {
		// indicate the index has released held reference for provided write
		entry.DecrementReaderWriterCount()
	}
	return successAlready
}

// TryMarkIndexGarbageCollected checks if the entry is eligible to be garbage collected
// from the index. If so, it marks the entry as GCed and returns true. Otherwise returns false.
func (entry *Entry) TryMarkIndexGarbageCollected() bool {
	// Since series insertions + index insertions are done separately async, it is possible for
	// a series to be in the index but not have data written yet, and so any series not in the
	// lookup yet we cannot yet consider empty.
	e, _, err := entry.Shard.TryRetrieveSeriesAndIncrementReaderWriterCount(entry.ID)
	if m3errors.Is(err, errShardNotOpen) {
		// Shard is closing, all entries which belonged to it should be gc'ed.
		entry.metrics.gcSuccessShardClosed.Inc(1)
		entry.IndexGarbageCollected.Store(true)
		return true
	}

	if err != nil {
		entry.metrics.noGcErr.Inc(1)
		return false
	}

	if e == nil {
		entry.metrics.noGcNil.Inc(1)
		return false
	}

	defer e.DecrementReaderWriterCount()

	// Was reconciled if the entry retrieved from the shard differs from the current.
	if e != entry {
		// If this entry needs further reconciliation, merge this entry's index
		// states into the
		entry.reverseIndex.RLock()
		e.MergeEntryIndexBlockStates(entry.reverseIndex.states)
		entry.reverseIndex.RUnlock()
	}

	// Consider non-empty if the entry is still being held since this could indicate
	// another thread holding a new series prior to writing to it.
	if e.ReaderWriterCount() > 1 {
		entry.metrics.noGcHasReaders.Inc(1)
		return false
	}

	// Series must be empty to be GCed. This happens when the data and index are flushed to disk and
	// so the series no longer has in-mem data.
	if !e.Series.IsEmpty() {
		entry.metrics.noGcNotEmptySeries.Inc(1)
		return false
	}

	// Mark as GCed from index so the entry can be safely cleaned up in the shard.
	// The reference to this entry from the index is removed by the code path that
	// marks this GCed bool.
	e.metrics.gcSuccessEmpty.Inc(1)
	e.IndexGarbageCollected.Store(true)

	if e != entry {
		entry.metrics.gcNeedsReconcile.Inc(1)
	} else {
		entry.metrics.gcNoReconcile.Inc(1)
	}

	return true
}

// mergeInto merges this entry index blocks into the provided index series.
func (entry *Entry) mergeInto(indexSeries doc.OnIndexSeries) {
	if entry == indexSeries {
		// NB: short circuit if attempting to merge an entry into itself.
		return
	}

	entry.reverseIndex.RLock()
	indexSeries.MergeEntryIndexBlockStates(entry.reverseIndex.states)
	entry.reverseIndex.RUnlock()
}

// TryReconcileDuplicates attempts to reconcile the index states of this entry.
func (entry *Entry) TryReconcileDuplicates() {
	// Since series insertions + index insertions are done separately async, it is possible for
	// a series to be in the index but not have data written yet, and so any series not in the
	// lookup yet we cannot yet consider empty.
	e, _, err := entry.Shard.TryRetrieveSeriesAndIncrementReaderWriterCount(entry.ID)
	if err != nil || e == nil {
		return
	}

	if e != entry {
		entry.mergeInto(e)
		entry.metrics.duplicateNeedsReconcile.Inc(1)
	} else {
		entry.metrics.duplicateNoReconcile.Inc(1)
	}

	e.DecrementReaderWriterCount()
}

// NeedsIndexGarbageCollected checks if the entry is eligible to be garbage collected
// from the index. Otherwise returns false.
func (entry *Entry) NeedsIndexGarbageCollected() bool {
	// This is a cheaper check that loading the entry from the shard again
	// which makes it cheaper to run frequently.
	// It may not be as accurate, but it's fine for an approximation since
	// only a single series in a segment needs to return true to trigger an
	// index segment to be garbage collected.
	if entry.insertTime.Load() == 0 {
		return false // Not inserted, does not need garbage collection.
	}

	// NB(antanas): Entries need to be GC'ed for closed shards.
	// Orphan entries will cause problems if same shard returns to the same node.
	if entry.Shard.Closed() {
		return true
	}
	// Check that a write is not potentially pending and the series is empty.
	return entry.ReaderWriterCount() == 0 && entry.Series.IsEmpty()
}

// SetInsertTime marks the entry as having been inserted into the shard at a given timestamp.
func (entry *Entry) SetInsertTime(t time.Time) {
	entry.insertTime.Store(t.UnixNano())
}

// Write writes a new value.
func (entry *Entry) Write(
	ctx context.Context,
	timestamp xtime.UnixNano,
	value float64,
	unit xtime.Unit,
	annotation []byte,
	wOpts series.WriteOptions,
) (bool, series.WriteType, error) {
	if err := entry.maybeIndex(timestamp); err != nil {
		return false, 0, err
	}
	return entry.Series.Write(
		ctx,
		timestamp,
		value,
		unit,
		annotation,
		wOpts,
	)
}

// LoadBlock loads a single block into the series.
func (entry *Entry) LoadBlock(
	block block.DatabaseBlock,
	writeType series.WriteType,
) error {
	// TODO(bodu): We can remove this once we have index snapshotting as index snapshots will
	// contained snapshotted index segments that cover snapshotted data.
	if err := entry.maybeIndex(block.StartTime()); err != nil {
		return err
	}
	return entry.Series.LoadBlock(block, writeType)
}

// UniqueIndex is the unique index for the series.
func (entry *Entry) UniqueIndex() uint64 {
	return entry.Series.UniqueIndex()
}

func (entry *Entry) maybeIndex(timestamp xtime.UnixNano) error {
	idx := entry.indexWriter
	if idx == nil {
		return nil
	}
	if !entry.NeedsIndexUpdate(idx.BlockStartForWriteTime(timestamp)) {
		return nil
	}
	entry.pendingIndexBatchSizeOne[0] = writes.PendingIndexInsert{
		Entry: index.WriteBatchEntry{
			Timestamp:     timestamp,
			OnIndexSeries: entry,
			EnqueuedAt:    entry.nowFn(),
		},
		Document: entry.Series.Metadata(),
	}
	entry.OnIndexPrepare(idx.BlockStartForWriteTime(timestamp))
	return idx.WritePending(entry.pendingIndexBatchSizeOne)
}

// SeriesRef returns the series read write ref.
func (entry *Entry) SeriesRef() (bootstrap.SeriesRef, error) {
	return entry, nil
}

// ReleaseRef must be called after using the series ref
// to release the reference count to the series so it can
// be expired by the owning shard eventually.
func (entry *Entry) ReleaseRef() {
	entry.DecrementReaderWriterCount()
}

// entryIndexState is used to capture the state of indexing for a single shard
// entry. It's used to prevent redundant indexing operations.
// NB(prateek): We need this amount of state because in the worst case, as we can have 3 active blocks being
// written to. Albeit that's an edge case due to bad configuration. Even outside of that, 2 blocks can
// be written to due to delayed, out of order writes. Consider an index block size of 2h, and buffer
// past of 10m. Say a write comes in at 2.05p (wallclock) for 2.05p (timestamp in the write), we'd index
// the entry, and update the entry to have a success for 4p. Now imagine another write
// comes in at 2.06p (wallclock) for 1.57p (timestamp in the write). We need to differentiate that we don't
// have a write for the 12-2p block from the 2-4p block, or we'd drop the late write.
type entryIndexState struct {
	sync.RWMutex
	states                   doc.EntryIndexBlockStates
	minIndexedT, maxIndexedT xtime.UnixNano
}

func newEntryIndexState() entryIndexState {
	return entryIndexState{
		states: make(doc.EntryIndexBlockStates, 4),
	}
}

func (s *entryIndexState) indexedRangeWithRLock() (xtime.UnixNano, xtime.UnixNano) {
	return s.minIndexedT, s.maxIndexedT
}

func (s *entryIndexState) indexedWithRLock(t xtime.UnixNano) bool {
	v, ok := s.states[t]
	if ok {
		return v.Success
	}
	return false
}

func (s *entryIndexState) indexedOrAttemptedWithRLock(t xtime.UnixNano) bool {
	v, ok := s.states[t]
	if ok {
		return v.Success || v.Attempt
	}
	return false
}

func (s *entryIndexState) setSuccessWithWLock(t xtime.UnixNano) {
	if s.indexedWithRLock(t) {
		return
	}

	// NB(r): If not inserted state yet that means we need to make an insertion,
	// this will happen if synchronously indexing and we haven't called
	// NeedIndexUpdate before we indexed the series.
	s.states[t] = doc.EntryIndexBlockState{
		Success: true,
	}

	if t > s.maxIndexedT {
		s.maxIndexedT = t
	}
	if t < s.minIndexedT || s.minIndexedT == 0 {
		s.minIndexedT = t
	}
}

func (s *entryIndexState) setAttemptWithWLock(t xtime.UnixNano, attempt bool) {
	v, ok := s.states[t]
	if ok {
		if v.Success {
			return // Attempt is not relevant if success.
		}
		v.Attempt = attempt
		s.states[t] = v
		return
	}

	s.states[t] = doc.EntryIndexBlockState{
		Attempt: attempt,
	}
}


<|start_filename|>src/query/api/v1/handler/prometheus/native/list_tags_test.go<|end_filename|>
// Copyright (c) 2019 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package native

import (
	"context"
	"errors"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/m3db/m3/src/query/api/v1/handler/prometheus/handleroptions"
	"github.com/m3db/m3/src/query/api/v1/options"
	"github.com/m3db/m3/src/query/block"
	"github.com/m3db/m3/src/query/models"
	"github.com/m3db/m3/src/query/storage"
	"github.com/m3db/m3/src/query/storage/m3/consolidators"
	"github.com/m3db/m3/src/x/headers"
	xtest "github.com/m3db/m3/src/x/test"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type listTagsMatcher struct {
	start, end xtime.UnixNano
}

func (m *listTagsMatcher) String() string { return "list tags query" }
func (m *listTagsMatcher) Matches(x interface{}) bool {
	q, ok := x.(*storage.CompleteTagsQuery)
	if !ok {
		return false
	}

	if !q.Start.Equal(m.start) {
		return false
	}

	// NB: end time for the query should be roughly `Now`
	if !q.End.Equal(m.end) {
		return false
	}

	if !q.CompleteNameOnly {
		return false
	}

	if len(q.FilterNameTags) != 0 {
		return false
	}

	if len(q.TagMatchers) != 1 {
		return false
	}

	return models.MatchAll == q.TagMatchers[0].Type
}

var _ gomock.Matcher = &listTagsMatcher{}

func b(s string) []byte { return []byte(s) }

func TestListTags(t *testing.T) {
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			testListTags(t, tt.meta, tt.ex)
		})
	}
}

func testListTags(t *testing.T, meta block.ResultMetadata, header string) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	// setup storage and handler
	store := storage.NewMockStorage(ctrl)
	storeResult := &consolidators.CompleteTagsResult{
		CompleteNameOnly: true,
		CompletedTags: []consolidators.CompletedTag{
			{Name: b("bar")},
			{Name: b("baz")},
			{Name: b("foo")},
		},

		Metadata: meta,
	}

	now := xtime.Now()
	nowFn := func() time.Time {
		return now.ToTime()
	}

	fb, err := handleroptions.NewFetchOptionsBuilder(
		handleroptions.FetchOptionsBuilderOptions{Timeout: 15 * time.Second})
	require.NoError(t, err)
	opts := options.EmptyHandlerOptions().
		SetStorage(store).
		SetFetchOptionsBuilder(fb).
		SetTagOptions(models.NewTagOptions()).
		SetNowFn(nowFn)
	h := NewListTagsHandler(opts)
	for _, method := range []string{"GET", "POST"} {
		testListTagsWithMatch(t, now, store, storeResult, method, header, h, false)
		testListTagsWithMatch(t, now, store, storeResult, method, header, h, true)
	}
}

func testListTagsWithMatch(
	t *testing.T,
	now xtime.UnixNano,
	store *storage.MockStorage,
	storeResult *consolidators.CompleteTagsResult,
	method string,
	header string,
	h http.Handler,
	withMatchOverride bool,
) {
	tagMatcher := models.Matchers{{Type: models.MatchAll}}
	target := "/labels"
	if withMatchOverride {
		tagMatcher = models.Matchers{{
			Type:  models.MatchEqual,
			Name:  []byte("__name__"),
			Value: []byte("testing"),
		}}
		target = "/labels?match[]=testing"
	}

	matcher := &storage.CompleteTagsQuery{
		CompleteNameOnly: true,
		TagMatchers:      tagMatcher,
		Start:            0,
		End:              now,
	}
	store.EXPECT().CompleteTags(gomock.Any(), gomock.Eq(matcher), gomock.Any()).
		Return(storeResult, nil)

	req := httptest.NewRequest(method, target, nil)
	w := httptest.NewRecorder()

	h.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Result().StatusCode) // nolint:bodyclose

	body := w.Result().Body
	defer body.Close() // nolint:errcheck

	r, err := ioutil.ReadAll(body)
	require.NoError(t, err)

	ex := `{"status":"success","data":["bar","baz","foo"]}`
	require.Equal(t, ex, string(r))

	actual := w.Header().Get(headers.LimitHeader)
	assert.Equal(t, header, actual)
}

func TestListErrorTags(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	// setup storage and handler
	store := storage.NewMockStorage(ctrl)
	fb, err := handleroptions.NewFetchOptionsBuilder(
		handleroptions.FetchOptionsBuilderOptions{Timeout: 15 * time.Second})
	require.NoError(t, err)
	opts := options.EmptyHandlerOptions().
		SetStorage(store).
		SetFetchOptionsBuilder(fb)
	handler := NewListTagsHandler(opts)
	for _, method := range []string{"GET", "POST"} {
		matcher := &listTagsMatcher{
			start: xtime.FromSeconds(100),
			end:   xtime.FromSeconds(1000),
		}
		store.EXPECT().CompleteTags(gomock.Any(), matcher, gomock.Any()).
			Return(nil, errors.New("err"))

		req := httptest.NewRequest(method, "/labels?start=100&end=1000", nil)
		w := httptest.NewRecorder()

		handler.ServeHTTP(w, req)
		body := w.Result().Body
		defer body.Close()

		r, err := ioutil.ReadAll(body)
		require.NoError(t, err)

		require.JSONEq(t, `{"status":"error","error":"err"}`, string(r))
	}
}

//nolint:dupl
func TestListTagsTimeout(t *testing.T) {
	ctrl := xtest.NewController(t)
	defer ctrl.Finish()

	req := httptest.NewRequest("GET", "/labels", nil)
	w := httptest.NewRecorder()
	h := NewListTagsHandler(storageSetup(t, ctrl, 1*time.Millisecond, expectTimeout))
	h.ServeHTTP(w, req)

	assert.Equal(t, 504, w.Code, "Status code not 504")
	assert.Contains(t, w.Body.String(), "context deadline exceeded")
}

//nolint:dupl
func TestListTagsUseRequestContext(t *testing.T) {
	ctrl := xtest.NewController(t)
	defer ctrl.Finish()

	cancelledCtx, cancel := context.WithCancel(context.Background())
	cancel()
	req := httptest.NewRequest("GET", "/labels", nil).WithContext(cancelledCtx)
	w := httptest.NewRecorder()
	h := NewListTagsHandler(storageSetup(t, ctrl, 15*time.Second, expectCancellation))
	h.ServeHTTP(w, req)

	assert.Equal(t, 499, w.Code, "Status code not 499")
	assert.Contains(t, w.Body.String(), "context canceled")
}


<|start_filename|>src/aggregator/aggregator/gauge_elem_gen.go<|end_filename|>
// Copyright (c) 2022 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// This file was automatically generated by genny.
// Any changes will be lost if this file is regenerated.
// see https://github.com/mauricelam/genny

package aggregator

import (
	"fmt"
	"math"
	"sync"
	"time"

	raggregation "github.com/m3db/m3/src/aggregator/aggregation"
	"github.com/m3db/m3/src/metrics/metadata"
	"github.com/m3db/m3/src/metrics/metric/aggregated"
	"github.com/m3db/m3/src/metrics/metric/unaggregated"
	"github.com/m3db/m3/src/metrics/transformation"
	"github.com/m3db/m3/src/x/instrument"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/willf/bitset"
	"go.uber.org/zap"
)

type lockedGaugeAggregation struct {
	aggregation   gaugeAggregation
	sourcesSeen   map[uint32]*bitset.BitSet
	mtx           sync.Mutex
	lastUpdatedAt xtime.UnixNano
	dirty         bool
	// resendEnabled is allowed to change while an aggregation is open, so it must be behind the lock.
	resendEnabled bool
	closed        bool
}

type timedGauge struct {
	lockedAgg  *lockedGaugeAggregation
	startAt    xtime.UnixNano // start time of an aggregation window
	prevStart  xtime.UnixNano
	nextStart  xtime.UnixNano
	inDirtySet bool
}

// close is called when the aggregation has been expired or the element is being closed.
func (ta *timedGauge) close() {
	ta.lockedAgg.close()
	ta.lockedAgg = nil
}

// GaugeElem is an element storing time-bucketed aggregations.
type GaugeElem struct {
	gaugeElemBase
	elemBase
	// startTime -> agg (new one per every resolution)
	values map[xtime.UnixNano]timedGauge
	// startTime -> state. this is local state to the flusher and does not need to guarded with a lock.
	// values and flushState should always have the exact same key set.
	flushState map[xtime.UnixNano]flushState
	// sorted start aligned times that have been written to since the last flush
	dirty []xtime.UnixNano

	// internal/no need for synchronization: small buffers to avoid memory allocations during consumption
	toConsume            []consumeState
	flushStateToExpire   []xtime.UnixNano
	forwardTimesToExpire []xtime.UnixNano
	// end internal state

	// min time in the values map. allows for iterating through map.
	minStartTime xtime.UnixNano
	// max time in the values map. allows for iterating through map.
	maxStartTime xtime.UnixNano
}

// NewGaugeElem returns a new GaugeElem.
func NewGaugeElem(data ElemData, opts ElemOptions) (*GaugeElem, error) {
	e := &GaugeElem{
		elemBase:   newElemBase(opts),
		dirty:      make([]xtime.UnixNano, 0, defaultNumAggregations), // in most cases values will have two entries
		values:     make(map[xtime.UnixNano]timedGauge),
		flushState: make(map[xtime.UnixNano]flushState),
	}
	if err := e.ResetSetData(data); err != nil {
		return nil, err
	}
	return e, nil
}

// MustNewGaugeElem returns a new GaugeElem and panics if an error occurs.
func MustNewGaugeElem(data ElemData, opts ElemOptions) *GaugeElem {
	elem, err := NewGaugeElem(data, opts)
	if err != nil {
		panic(fmt.Errorf("unable to create element: %v", err))
	}
	return elem
}

// ResetSetData resets the element and sets data.
func (e *GaugeElem) ResetSetData(data ElemData) error {
	useDefaultAggregation := data.AggTypes.IsDefault()
	if useDefaultAggregation {
		data.AggTypes = e.DefaultAggregationTypes(e.aggTypesOpts)
	}
	if err := e.elemBase.resetSetData(data, useDefaultAggregation); err != nil {
		return err
	}
	return e.gaugeElemBase.ResetSetData(e.aggTypesOpts, data.AggTypes, useDefaultAggregation)
}

// AddUnion adds a metric value union at a given timestamp.
func (e *GaugeElem) AddUnion(timestamp time.Time, mu unaggregated.MetricUnion, resendEnabled bool) error {
	alignedStart := timestamp.Truncate(e.sp.Resolution().Window)
	lockedAgg, err := e.findOrCreate(alignedStart.UnixNano(), createAggregationOptions{})
	if err != nil {
		return err
	}
	lockedAgg.mtx.Lock()
	if lockedAgg.closed {
		// Note: this might have created an entry in the dirty set for lockedAgg when calling findOrCreate, even though
		// it's already closed. The Consume loop will detect this and clean it up.
		aggResendEnabled := lockedAgg.resendEnabled
		lockedAgg.mtx.Unlock()
		if !aggResendEnabled && resendEnabled {
			return errClosedBeforeResendEnabledMigration
		}
		return errAggregationClosed
	}
	lockedAgg.aggregation.AddUnion(timestamp, mu)
	lockedAgg.dirty = true
	lockedAgg.lastUpdatedAt = xtime.Now()
	lockedAgg.resendEnabled = resendEnabled
	lockedAgg.mtx.Unlock()
	return nil
}

// AddValue adds a metric value at a given timestamp.
func (e *GaugeElem) AddValue(timestamp time.Time, value float64, annotation []byte) error {
	alignedStart := timestamp.Truncate(e.sp.Resolution().Window).UnixNano()
	lockedAgg, err := e.findOrCreate(alignedStart, createAggregationOptions{})
	if err != nil {
		return err
	}
	lockedAgg.mtx.Lock()
	if lockedAgg.closed {
		lockedAgg.mtx.Unlock()
		return errAggregationClosed
	}
	lockedAgg.aggregation.Add(timestamp, value, annotation)
	lockedAgg.dirty = true
	lockedAgg.lastUpdatedAt = xtime.Now()
	lockedAgg.mtx.Unlock()
	return nil
}

// AddUnique adds a metric value from a given source at a given timestamp.
// If previous values from the same source have already been added to the
// same aggregation, the incoming value is discarded.
//nolint: dupl
func (e *GaugeElem) AddUnique(
	timestamp time.Time,
	metric aggregated.ForwardedMetric,
	metadata metadata.ForwardMetadata,
) error {
	alignedStart := timestamp.Truncate(e.sp.Resolution().Window).UnixNano()
	lockedAgg, err := e.findOrCreate(alignedStart, createAggregationOptions{
		initSourceSet: true,
	})
	if err != nil {
		return err
	}
	lockedAgg.mtx.Lock()
	if lockedAgg.closed {
		lockedAgg.mtx.Unlock()
		return errAggregationClosed
	}
	versionsSeen := lockedAgg.sourcesSeen[metadata.SourceID]
	if versionsSeen == nil {
		// N.B - these bitsets will be transitively cached through the cached sources seen.
		versionsSeen = bitset.New(defaultNumVersions)
		lockedAgg.sourcesSeen[metadata.SourceID] = versionsSeen
	}
	version := uint(metric.Version)
	if versionsSeen.Test(version) {
		lockedAgg.mtx.Unlock()
		return errDuplicateForwardingSource
	}
	versionsSeen.Set(version)

	if metric.Version > 0 {
		e.writeMetrics.updatedValues.Inc(1)
		for i := range metric.Values {
			if err := lockedAgg.aggregation.UpdateVal(timestamp, metric.Values[i], metric.PrevValues[i]); err != nil {
				return err
			}
		}
	} else {
		for _, v := range metric.Values {
			lockedAgg.aggregation.Add(timestamp, v, metric.Annotation)
		}
	}
	lockedAgg.dirty = true
	lockedAgg.lastUpdatedAt = xtime.Now()
	lockedAgg.resendEnabled = metadata.ResendEnabled
	lockedAgg.mtx.Unlock()
	return nil
}

// remove expired aggregations from the values map.
func (e *GaugeElem) expireValuesWithLock(
	targetNanos int64,
	isEarlierThanFn isEarlierThanFn,
	flushMetrics *flushMetrics,
) {
	var expiredCount int64
	e.flushStateToExpire = e.flushStateToExpire[:0]
	if len(e.values) == 0 {
		return
	}
	resolution := e.sp.Resolution().Window

	currAgg := e.values[e.minStartTime]
	resendExpire := targetNanos - int64(e.bufferForPastTimedMetricFn(resolution))
	for isEarlierThanFn(int64(currAgg.startAt), resolution, targetNanos) {
		if e.flushState[currAgg.startAt].latestResendEnabled {
			// if resend enabled we want to keep this value until it is outside the buffer past period.
			if !isEarlierThanFn(int64(currAgg.startAt), resolution, resendExpire) {
				break
			}
		}

		// close the agg to prevent any more writes.
		dirty := false
		currAgg.lockedAgg.mtx.Lock()
		if currAgg.lockedAgg.resendEnabled != e.flushState[currAgg.startAt].latestResendEnabled {
			// the aggregation migrated to resendEnabled after the flusher read the resendEnabled state.
			// keep the aggregation for now and try to expire on the next flush.
			currAgg.lockedAgg.mtx.Unlock()
			break
		}
		currAgg.lockedAgg.closed = true
		dirty = currAgg.lockedAgg.dirty
		currAgg.lockedAgg.mtx.Unlock()
		if dirty {
			// a race occurred and a write happened before we could close the aggregation. will expire next time.
			break
		}

		// if this current value is closed and clean it will no longer be flushed. this means it's safe
		// to remove the previous value since it will no longer be needed for binary transformations. when the
		// next value is eligible to be expired, this current value will actually be removed.
		// if we're currently pointing at the start skip this because there is no previous for the start. this
		// ensures we always keep at least one value in the map for binary transformations.
		if prevAgg, ok := e.prevAggWithLock(currAgg); ok && currAgg.startAt != e.minStartTime {
			// can't expire flush state until after the flushing, so we save the time to expire later.
			e.flushStateToExpire = append(e.flushStateToExpire, e.minStartTime)
			delete(e.values, e.minStartTime)
			e.minStartTime = currAgg.startAt
			expiredCount++

			// it's safe to access this outside the agg lock since it was closed in a previous iteration.
			// This is to make sure there aren't too many cached source sets taking up
			// too much space.
			if prevAgg.lockedAgg.sourcesSeen != nil && len(e.cachedSourceSets) < e.opts.MaxNumCachedSourceSets() {
				e.cachedSourceSets = append(e.cachedSourceSets, prevAgg.lockedAgg.sourcesSeen)
			}
			prevAgg.close()
		}
		var ok bool
		currAgg, ok = e.nextAggWithLock(currAgg)
		if !ok {
			break
		}
	}
	flushMetrics.valuesExpired.Inc(expiredCount)
}

func (e *GaugeElem) expireFlushState() {
	for _, t := range e.flushStateToExpire {
		fState, ok := e.flushState[t]
		if !ok {
			ts := t.ToTime()
			instrument.EmitAndLogInvariantViolation(e.opts.InstrumentOptions(), func(l *zap.Logger) {
				l.Error("expire time not in state map", zap.Time("ts", ts))
			})
			continue
		}
		fState.close()
		delete(e.flushState, t)
	}
}

// return the previous aggregation before the provided time. returns false if the provided time is the
// earliest time or the map is empty.
func (e *GaugeElem) prevAggWithLock(agg timedGauge) (timedGauge, bool) {
	if len(e.values) == 0 {
		return timedGauge{}, false
	}
	if agg.prevStart != 0 {
		prevAgg, ok := e.values[agg.prevStart]
		return prevAgg, ok
	}

	resolution := e.sp.Resolution().Window
	startTime := agg.startAt.Add(-resolution)
	for !startTime.Before(e.minStartTime) {
		agg, ok := e.values[startTime]
		if ok {
			return agg, true
		}
		startTime = startTime.Add(-resolution)
	}
	return timedGauge{}, false
}

// return the next aggregation after the provided time. returns false if the provided time is the
// largest time or the map is empty.
func (e *GaugeElem) nextAggWithLock(agg timedGauge) (timedGauge, bool) {
	if len(e.values) == 0 {
		return timedGauge{}, false
	}
	if agg.nextStart != 0 {
		nextAgg, ok := e.values[agg.nextStart]
		return nextAgg, ok
	}
	resolution := e.sp.Resolution().Window
	start := agg.startAt.Add(resolution)
	for !start.After(e.maxStartTime) {
		agg, ok := e.values[start]
		if ok {
			return agg, true
		}
		start = start.Add(resolution)
	}
	return timedGauge{}, false
}

// Consume consumes values before a given time and removes them from the element
// after they are consumed, returning whether the element can be collected after
// the consumption is completed.
// NB: Consume is not thread-safe and must be called within a single goroutine
// to avoid race conditions.
func (e *GaugeElem) Consume(
	targetNanos int64,
	isEarlierThanFn isEarlierThanFn,
	timestampNanosFn timestampNanosFn,
	targetNanosFn targetNanosFn,
	flushLocalFn flushLocalMetricFn,
	flushForwardedFn flushForwardedMetricFn,
	onForwardedFlushedFn onForwardingElemFlushedFn,
	jitter time.Duration,
	flushType flushType,
) bool {
	resolution := e.sp.Resolution().Window
	fMetrics := e.flushMetrics(resolution, flushType)
	fMetrics.elemsScanned.Inc(1)

	// reverse engineer the allowed lateness.
	latenessAllowed := time.Duration(targetNanos - targetNanosFn(targetNanos))
	e.Lock()
	if e.closed {
		e.Unlock()
		return false
	}

	// move currently dirty aggs to toConsume to process next.
	e.dirtyToConsumeWithLock(targetNanos, resolution, isEarlierThanFn)

	// expire the values and aggregations while we still hold the lock.
	e.expireValuesWithLock(targetNanos, isEarlierThanFn, fMetrics)
	canCollect := len(e.dirty) == 0 && e.tombstoned
	e.Unlock()

	// Process the aggregations that are ready for consumption.
	for _, cState := range e.toConsume {
		e.processValue(cState,
			timestampNanosFn,
			flushLocalFn,
			flushForwardedFn,
			resolution,
			latenessAllowed,
			jitter,
			fMetrics,
		)
	}
	fMetrics.valuesProcessed.Inc(int64(len(e.toConsume)))

	// expire the flush state after processing since it's needed in the processing.
	e.expireFlushState()

	if e.parsedPipeline.HasRollup {
		forwardedAggregationKey, _ := e.ForwardedAggregationKey()
		e.forwardTimesToExpire = e.forwardTimesToExpire[:0]
		for _, startTime := range e.flushStateToExpire {
			// the forward writer uses the timestamp of the aggregation, so need to convert the start aligned time
			// to a timestamp.
			e.forwardTimesToExpire = append(e.forwardTimesToExpire,
				xtime.UnixNano(timestampNanosFn(int64(startTime), resolution)))
		}
		onForwardedFlushedFn(e.onForwardedAggregationWrittenFn, forwardedAggregationKey, e.forwardTimesToExpire)
	}

	return canCollect
}

func (e *GaugeElem) dirtyToConsumeWithLock(targetNanos int64,
	resolution time.Duration,
	isEarlierThanFn isEarlierThanFn) {
	e.toConsume = e.toConsume[:0]
	// Evaluate and GC expired items.
	dirtyTimes := e.dirty
	e.dirty = e.dirty[:0]
	for i, dirtyTime := range dirtyTimes {
		if !isEarlierThanFn(int64(dirtyTime), resolution, targetNanos) {
			// not ready yet
			e.dirty = append(e.dirty, dirtyTime)
			continue
		}
		agg, ok := e.values[dirtyTime]
		if !ok {
			// there is a race where a writer adds a closed aggregation to the dirty set. eventually the closed
			// aggregation is expired and removed from the values map. ok to skip.
			continue
		}

		var dirty bool
		e.toConsume, dirty = e.appendConsumeStateWithLock(agg, e.toConsume, isDirty)
		if !dirty {
			// there is a race where the value was added to the dirty set, but the writer didn't actually update the
			// value yet (by marking dirty). add back to the dirty set so it can be processed in the next round once
			// the value has been updated.
			e.dirty = append(e.dirty, dirtyTime)
			continue
		}
		val := e.values[dirtyTime]
		val.inDirtySet = false
		e.values[dirtyTime] = val
		cState := e.toConsume[len(e.toConsume)-1]

		// potentially consume the nextAgg as well in case we need to cascade an update to the nextAgg.
		// this is necessary for binary transformations that rely on the previous aggregation value for calculating the
		// current aggregation value. if the nextAgg was already flushed, it used an outdated value for the previous
		// value (this agg). this can only happen when we allow updating previously flushed data (i.e resendEnabled).
		if cState.resendEnabled {
			nextAgg, ok := e.nextAggWithLock(agg)
			// only need to add if not already in the dirty set (since it will be added in a subsequent iteration).
			if ok &&
				// at the end of the dirty times OR the next dirty time does not match.
				(i == len(dirtyTimes)-1 || dirtyTimes[i+1] != nextAgg.startAt) {
				// only need to add if it was previously flushed.
				e.toConsume, _ = e.appendConsumeStateWithLock(nextAgg, e.toConsume, e.isFlushed)
			}
		}
	}
}

func (e *GaugeElem) isFlushed(c *consumeState) bool {
	return e.flushState[c.startAt].flushed
}

// append the consumeState for the timedGauge to the provided slice if it matches the provided filter.
// returns the updated slice and true if added.
func (e *GaugeElem) appendConsumeStateWithLock(
	agg timedGauge,
	toConsume []consumeState,
	includeFilter func(*consumeState) bool,
) ([]consumeState, bool) {
	// try reusing memory already allocated in the slice.
	if cap(toConsume) >= len(toConsume)+1 {
		toConsume = toConsume[:len(toConsume)+1]
	} else {
		toConsume = append(toConsume, consumeState{
			values: make([]float64, 0, len(e.aggTypes)),
		})
	}
	cState := &toConsume[len(toConsume)-1]
	cState.Reset()
	// copy the lockedAgg data while holding the lock.
	agg.lockedAgg.mtx.Lock()
	cState.dirty = agg.lockedAgg.dirty
	cState.lastUpdatedAt = agg.lockedAgg.lastUpdatedAt
	cState.resendEnabled = agg.lockedAgg.resendEnabled
	for _, aggType := range e.aggTypes {
		cState.values = append(cState.values, agg.lockedAgg.aggregation.ValueOf(aggType))
	}
	cState.annotation = raggregation.MaybeReplaceAnnotation(cState.annotation, agg.lockedAgg.aggregation.Annotation())
	agg.lockedAgg.dirty = false
	agg.lockedAgg.mtx.Unlock()

	// update with everything else.
	prevAgg, ok := e.prevAggWithLock(agg)
	if ok {
		cState.prevStartTime = prevAgg.startAt
	} else {
		cState.prevStartTime = 0
	}
	cState.startAt = agg.startAt
	// update the flush state with the latestResendEnabled since expireValuesWithLock needs it before actual processing.
	fState := e.flushState[cState.startAt]
	fState.latestResendEnabled = cState.resendEnabled
	e.flushState[cState.startAt] = fState

	if includeFilter != nil && !includeFilter(cState) {
		// since we eagerly appended, we need to remove if it should not be included.
		toConsume = toConsume[0 : len(toConsume)-1]
		return toConsume, false
	}
	return toConsume, true
}

// Close closes the element.
func (e *GaugeElem) Close() {
	e.Lock()
	if e.closed {
		e.Unlock()
		return
	}
	e.closed = true
	e.id = nil
	e.parsedPipeline = parsedPipeline{}
	e.writeForwardedMetricFn = nil
	e.onForwardedAggregationWrittenFn = nil
	for idx := range e.cachedSourceSets {
		e.cachedSourceSets[idx] = nil
	}
	e.cachedSourceSets = nil

	// note: this is not in the hot path so it's ok to iterate over the map.
	// this allows to catch any bugs with unexpected entries still in the map.
	minStartTime := e.minStartTime
	for k, v := range e.values {
		if k < minStartTime {
			k := k
			ts := e.minStartTime.ToTime()
			instrument.EmitAndLogInvariantViolation(e.opts.InstrumentOptions(), func(l *zap.Logger) {
				l.Error("value timestamp is less than min start time",
					zap.Time("ts", k.ToTime()),
					zap.Time("min", ts))
			})
		}
		v.close()
		delete(e.values, k)
		fState, ok := e.flushState[k]
		if ok {
			fState.close()
		}
		delete(e.flushState, k)
	}
	// clean up any dangling flush state that should never exist.
	for k, v := range e.flushState {
		ts := k.ToTime()
		instrument.EmitAndLogInvariantViolation(e.opts.InstrumentOptions(), func(l *zap.Logger) {
			l.Error("dangling state timestamp", zap.Time("ts", ts))
		})
		v.close()
		delete(e.flushState, k)
	}
	e.gaugeElemBase.Close()
	aggTypesPool := e.aggTypesOpts.TypesPool()
	pool := e.ElemPool(e.opts)
	e.dirty = e.dirty[:0]
	e.toConsume = e.toConsume[:0]
	e.flushStateToExpire = e.flushStateToExpire[:0]
	e.minStartTime = 0
	e.Unlock()

	if !e.useDefaultAggregation {
		aggTypesPool.Put(e.aggTypes)
	}
	pool.Put(e)
}

func (e *GaugeElem) insertDirty(alignedStart xtime.UnixNano) {
	numValues := len(e.dirty)

	// Optimize for the common case.
	if numValues > 0 && e.dirty[numValues-1] == alignedStart {
		return
	}
	// Binary search for the unusual case. We intentionally do not
	// use the sort.Search() function because it requires passing
	// in a closure.
	left, right := 0, numValues
	for left < right {
		mid := left + (right-left)/2 // avoid overflow
		if e.dirty[mid] < alignedStart {
			left = mid + 1
		} else {
			right = mid
		}
	}
	// If the current timestamp is equal to or larger than the target time,
	// return the index as is.
	if left < numValues && e.dirty[left] == alignedStart {
		return
	}

	e.dirty = append(e.dirty, 0)
	copy(e.dirty[left+1:numValues+1], e.dirty[left:numValues])
	e.dirty[left] = alignedStart
}

// find finds the aggregation for a given time, or returns nil.
//nolint: dupl
func (e *GaugeElem) find(alignedStartNanos xtime.UnixNano) (timedGauge, error) {
	e.RLock()
	if e.closed {
		e.RUnlock()
		return timedGauge{}, errElemClosed
	}
	timedAgg, ok := e.values[alignedStartNanos]
	if ok {
		e.RUnlock()
		return timedAgg, nil
	}
	e.RUnlock()
	return timedGauge{}, nil
}

// findOrCreate finds the aggregation for a given time, or creates one
// if it doesn't exist.
//nolint: dupl
func (e *GaugeElem) findOrCreate(
	alignedStartNanos int64,
	createOpts createAggregationOptions,
) (*lockedGaugeAggregation, error) {
	e.writeMetrics.writes.Inc(1)
	alignedStart := xtime.UnixNano(alignedStartNanos)
	found, err := e.find(alignedStart)
	if err != nil {
		return nil, err
	}
	// if the aggregation is found and does not need to be updated, return as is.
	if found.lockedAgg != nil && found.inDirtySet {
		return found.lockedAgg, err
	}

	e.Lock()
	if e.closed {
		e.Unlock()
		return nil, errElemClosed
	}

	timedAgg, ok := e.values[alignedStart]
	if ok {
		// add to dirty set so it will be flushed.
		if !timedAgg.inDirtySet {
			timedAgg.inDirtySet = true
			e.insertDirty(alignedStart)
			e.values[alignedStart] = timedAgg
		}
		e.Unlock()
		return timedAgg.lockedAgg, nil
	}

	var sourcesSeen map[uint32]*bitset.BitSet
	if createOpts.initSourceSet {
		if numCachedSourceSets := len(e.cachedSourceSets); numCachedSourceSets > 0 {
			sourcesSeen = e.cachedSourceSets[numCachedSourceSets-1]
			e.cachedSourceSets[numCachedSourceSets-1] = nil
			e.cachedSourceSets = e.cachedSourceSets[:numCachedSourceSets-1]
			for _, bs := range sourcesSeen {
				bs.ClearAll()
			}
		} else {
			sourcesSeen = make(map[uint32]*bitset.BitSet)
		}
	}
	// NB(vytenis): lockedGaugeAggregation will be returned to pool on timedGauge close.
	// this is a bit different from regular pattern of using a pool object due to codegen with Genny limitations,
	// so we can avoid writing more boilerplate.
	// timedGauge itself is always pass-by-value, but lockedGaugeAggregation incurs an expensive allocation on heap
	// in the critical path (30%+, depending on workload as of 2020-05-01): see https://github.com/m3db/m3/pull/4109
	timedAgg = timedGauge{
		startAt: alignedStart,
		lockedAgg: lockedGaugeAggregationFromPool(
			e.NewAggregation(e.opts, e.aggOpts),
			sourcesSeen,
		),
		inDirtySet: true,
	}

	if len(e.values) == 0 || e.minStartTime > alignedStart {
		e.minStartTime = alignedStart
	}
	prevMaxStart := e.maxStartTime
	if len(e.values) == 0 || alignedStart > e.maxStartTime {
		e.maxStartTime = alignedStart
	}

	if len(e.values) > 0 {
		if e.maxStartTime == alignedStart {
			// common case we are adding the latest start time.
			timedAgg.prevStart = prevMaxStart
			prevAgg := e.values[prevMaxStart]
			prevAgg.nextStart = alignedStart
			e.values[prevMaxStart] = prevAgg
		} else {
			// look up
			prevAgg, ok := e.prevAggWithLock(timedAgg)
			if ok {
				timedAgg.prevStart = prevAgg.startAt
				prevAgg.nextStart = alignedStart
				e.values[prevAgg.startAt] = prevAgg
			}
			nextAgg, ok := e.nextAggWithLock(timedAgg)
			if ok {
				timedAgg.nextStart = nextAgg.startAt
				nextAgg.prevStart = alignedStart
				e.values[nextAgg.startAt] = nextAgg
			}
		}
	}

	e.values[alignedStart] = timedAgg
	e.insertDirty(alignedStart)
	e.Unlock()
	return timedAgg.lockedAgg, nil
}

// returns true if a datapoint is emitted.
func (e *GaugeElem) processValue(
	cState consumeState,
	timestampNanosFn timestampNanosFn,
	flushLocalFn flushLocalMetricFn,
	flushForwardedFn flushForwardedMetricFn,
	resolution time.Duration,
	latenessAllowed time.Duration,
	jitter time.Duration,
	flushMetrics *flushMetrics,
) {
	var (
		transformations  = e.parsedPipeline.Transformations
		discardNaNValues = e.opts.DiscardNaNAggregatedValues()
		timestamp        = xtime.UnixNano(timestampNanosFn(int64(cState.startAt), resolution))
		prevTimestamp    = xtime.UnixNano(timestampNanosFn(int64(cState.prevStartTime), resolution))
		// expectedProcessingTime should be the next resolution window after the aggregation was updated.
		expectedProcessingTime = cState.lastUpdatedAt.Truncate(resolution).Add(resolution)
	)
	fState := e.flushState[cState.startAt]
	if cState.dirty && fState.flushed && !cState.resendEnabled {
		cState := cState
		instrument.EmitAndLogInvariantViolation(e.opts.InstrumentOptions(), func(l *zap.Logger) {
			l.Error("reflushing aggregation without resendEnabled", zap.Any("consumeState", cState))
		})
	}

	for aggTypeIdx, aggType := range e.aggTypes {
		var extraDp transformation.Datapoint
		value := cState.values[aggTypeIdx]
		for _, transformOp := range transformations {
			unaryOp, isUnaryOp := transformOp.UnaryTransform()
			binaryOp, isBinaryOp := transformOp.BinaryTransform()
			unaryMultiOp, isUnaryMultiOp := transformOp.UnaryMultiOutputTransform()
			switch {
			case isUnaryOp:
				curr := transformation.Datapoint{
					TimeNanos: int64(timestamp),
					Value:     value,
				}

				res := unaryOp.Evaluate(curr)

				value = res.Value

			case isBinaryOp:
				prev := transformation.Datapoint{
					Value: nan,
				}
				if cState.prevStartTime > 0 {
					prevFlushState, ok := e.flushState[cState.prevStartTime]
					if !ok {
						ts := cState.prevStartTime.ToTime()
						instrument.EmitAndLogInvariantViolation(e.opts.InstrumentOptions(), func(l *zap.Logger) {
							l.Error("previous start time not in state map",
								zap.Time("ts", ts))
						})
					} else {
						prev.Value = prevFlushState.consumedValues[aggTypeIdx]
						prev.TimeNanos = int64(prevTimestamp)
					}
				}
				curr := transformation.Datapoint{
					TimeNanos: int64(timestamp),
					Value:     value,
				}
				res := binaryOp.Evaluate(prev, curr, transformation.FeatureFlags{})

				// NB: we only need to record the value needed for derivative transformations.
				// We currently only support first-order derivative transformations so we only
				// need to keep one value. In the future if we need to support higher-order
				// derivative transformations, we need to store an array of values here.
				if fState.consumedValues == nil {
					fState.consumedValues = make([]float64, len(e.aggTypes))
				}
				fState.consumedValues[aggTypeIdx] = curr.Value
				value = res.Value
			case isUnaryMultiOp:
				curr := transformation.Datapoint{
					TimeNanos: int64(timestamp),
					Value:     value,
				}

				var res transformation.Datapoint
				res, extraDp = unaryMultiOp.Evaluate(curr, resolution)
				value = res.Value
			}
		}

		if discardNaNValues && math.IsNaN(value) {
			continue
		}

		// It's ok to send a 0 prevValue on the first forward because it's not used in AddUnique unless it's a
		// resend (version > 0)
		var prevValue float64
		if fState.emittedValues == nil {
			fState.emittedValues = make([]float64, len(e.aggTypes))
		} else {
			prevValue = fState.emittedValues[aggTypeIdx]
		}
		fState.emittedValues[aggTypeIdx] = value
		if fState.flushed {
			// no need to resend a value that hasn't changed.
			if (math.IsNaN(prevValue) && math.IsNaN(value)) || (prevValue == value) {
				continue
			}
		}

		fwdType := forwardTypeRemote
		if !e.parsedPipeline.HasRollup {
			fwdType = forwardTypeLocal
			toFlush := make([]transformation.Datapoint, 0, 2)
			toFlush = append(toFlush, transformation.Datapoint{
				TimeNanos: int64(timestamp),
				Value:     value,
			})
			if extraDp.TimeNanos != 0 {
				toFlush = append(toFlush, extraDp)
			}
			for _, point := range toFlush {
				switch e.idPrefixSuffixType {
				case NoPrefixNoSuffix:
					flushLocalFn(nil, e.id, nil, point.TimeNanos, point.Value, cState.annotation,
						e.sp)
				case WithPrefixWithSuffix:
					flushLocalFn(e.FullPrefix(e.opts), e.id, e.TypeStringFor(e.aggTypesOpts, aggType),
						point.TimeNanos, point.Value, cState.annotation, e.sp)
				}
			}
		} else {
			forwardedAggregationKey, _ := e.ForwardedAggregationKey()
			flushForwardedFn(e.writeForwardedMetricFn, forwardedAggregationKey,
				int64(timestamp), value, prevValue, cState.annotation, cState.resendEnabled)
		}
		// add latenessAllowed and jitter to the timestamp of the aggregation, since those should not be
		// counted towards the processing lag.
		// forward lag = current time - (agg timestamp + lateness allowed + jitter)
		// use expectedProcessingTime instead of the aggregation timestamp since the aggregation timestamp could be
		// in the past for updated aggregations (resendEnabled).
		lag := xtime.Since(expectedProcessingTime.Add(latenessAllowed))
		flushMetrics.forwardLag(forwardKey{fwdType: fwdType, jitter: false}).
			RecordDuration(lag)
		flushMetrics.forwardLag(forwardKey{fwdType: fwdType, jitter: true}).
			RecordDuration(lag + jitter)
	}
	fState.flushed = true
	e.flushState[cState.startAt] = fState
}


<|start_filename|>src/query/storage/m3/cluster_resolver_test.go<|end_filename|>
// Copyright (c) 2019 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package m3

import (
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/client"
	"github.com/m3db/m3/src/metrics/policy"
	"github.com/m3db/m3/src/query/storage"
	"github.com/m3db/m3/src/query/storage/m3/consolidators"
	"github.com/m3db/m3/src/query/storage/m3/storagemetadata"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/ident"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestFanoutAggregatedOptimizationDisabledGivesAllClustersAsPartial(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	now := xtime.Now()
	s, _ := setup(t, ctrl)
	store, ok := s.(*m3storage)
	assert.True(t, ok)
	var r reusedAggregatedNamespaceSlices
	opts := &storage.FanoutOptions{
		FanoutAggregatedOptimized: storage.FanoutForceDisable,
	}

	clusters := store.clusters.ClusterNamespaces()
	r = aggregatedNamespaces(clusters, r, nil, now, now, opts)
	assert.Equal(t, 0, len(r.completeAggregated))
	assert.Equal(t, 4, len(r.partialAggregated))
}

func TestFanoutUnaggregatedDisableReturnsAggregatedNamespaces(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	s, _ := setup(t, ctrl)
	store, ok := s.(*m3storage)
	assert.True(t, ok)
	opts := &storage.FanoutOptions{
		FanoutUnaggregated: storage.FanoutForceDisable,
	}

	start := xtime.Now()
	end := start.Add(time.Hour * 24 * -90)
	_, clusters, err := resolveClusterNamespacesForQuery(start, start, end, store.clusters, opts,
		nil, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(clusters))
	assert.Equal(t, "metrics_aggregated_1m:30d", clusters[0].NamespaceID().String())
}

func TestFanoutUnaggregatedEnabledReturnsUnaggregatedNamespaces(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	s, _ := setup(t, ctrl)
	store, ok := s.(*m3storage)
	assert.True(t, ok)
	opts := &storage.FanoutOptions{
		FanoutUnaggregated: storage.FanoutForceEnable,
	}

	start := xtime.Now()
	end := start.Add(time.Hour * 24 * -90)
	_, clusters, err := resolveClusterNamespacesForQuery(start,
		start, end, store.clusters, opts, nil, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(clusters))
	assert.Equal(t, "metrics_unaggregated", clusters[0].NamespaceID().String())
}

func TestGraphitePath(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	s, _ := setup(t, ctrl)
	store, ok := s.(*m3storage)
	assert.True(t, ok)
	opts := &storage.FanoutOptions{
		FanoutUnaggregated:        storage.FanoutForceDisable,
		FanoutAggregated:          storage.FanoutForceEnable,
		FanoutAggregatedOptimized: storage.FanoutForceDisable,
	}

	start := xtime.Now()
	end := start.Add(time.Second * -30)
	_, clusters, err := resolveClusterNamespacesForQuery(start, start, end, store.clusters, opts,
		nil, nil)
	require.NoError(t, err)
	require.Equal(t, 4, len(clusters))
	expected := []string{
		"metrics_aggregated_1m:30d", "metrics_aggregated_5m:90d",
		"metrics_aggregated_partial_1m:180d", "metrics_aggregated_10m:365d",
	}

	for i, cluster := range clusters {
		assert.Equal(t, expected[i], cluster.NamespaceID().String())
	}
}

var (
	genRetentionFiltered, genRetentionUnfiltered = time.Hour, time.Hour * 10
	genResolution                                = time.Minute
)

func generateClusters(t *testing.T, ctrl *gomock.Controller) Clusters {
	session := client.NewMockSession(ctrl)

	clusters, err := NewClusters(UnaggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("UNAGG"),
		Retention:   genRetentionFiltered + time.Minute,
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_FILTERED"),
		Retention:   genRetentionFiltered,
		Resolution:  genResolution,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_NO_FILTER"),
		Retention:   genRetentionUnfiltered,
		Resolution:  genResolution,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_FILTERED_COMPLETE"),
		Retention:   genRetentionFiltered,
		Resolution:  genResolution + time.Second,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_NO_FILTER_COMPLETE"),
		Retention:   genRetentionUnfiltered,
		Resolution:  genResolution + time.Second,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
		Session:     session,
	})

	require.NoError(t, err)
	return clusters
}

// used to generate storage.RelatedQueries during a test
type testTimeOffsets struct {
	startOffset time.Duration
	endOffset   time.Duration
}

var testCases = []struct {
	name                     string
	queryLength              time.Duration
	opts                     *storage.FanoutOptions
	restrict                 *storage.RestrictQueryOptions
	expectedType             consolidators.QueryFanoutType
	relatedQueryOffsets      []testTimeOffsets
	expectedClusterNames     []string
	expectedErr              error
	expectedErrContains      string
	expectedErrInvalidParams bool
}{
	{
		name: "all disabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceDisable,
			FanoutAggregated:          storage.FanoutForceDisable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType: consolidators.NamespaceInvalid,
		expectedErr:  errUnaggregatedAndAggregatedDisabled,
	},
	{
		name: "optimize enabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceDisable,
			FanoutAggregated:          storage.FanoutForceDisable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		expectedType: consolidators.NamespaceInvalid,
		expectedErr:  errUnaggregatedAndAggregatedDisabled,
	},
	{
		name: "unagg enabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceDisable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name:        "unagg enabled short range",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceDisable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType:         consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name: "unagg optimized enabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceDisable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name:        "unagg optimized enabled short range",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceDisable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		expectedType:         consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name: "agg enabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceDisable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType: consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{
			"AGG_FILTERED", "AGG_NO_FILTER",
			"AGG_FILTERED_COMPLETE", "AGG_NO_FILTER_COMPLETE",
		},
	},
	{
		name:        "agg enabled short range",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceDisable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType: consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{
			"AGG_FILTERED", "AGG_NO_FILTER",
			"AGG_FILTERED_COMPLETE", "AGG_NO_FILTER_COMPLETE",
		},
	},
	{
		name: "unagg and agg enabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType: consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{
			"UNAGG", "AGG_FILTERED", "AGG_NO_FILTER",
			"AGG_FILTERED_COMPLETE", "AGG_NO_FILTER_COMPLETE",
		},
	},
	{
		name:        "unagg and agg enabled short range",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		expectedType:         consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name: "agg and optimization enabled",
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceDisable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		expectedType:         consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{"AGG_NO_FILTER", "AGG_NO_FILTER_COMPLETE"},
	},
	{
		name:        "all enabled short range",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		expectedType:         consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name:        "all enabled long range",
		queryLength: time.Hour * 1000,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"AGG_NO_FILTER", "AGG_NO_FILTER_COMPLETE"},
	},
	{
		name:        "restrict to unaggregated",
		queryLength: time.Hour * 1000,
		restrict: &storage.RestrictQueryOptions{
			RestrictByType: &storage.RestrictByType{
				MetricsType: storagemetadata.UnaggregatedMetricsType,
			},
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"UNAGG"},
	},
	{
		name:        "restrict to aggregate filtered",
		queryLength: time.Hour * 1000,
		restrict: &storage.RestrictQueryOptions{
			RestrictByType: &storage.RestrictByType{
				MetricsType: storagemetadata.AggregatedMetricsType,
				StoragePolicy: policy.MustParseStoragePolicy(
					genResolution.String() + ":" + genRetentionFiltered.String()),
			},
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"AGG_FILTERED"},
	},
	{
		name:        "restrict to aggregate unfiltered",
		queryLength: time.Hour * 1000,
		restrict: &storage.RestrictQueryOptions{
			RestrictByType: &storage.RestrictByType{
				MetricsType: storagemetadata.AggregatedMetricsType,
				StoragePolicy: policy.MustParseStoragePolicy(
					genResolution.String() + ":" + genRetentionUnfiltered.String()),
			},
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"AGG_NO_FILTER"},
	},
	{
		name:        "restrict with unknown metrics type",
		queryLength: time.Hour * 1000,
		restrict: &storage.RestrictQueryOptions{
			RestrictByType: &storage.RestrictByType{
				MetricsType: storagemetadata.UnknownMetricsType,
			},
		},
		expectedErrContains:      "unrecognized metrics type:",
		expectedErrInvalidParams: true,
	},
	{
		name:        "restrict with unknown storage policy",
		queryLength: time.Hour * 1000,
		restrict: &storage.RestrictQueryOptions{
			RestrictByType: &storage.RestrictByType{
				MetricsType:   storagemetadata.AggregatedMetricsType,
				StoragePolicy: policy.MustParseStoragePolicy("1s:100d"),
			},
		},
		expectedErrContains:      "could not find namespace for storage policy:",
		expectedErrInvalidParams: true,
	},
	{
		name:        "restrict to multiple aggregate",
		queryLength: time.Hour * 1000,
		restrict: &storage.RestrictQueryOptions{
			RestrictByTypes: []*storage.RestrictByType{
				{
					MetricsType: storagemetadata.AggregatedMetricsType,
					StoragePolicy: policy.MustParseStoragePolicy(
						genResolution.String() + ":" + genRetentionFiltered.String()),
				},
				{
					MetricsType: storagemetadata.AggregatedMetricsType,
					StoragePolicy: policy.MustParseStoragePolicy(
						(genResolution + time.Second).String() + ":" + genRetentionUnfiltered.String()),
				},
			},
		},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"AGG_FILTERED", "AGG_NO_FILTER_COMPLETE"},
	},
	{
		name:        "restrict to multiple aggregate all range",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceDisable,
			FanoutAggregated:          storage.FanoutDefault,
			FanoutAggregatedOptimized: storage.FanoutForceDisable,
		},
		restrict: &storage.RestrictQueryOptions{
			RestrictByTypes: []*storage.RestrictByType{
				{
					MetricsType: storagemetadata.AggregatedMetricsType,
					StoragePolicy: policy.MustParseStoragePolicy(
						(genResolution + time.Second).String() + ":" + genRetentionFiltered.String()),
				},
				{
					MetricsType: storagemetadata.AggregatedMetricsType,
					StoragePolicy: policy.MustParseStoragePolicy(
						(genResolution + time.Second).String() + ":" + genRetentionUnfiltered.String()),
				},
			},
		},
		expectedType:         consolidators.NamespaceCoversAllQueryRange,
		expectedClusterNames: []string{"AGG_FILTERED_COMPLETE", "AGG_NO_FILTER_COMPLETE"},
	},
	{
		name:        "all enabled short range w/ related queries",
		queryLength: time.Minute,
		opts: &storage.FanoutOptions{
			FanoutUnaggregated:        storage.FanoutForceEnable,
			FanoutAggregated:          storage.FanoutForceEnable,
			FanoutAggregatedOptimized: storage.FanoutForceEnable,
		},
		relatedQueryOffsets:  []testTimeOffsets{{startOffset: time.Hour * 1000, endOffset: 0}},
		expectedType:         consolidators.NamespaceCoversPartialQueryRange,
		expectedClusterNames: []string{"AGG_NO_FILTER", "AGG_NO_FILTER_COMPLETE"},
	},
}

func TestResolveClusterNamespacesForQueryWithOptions(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	clusters := generateClusters(t, ctrl)

	now := xtime.Now()
	end := now
	for _, tt := range testCases {
		t.Run(tt.name, func(t *testing.T) {
			start := now.Add(tt.queryLength * -1)
			if tt.queryLength == 0 {
				// default case
				start = start.Add(time.Hour * -2)
			}

			relatedQueries := make([]storage.QueryTimespan, 0, len(tt.relatedQueryOffsets))
			for _, offset := range tt.relatedQueryOffsets {
				timespan := storage.QueryTimespan{
					Start: now.Add(-offset.startOffset),
					End:   now.Add(-offset.endOffset),
				}
				relatedQueries = append(relatedQueries, timespan)
			}

			fanoutType, clusters, err := resolveClusterNamespacesForQuery(now,
				start, end, clusters, tt.opts, tt.restrict,
				&storage.RelatedQueryOptions{Timespans: relatedQueries})
			if tt.expectedErr != nil {
				assert.Error(t, err)
				assert.Equal(t, tt.expectedErr, err)
				assert.Nil(t, clusters)
				return
			}

			if substr := tt.expectedErrContains; substr != "" {
				assert.Error(t, err)
				assert.True(t, strings.Contains(err.Error(), substr))
				assert.Nil(t, clusters)
				invalidParams := xerrors.IsInvalidParams(err)
				assert.Equal(t, tt.expectedErrInvalidParams, invalidParams)
				return
			}

			require.NoError(t, err)
			actualNames := make([]string, len(clusters))
			for i, c := range clusters {
				actualNames[i] = c.NamespaceID().String()
			}

			// NB: order does not matter.
			sort.Strings(actualNames)
			sort.Strings(tt.expectedClusterNames)
			assert.Equal(t, tt.expectedClusterNames, actualNames)
			assert.Equal(t, tt.expectedType, fanoutType)
		})
	}
}

func TestLongUnaggregatedRetention(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	session := client.NewMockSession(ctrl)
	retentionFiltered, retentionUnfiltered := time.Hour, time.Hour*10
	resolution := time.Minute

	clusters, err := NewClusters(UnaggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("UNAGG"),
		Retention:   retentionUnfiltered,
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_FILTERED"),
		Retention:   retentionFiltered,
		Resolution:  resolution,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_NO_FILTER"),
		Retention:   retentionUnfiltered + time.Minute,
		Resolution:  resolution,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_FILTERED_COMPLETE"),
		Retention:   retentionFiltered,
		Resolution:  resolution + time.Second,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("AGG_NO_FILTER_COMPLETE"),
		Retention:   retentionUnfiltered,
		Resolution:  resolution + time.Second,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
		Session:     session,
	})

	require.NoError(t, err)
	now := xtime.Now()
	end := now
	start := now.Add(time.Hour * -100)
	opts := &storage.FanoutOptions{
		FanoutUnaggregated:        storage.FanoutForceEnable,
		FanoutAggregated:          storage.FanoutForceEnable,
		FanoutAggregatedOptimized: storage.FanoutForceEnable,
	}

	fanoutType, ns, err := resolveClusterNamespacesForQuery(now, start, end, clusters,
		opts, nil, nil)

	require.NoError(t, err)
	actualNames := make([]string, len(ns))
	for i, c := range ns {
		actualNames[i] = c.NamespaceID().String()
	}

	expected := []string{"UNAGG", "AGG_NO_FILTER"}
	// NB: order does not matter.
	sort.Strings(actualNames)
	sort.Strings(expected)
	assert.Equal(t, expected, actualNames)
	assert.Equal(t, consolidators.NamespaceCoversPartialQueryRange, fanoutType)
}

func TestExampleCase(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	session := client.NewMockSession(ctrl)
	ns, err := NewClusters(
		UnaggregatedClusterNamespaceDefinition{
			NamespaceID: ident.StringID("metrics_10s_24h"),
			Retention:   24 * time.Hour,
			Session:     session,
		}, AggregatedClusterNamespaceDefinition{
			NamespaceID: ident.StringID("metrics_180s_360h"),
			Retention:   360 * time.Hour,
			Resolution:  120 * time.Second,
			Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
			Session:     session,
		}, AggregatedClusterNamespaceDefinition{
			NamespaceID: ident.StringID("metrics_600s_17520h"),
			Retention:   17520 * time.Hour,
			Resolution:  600 * time.Second,
			Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
			Session:     session,
		},
	)
	require.NoError(t, err)

	now := xtime.Now()
	end := now

	for i := 27; i < 17520; i++ {
		start := now.Add(time.Hour * -1 * time.Duration(i))
		fanoutType, clusters, err := resolveClusterNamespacesForQuery(now, start, end, ns,
			&storage.FanoutOptions{}, nil, nil)

		require.NoError(t, err)
		actualNames := make([]string, len(clusters))
		for i, c := range clusters {
			actualNames[i] = c.NamespaceID().String()
		}

		// NB: order does not matter.
		sort.Strings(actualNames)
		assert.Equal(t, []string{
			"metrics_10s_24h",
			"metrics_180s_360h", "metrics_600s_17520h",
		}, actualNames)
		assert.Equal(t, consolidators.NamespaceCoversPartialQueryRange, fanoutType)
	}
}

func TestDeduplicatePartialAggregateNamespaces(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	session := client.NewMockSession(ctrl)
	ns, err := NewClusters(UnaggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("default"),
		Retention:   24 * time.Hour,
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("aggregated_block_6h"),
		Retention:   360 * time.Hour,
		Resolution:  10 * time.Minute,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: false},
		Session:     session,
	}, AggregatedClusterNamespaceDefinition{
		NamespaceID: ident.StringID("aggregated_block_6h"),
		Retention:   360 * time.Hour,
		Resolution:  1 * time.Hour,
		Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
		Session:     session,
	})
	require.NoError(t, err)

	now := xtime.Now()
	end := now

	start := now.Add(-48 * time.Hour)
	fanoutType, clusters, err := resolveClusterNamespacesForQuery(now, start, end, ns,
		&storage.FanoutOptions{}, nil, nil)
	require.NoError(t, err)

	actualNames := make([]string, len(clusters))
	for i, c := range clusters {
		actualNames[i] = c.NamespaceID().String()
	}

	// NB: order does not matter.
	sort.Strings(actualNames)
	assert.Equal(t, []string{"aggregated_block_6h"}, actualNames)
	assert.Equal(t, consolidators.NamespaceCoversAllQueryRange, fanoutType)
}

func TestResolveNamespaceWithDataLatency(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	dataLatency := 10 * time.Hour

	session := client.NewMockSession(ctrl)
	ns, err := NewClusters(
		UnaggregatedClusterNamespaceDefinition{
			NamespaceID: ident.StringID("default"),
			Retention:   24 * time.Hour,
			Session:     session,
		},
		AggregatedClusterNamespaceDefinition{
			NamespaceID: ident.StringID("aggregated_30d"),
			Retention:   30 * 24 * time.Hour,
			Resolution:  5 * time.Minute,
			Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
			Session:     session,
		},
		AggregatedClusterNamespaceDefinition{
			NamespaceID: ident.StringID("aggregated_60d"),
			Retention:   60 * 24 * time.Hour,
			Resolution:  10 * time.Minute,
			Downsample:  &ClusterNamespaceDownsampleOptions{All: true},
			DataLatency: dataLatency,
			Session:     session,
		},
	)
	require.NoError(t, err)

	var (
		now   = xtime.Now()
		start = now.Add(-40 * 24 * time.Hour)
		end   = now.Add(-3 * time.Hour)
	)

	fanoutType, clusters, err := resolveClusterNamespacesForQuery(now, start, end, ns,
		&storage.FanoutOptions{}, nil, nil)
	require.NoError(t, err)

	actualNamespaces := make(map[string]narrowing)
	for _, c := range clusters {
		actualNamespaces[c.NamespaceID().String()] = c.narrowing
	}

	stitchAt := now.Add(-dataLatency).Truncate(10 * time.Minute)
	expectedNamespaces := map[string]narrowing{
		"default":        {start: stitchAt},
		"aggregated_60d": {end: stitchAt},
	}

	assert.Equal(t, expectedNamespaces, actualNamespaces)
	assert.Equal(t, consolidators.NamespaceCoversAllQueryRange, fanoutType)
}


<|start_filename|>src/query/api/v1/handler/prom/prom.go<|end_filename|>
// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package prom

import (
	"errors"
	"net/http"

	"github.com/m3db/m3/src/query/api/v1/options"
	"github.com/m3db/m3/src/query/block"
	"github.com/m3db/m3/src/query/storage/prometheus"

	"github.com/prometheus/prometheus/promql/parser"
	promstorage "github.com/prometheus/prometheus/storage"
)

// opts defines options for PromQL handler.
type opts struct {
	instant    bool
	queryable  promstorage.Queryable
	newQueryFn NewQueryFn
}

// Option is a Prometheus handler option.
type Option func(*opts) error

// WithEngine sets the PromQL engine.
func WithEngine(promQLEngineFn options.PromQLEngineFn) Option {
	return withEngine(promQLEngineFn, false)
}

// WithInstantEngine sets the PromQL instant engine.
func WithInstantEngine(promQLEngineFn options.PromQLEngineFn) Option {
	return withEngine(promQLEngineFn, true)
}

func withEngine(promQLEngineFn options.PromQLEngineFn, instant bool) Option {
	return func(o *opts) error {
		if promQLEngineFn == nil {
			return errors.New("invalid engine fn")
		}
		o.instant = instant
		o.newQueryFn = newRangeQueryFn(promQLEngineFn, o.queryable)
		if instant {
			o.newQueryFn = newInstantQueryFn(promQLEngineFn, o.queryable)
		}
		return nil
	}
}

func newDefaultOptions(hOpts options.HandlerOptions) opts {
	queryable := prometheus.NewPrometheusQueryable(
		prometheus.PrometheusOptions{
			Storage:           hOpts.Storage(),
			InstrumentOptions: hOpts.InstrumentOpts(),
		})
	return opts{
		queryable:  queryable,
		instant:    false,
		newQueryFn: newRangeQueryFn(hOpts.PrometheusEngineFn(), queryable),
	}
}

// NewReadHandler creates a handler to handle PromQL requests.
func NewReadHandler(hOpts options.HandlerOptions, options ...Option) (http.Handler, error) {
	opts := newDefaultOptions(hOpts)
	for _, optionFn := range options {
		if err := optionFn(&opts); err != nil {
			return nil, err
		}
	}
	return newReadHandler(hOpts, opts)
}

// ApplyRangeWarnings applies warnings encountered during execution.
func ApplyRangeWarnings(
	query string, meta *block.ResultMetadata,
) error {
	expr, err := parser.ParseExpr(query)
	if err != nil {
		return err
	}

	parser.Inspect(expr, func(node parser.Node, path []parser.Node) error {
		if n, ok := node.(*parser.MatrixSelector); ok {
			meta.VerifyTemporalRange(n.Range)
		}

		return nil
	})

	return nil
}


<|start_filename|>src/dbnode/integration/index_active_block_rotate_test.go<|end_filename|>
//go:build integration
// +build integration

//
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/integration/generate"
	"github.com/m3db/m3/src/dbnode/namespace"
	"github.com/m3db/m3/src/dbnode/retention"
	"github.com/m3db/m3/src/dbnode/storage/index/compaction"
	xclock "github.com/m3db/m3/src/x/clock"
	"github.com/m3db/m3/src/x/ident"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
	"go.uber.org/atomic"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
)

func TestIndexActiveBlockRotate(t *testing.T) {
	var (
		testNsID        = ident.StringID("testns")
		numWrites       = 50
		numTags         = 10
		blockSize       = 2 * time.Hour
		indexBlockSize  = blockSize * 2
		retentionPeriod = 12 * blockSize
		bufferPast      = 10 * time.Minute
		rOpts           = retention.NewOptions().
				SetRetentionPeriod(retentionPeriod).
				SetBlockSize(blockSize).
				SetBufferPast(bufferPast)

		idxOpts = namespace.NewIndexOptions().SetEnabled(true).SetBlockSize(indexBlockSize)
		nsOpts  = namespace.NewOptions().
			SetRetentionOptions(rOpts).
			SetIndexOptions(idxOpts).
			SetColdWritesEnabled(true)

		defaultTimeout = time.Minute
		// verifyTimeout  = time.Minute
	)
	ns, err := namespace.NewMetadata(testNsID, nsOpts)
	require.NoError(t, err)

	// Set time to next warm flushable block transition
	// (i.e. align by block + bufferPast - time.Second)
	currTime := time.Now().UTC()
	progressTime := false
	progressTimeDelta := time.Duration(0)
	lockTime := sync.RWMutex{}
	setTime := func(t time.Time) {
		lockTime.Lock()
		defer lockTime.Unlock()
		progressTime = false
		currTime = t.UTC()
	}
	setProgressTime := func() {
		lockTime.Lock()
		defer lockTime.Unlock()
		progressTime = true
		actualNow := time.Now().UTC()
		progressTimeDelta = currTime.Sub(actualNow)
	}
	nowFn := func() time.Time {
		lockTime.RLock()
		at := currTime
		progress := progressTime
		progressDelta := progressTimeDelta
		lockTime.RUnlock()
		if progress {
			return time.Now().UTC().Add(progressDelta)
		}
		return at
	}

	testOpts := NewTestOptions(t).
		SetNamespaces([]namespace.Metadata{ns}).
		SetWriteNewSeriesAsync(true).
		SetNowFn(nowFn)

	testSetup, err := NewTestSetup(t, testOpts, nil)
	require.NoError(t, err)
	defer testSetup.Close()

	// Write test data to disk so that there's some blocks on disk to simulate
	// some index blocks already having on disk segments already in them.
	require.NoError(t, testSetup.InitializeBootstrappers(InitializeBootstrappersOptions{
		WithFileSystem: true,
	}))
	now := testSetup.NowFn()()
	fooSeries := generate.Series{
		ID:   ident.StringID("foo"),
		Tags: ident.NewTags(ident.StringTag("city", "new_york")),
	}
	barSeries := generate.Series{
		ID:   ident.StringID("bar"),
		Tags: ident.NewTags(ident.StringTag("city", "new_jersey")),
	}
	seriesMaps := generate.BlocksByStart([]generate.BlockConfig{
		{
			IDs:       []string{fooSeries.ID.String()},
			Tags:      fooSeries.Tags,
			NumPoints: 100,
			Start:     now.Add(-3 * blockSize),
		},
		{
			IDs:       []string{barSeries.ID.String()},
			Tags:      barSeries.Tags,
			NumPoints: 100,
			Start:     now.Add(-3 * blockSize),
		},
	})
	require.NoError(t, writeTestDataToDisk(ns, testSetup, seriesMaps, 0))

	// Set foreground compaction planner options to force index compaction.
	minCompactSize := 10
	foregroundCompactionOpts := compaction.DefaultOptions
	foregroundCompactionOpts.Levels = []compaction.Level{
		{
			MinSizeInclusive: 0,
			MaxSizeExclusive: int64(minCompactSize),
		},
	}
	indexOpts := testSetup.StorageOpts().IndexOptions().
		SetForegroundCompactionPlannerOptions(foregroundCompactionOpts)
	testSetup.SetStorageOpts(testSetup.StorageOpts().SetIndexOptions(indexOpts))

	// Configure log capture
	log := testSetup.StorageOpts().InstrumentOptions().Logger()
	captureCore, logs := observer.New(zapcore.ErrorLevel)
	zapOpt := zap.WrapCore(func(existingCore zapcore.Core) zapcore.Core {
		return zapcore.NewTee(existingCore, captureCore)
	})
	log = log.WithOptions(zapOpt)

	// Wire up logger.
	instrumentOpts := testSetup.StorageOpts().InstrumentOptions().
		SetLogger(log)
	testSetup.SetStorageOpts(testSetup.StorageOpts().SetInstrumentOptions(instrumentOpts))
	scope := testSetup.Scope()

	// Start the server.
	require.NoError(t, testSetup.StartServer())

	// Stop the server.
	defer func() {
		require.NoError(t, testSetup.StopServer())
		log.Debug("server is now down")
	}()

	// Write test data.
	session, err := testSetup.M3DBClient().DefaultSession()
	require.NoError(t, err)

	var (
		metricGCSeries   = "index.block.active-block.gc-series+namespace=" + testNsID.String()
		metricFlushIndex = "database.flushIndex.success+namespace=" + testNsID.String()
	)
	prevWarmFlushes := counterValue(t, scope, metricFlushIndex)
	prevNumGCSeries := 0
	numGCSeries := counterValue(t, scope, metricGCSeries)
	require.Equal(t, 0, numGCSeries)

	prevLog := log
	for i := 0; i < 4; i++ {
		log = prevLog.With(zap.Int("checkIteration", i))

		// Progress to next time just before a flush and freeze (using setTime).
		prevTime := nowFn()
		newTime := prevTime.
			Truncate(indexBlockSize).
			Add(2 * indexBlockSize)
		setTime(newTime)
		log.Info("progressing time to before next block edge",
			zap.Stringer("prevTime", prevTime),
			zap.Stringer("newTime", newTime))

		start := time.Now()
		log.Info("writing test data")

		t0 := xtime.ToUnixNano(newTime.Add(-1 * (bufferPast / 2)))
		t1 := xtime.ToUnixNano(newTime)
		writesPeriodIter := GenerateTestIndexWrite(i, numWrites, numTags, t0, t1)
		writesPeriodIter.Write(t, testNsID, session)
		log.Info("test data written", zap.Duration("took", time.Since(start)))

		log.Info("waiting till data is indexed")
		indexed := xclock.WaitUntil(func() bool {
			indexedPeriod := writesPeriodIter.NumIndexed(t, testNsID, session)
			return indexedPeriod == len(writesPeriodIter)
		}, 15*time.Second)
		require.True(t, indexed,
			fmt.Sprintf("unexpected data indexed: actual=%d, expected=%d",
				writesPeriodIter.NumIndexedWithOptions(t, testNsID, session, NumIndexedOptions{Logger: log}),
				len(writesPeriodIter)))
		log.Info("verified data is indexed", zap.Duration("took", time.Since(start)))

		newTime = prevTime.
			Truncate(indexBlockSize).
			Add(2 * indexBlockSize).
			Add(bufferPast).
			Add(-100 * time.Millisecond)
		setTime(newTime)
		log.Info("progressing time to before next flush",
			zap.Stringer("prevTime", prevTime),
			zap.Stringer("newTime", newTime))

		log.Info("waiting till warm flush occurs")

		// Resume time progressing by wall clock.
		setProgressTime()

		// Start checks to ensure metrics are visible the whole time.
		checkFailed := atomic.NewUint64(0)
		checkIndexable := func() {
			numGCSeriesBefore := counterValue(t, scope, metricGCSeries)
			indexedPeriod := writesPeriodIter.NumIndexed(t, testNsID, session)
			numGCSeriesAfter := counterValue(t, scope, metricGCSeries)
			if len(writesPeriodIter) != indexedPeriod {
				assert.Equal(t, len(writesPeriodIter), indexedPeriod,
					fmt.Sprintf("some metrics not indexed/visible: actual=%d, expected=%d, numGCBefore=%d, numGCAfter=%d",
						writesPeriodIter.NumIndexedWithOptions(t, testNsID, session, NumIndexedOptions{Logger: log}),
						len(writesPeriodIter),
						numGCSeriesBefore,
						numGCSeriesAfter))
				checkFailed.Inc()
			}
		}

		ticker := time.NewTicker(10 * time.Millisecond)
		stopTickCh := make(chan struct{})
		closedTickCh := make(chan struct{})
		go func() {
			defer func() {
				ticker.Stop()
				close(closedTickCh)
			}()

			for {
				select {
				case <-ticker.C:
					checkIndexable()
				case <-stopTickCh:
					return
				}
			}
		}()

		start = time.Now()
		warmFlushed := xclock.WaitUntil(func() bool {
			return counterValue(t, scope, metricFlushIndex)-prevWarmFlushes > 0
		}, defaultTimeout)
		counter := counterValue(t, scope, metricFlushIndex)
		require.True(t, warmFlushed,
			fmt.Sprintf("warm flush stats: current=%d, previous=%d", counter, prevWarmFlushes))
		log.Info("verified data has been warm flushed", zap.Duration("took", time.Since(start)))
		prevWarmFlushes = counter

		start = time.Now()
		log.Info("waiting for GC of series")

		expectedNumGCSeries := prevNumGCSeries + numWrites - minCompactSize
		gcSeries := xclock.WaitUntil(func() bool {
			// Run background compaction path to check that this path correctly
			// identifies these series as "empty" post the warm flush above.
			// Note: typically this path gets called from just WriteBatch calls but
			// for this test we just explcitly invoke the background compact path.
			for _, ns := range testSetup.DB().Namespaces() {
				idx, err := ns.Index()
				require.NoError(t, err)

				idx.BackgroundCompact()
			}

			numGCSeries := counterValue(t, scope, metricGCSeries)
			return numGCSeries >= expectedNumGCSeries
		}, defaultTimeout)
		numGCSeries := counterValue(t, scope, metricGCSeries)
		require.True(t, gcSeries,
			fmt.Sprintf("unexpected num gc series: actual=%d, expected=%d",
				numGCSeries, expectedNumGCSeries))
		require.True(t, numGCSeries >= expectedNumGCSeries)
		log.Info("verified series have been GC'd", zap.Duration("took", time.Since(start)))
		prevNumGCSeries = numGCSeries

		require.Equal(t, 0, logs.Len(), "errors found in logs during flush/indexing")

		// Keep running indexable check for a few seconds, then progress next iter.
		time.Sleep(5 * time.Second)
		close(stopTickCh)
		<-closedTickCh

		// Ensure check did not fail.
		require.True(t, checkFailed.Load() == 0,
			fmt.Sprintf("check indexable errors: %d", checkFailed.Load()))
	}

	log.Info("checks passed")
}

func counterValue(t *testing.T, r tally.TestScope, key string) int {
	v, ok := r.Snapshot().Counters()[key]
	require.True(t, ok)
	return int(v.Value())
}


<|start_filename|>src/cluster/client/etcd/types.go<|end_filename|>
// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package etcd

import (
	"crypto/tls"
	"os"
	"time"

	"github.com/m3db/m3/src/cluster/client"
	"github.com/m3db/m3/src/cluster/services"
	"github.com/m3db/m3/src/x/instrument"
	"github.com/m3db/m3/src/x/retry"

	"google.golang.org/grpc"
)

// Options is the Options to create a config service client.
type Options interface {
	// RequestTimeout is the timeout for etcd requests
	RequestTimeout() time.Duration
	// SetRequestTimeout sets the RequestTimeout
	SetRequestTimeout(t time.Duration) Options

	Env() string
	SetEnv(e string) Options

	Zone() string
	SetZone(z string) Options

	Service() string
	SetService(id string) Options

	CacheDir() string
	SetCacheDir(dir string) Options

	ServicesOptions() services.Options
	SetServicesOptions(opts services.Options) Options

	Clusters() []Cluster
	SetClusters(clusters []Cluster) Options
	ClusterForZone(z string) (Cluster, bool)

	InstrumentOptions() instrument.Options
	SetInstrumentOptions(iopts instrument.Options) Options

	RetryOptions() retry.Options
	SetRetryOptions(retryOpts retry.Options) Options

	// WatchChanCheckInterval will be used to periodically check if a watch chan
	// is no longer being subscribed and should be closed
	WatchChanCheckInterval() time.Duration
	// SetWatchChanCheckInterval sets the WatchChanCheckInterval
	SetWatchChanCheckInterval(t time.Duration) Options

	// WatchChanResetInterval is the delay before resetting the etcd watch chan
	WatchChanResetInterval() time.Duration
	// SetWatchChanResetInterval sets the WatchChanResetInterval
	SetWatchChanResetInterval(t time.Duration) Options

	// WatchChanInitTimeout is the timeout for a watchChan initialization
	WatchChanInitTimeout() time.Duration
	// SetWatchChanInitTimeout sets the WatchChanInitTimeout
	SetWatchChanInitTimeout(t time.Duration) Options

	WatchWithRevision() int64
	SetWatchWithRevision(rev int64) Options

	// EnableFastGets returns whether to use clientv3.WithSerializable() option to speed up gets.
	EnableFastGets() bool
	// SetEnableFastGets sets clientv3.WithSerializable() to speed up gets, but can fetch stale data.
	SetEnableFastGets(enabled bool) Options

	SetNewDirectoryMode(fm os.FileMode) Options
	NewDirectoryMode() os.FileMode

	Validate() error
}

// KeepAliveOptions provide a set of client-side keepAlive options.
type KeepAliveOptions interface {
	// KeepAliveEnabled determines whether keepAlives are enabled.
	KeepAliveEnabled() bool

	// SetKeepAliveEnabled sets whether keepAlives are enabled.
	SetKeepAliveEnabled(value bool) KeepAliveOptions

	// KeepAlivePeriod is the duration of time after which if the client doesn't see
	// any activity the client pings the server to see if transport is alive.
	KeepAlivePeriod() time.Duration

	// SetKeepAlivePeriod sets the duration of time after which if the client doesn't see
	// any activity the client pings the server to see if transport is alive.
	SetKeepAlivePeriod(value time.Duration) KeepAliveOptions

	// KeepAlivePeriodMaxJitter is used to add some jittering to keep alive period
	// to avoid a large number of clients all sending keepalive probes at the
	// same time.
	KeepAlivePeriodMaxJitter() time.Duration

	// SetKeepAlivePeriodMaxJitter sets the maximum jittering to keep alive period
	// to avoid a large number of clients all sending keepalive probes at the
	// same time.
	SetKeepAlivePeriodMaxJitter(value time.Duration) KeepAliveOptions

	// KeepAliveTimeout is the time that the client waits for a response for the
	// keep-alive probe. If the response is not received in this time, the connection is closed.
	KeepAliveTimeout() time.Duration

	// SetKeepAliveTimeout sets the time that the client waits for a response for the
	// keep-alive probe. If the response is not received in this time, the connection is closed.
	SetKeepAliveTimeout(value time.Duration) KeepAliveOptions
}

// TLSOptions defines the options for TLS.
type TLSOptions interface {
	CrtPath() string
	SetCrtPath(string) TLSOptions

	KeyPath() string
	SetKeyPath(string) TLSOptions

	CACrtPath() string
	SetCACrtPath(string) TLSOptions

	Config() (*tls.Config, error)
}

// Cluster defines the configuration for a etcd cluster.
type Cluster interface {
	Zone() string
	SetZone(z string) Cluster

	Endpoints() []string
	SetEndpoints(endpoints []string) Cluster

	KeepAliveOptions() KeepAliveOptions
	SetKeepAliveOptions(value KeepAliveOptions) Cluster

	TLSOptions() TLSOptions
	SetTLSOptions(TLSOptions) Cluster

	AutoSyncInterval() time.Duration
	SetAutoSyncInterval(value time.Duration) Cluster

	DialTimeout() time.Duration
	SetDialTimeout(value time.Duration) Cluster

	// DialOptions are used by the etcd client to connect to etcd peers. They can be used to customize
	// the low level networking behavior of the client (e.g. to forward connections over a proxy).
	DialOptions() []grpc.DialOption
	SetDialOptions(opts []grpc.DialOption) Cluster
}

// Client is an etcd-backed m3cluster client.
type Client interface {
	client.Client

	Clients() []ZoneClient
}


<|start_filename|>src/aggregator/integration/topic.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

import (
	"testing"
	"time"

	clusterclient "github.com/m3db/m3/src/cluster/client"
	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/services"
	"github.com/m3db/m3/src/msg/topic"

	"github.com/stretchr/testify/require"
)

func initializeTopic(
	topicName string, //nolint:unparam
	clusterClient clusterclient.Client,
	numShards int,
) (topic.Service, error) {
	serviceID := services.NewServiceID().SetName(defaultServiceName)
	cs := topic.NewConsumerService().
		SetServiceID(serviceID).
		SetConsumptionType(topic.Replicated).
		SetMessageTTLNanos(time.Minute.Nanoseconds())
	ingestTopic := topic.NewTopic().
		SetName(topicName).
		SetNumberOfShards(uint32(numShards)).
		SetConsumerServices([]topic.ConsumerService{cs})

	topicServiceOpts := topic.NewServiceOptions().
		SetConfigService(clusterClient)
	topicService, err := topic.NewService(topicServiceOpts)
	if err != nil {
		return topicService, err
	}

	_, err = topicService.CheckAndSet(ingestTopic, 0)

	return topicService, err
}

func removeAllTopicConsumers(
	topicService topic.Service,
	topicName string, //nolint:unparam
) error {
	topic, err := topicService.Get(topicName)
	if err != nil {
		return err
	}
	for len(topic.ConsumerServices()) > 0 {
		topic, err = topic.RemoveConsumerService(topic.ConsumerServices()[0].ServiceID())
		if err != nil {
			return err
		}
	}
	_, err = topicService.CheckAndSet(topic, topic.Version())
	return err
}

func setupTopic(t *testing.T, serverOpts testServerOptions, placement placement.Placement) testServerOptions {
	topicService, err := initializeTopic(defaultTopicName, serverOpts.ClusterClient(), placement.NumShards())
	require.NoError(t, err)
	return serverOpts.
		SetTopicService(topicService).
		SetTopicName(defaultTopicName)
}


<|start_filename|>src/query/api/v1/httpd/handler.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package httpd

import (
	"encoding/json"
	"fmt"
	"net/http"
	// needed for pprof handler registration
	_ "net/http/pprof"
	"time"

	"github.com/m3db/m3/src/cluster/placementhandler"
	"github.com/m3db/m3/src/cluster/placementhandler/handleroptions"
	"github.com/m3db/m3/src/cmd/services/m3query/config"
	"github.com/m3db/m3/src/query/api/v1/handler"
	"github.com/m3db/m3/src/query/api/v1/handler/database"
	"github.com/m3db/m3/src/query/api/v1/handler/graphite"
	"github.com/m3db/m3/src/query/api/v1/handler/influxdb"
	m3json "github.com/m3db/m3/src/query/api/v1/handler/json"
	"github.com/m3db/m3/src/query/api/v1/handler/namespace"
	"github.com/m3db/m3/src/query/api/v1/handler/openapi"
	"github.com/m3db/m3/src/query/api/v1/handler/prom"
	"github.com/m3db/m3/src/query/api/v1/handler/prometheus/native"
	"github.com/m3db/m3/src/query/api/v1/handler/prometheus/remote"
	"github.com/m3db/m3/src/query/api/v1/handler/topic"
	"github.com/m3db/m3/src/query/api/v1/middleware"
	"github.com/m3db/m3/src/query/api/v1/options"
	"github.com/m3db/m3/src/query/api/v1/route"
	"github.com/m3db/m3/src/query/parser/promql"
	"github.com/m3db/m3/src/query/util/queryhttp"
	xdebug "github.com/m3db/m3/src/x/debug"
	xhttp "github.com/m3db/m3/src/x/net/http"

	"github.com/gorilla/mux"
	"github.com/jonboulle/clockwork"
	"go.uber.org/zap"
)

const (
	healthURL = "/health"
	routesURL = "/routes"

	// EngineURLParam defines query url parameter which is used to switch between
	// prometheus and m3query engines.
	EngineURLParam = "engine"
)

var (
	remoteSource = map[string]string{"source": "remote"}
	nativeSource = map[string]string{"source": "native"}

	v1APIGroup = map[string]string{"api_group": "v1"}
)

// Handler represents the top-level HTTP handler.
type Handler struct {
	registry         *queryhttp.EndpointRegistry
	handler          *mux.Router
	options          options.HandlerOptions
	customHandlers   []options.CustomHandler
	logger           *zap.Logger
	middlewareConfig config.MiddlewareConfiguration
}

// Router returns the http handler registered with all relevant routes for query.
func (h *Handler) Router() http.Handler {
	return h.handler
}

// NewHandler returns a new instance of handler with routes.
func NewHandler(
	handlerOptions options.HandlerOptions,
	middlewareConfig config.MiddlewareConfiguration,
	customHandlers ...options.CustomHandler,
) *Handler {
	var (
		r        = mux.NewRouter()
		logger   = handlerOptions.InstrumentOpts().Logger()
		registry = queryhttp.NewEndpointRegistry(r)
	)
	return &Handler{
		registry:         registry,
		handler:          r,
		options:          handlerOptions,
		customHandlers:   customHandlers,
		logger:           logger,
		middlewareConfig: middlewareConfig,
	}
}

// RegisterRoutes registers all http routes.
func (h *Handler) RegisterRoutes() error {
	instrumentOpts := h.options.InstrumentOpts()

	// OpenAPI.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    openapi.URL,
		Handler: openapi.NewDocHandler(instrumentOpts),
		Methods: methods(openapi.HTTPMethod),
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		PathPrefix: openapi.StaticURLPrefix,
		Handler:    openapi.StaticHandler(),
	}); err != nil {
		return err
	}

	// Prometheus remote read/write endpoints.
	remoteSourceOpts := h.options.SetInstrumentOpts(instrumentOpts.
		SetMetricsScope(instrumentOpts.MetricsScope().
			Tagged(remoteSource).
			Tagged(v1APIGroup),
		))

	promRemoteReadHandler := remote.NewPromReadHandler(remoteSourceOpts)
	promRemoteWriteHandler, err := remote.NewPromWriteHandler(remoteSourceOpts)
	if err != nil {
		return err
	}

	nativeSourceOpts := h.options.SetInstrumentOpts(instrumentOpts.
		SetMetricsScope(instrumentOpts.MetricsScope().
			Tagged(nativeSource).
			Tagged(v1APIGroup),
		))

	promqlQueryHandler, err := prom.NewReadHandler(nativeSourceOpts,
		prom.WithEngine(h.options.PrometheusEngineFn()))
	if err != nil {
		return err
	}
	promqlInstantQueryHandler, err := prom.NewReadHandler(nativeSourceOpts,
		prom.WithInstantEngine(h.options.PrometheusEngineFn()))
	if err != nil {
		return err
	}
	nativePromReadHandler := native.NewPromReadHandler(nativeSourceOpts)
	nativePromReadInstantHandler := native.NewPromReadInstantHandler(nativeSourceOpts)

	h.options.QueryRouter().Setup(options.QueryRouterOptions{
		DefaultQueryEngine: h.options.DefaultQueryEngine(),
		PromqlHandler:      promqlQueryHandler.ServeHTTP,
		M3QueryHandler:     nativePromReadHandler.ServeHTTP,
	})

	h.options.InstantQueryRouter().Setup(options.QueryRouterOptions{
		DefaultQueryEngine: h.options.DefaultQueryEngine(),
		PromqlHandler:      promqlInstantQueryHandler.ServeHTTP,
		M3QueryHandler:     nativePromReadInstantHandler.ServeHTTP,
	})

	// Query routable endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               native.PromReadURL,
		Handler:            h.options.QueryRouter(),
		Methods:            native.PromReadHTTPMethods,
		MiddlewareOverride: native.WithRangeQueryParamsAndRangeRewriting,
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               native.PromReadInstantURL,
		Handler:            h.options.InstantQueryRouter(),
		Methods:            native.PromReadInstantHTTPMethods,
		MiddlewareOverride: native.WithInstantQueryParamsAndRangeRewriting,
	}); err != nil {
		return err
	}

	// Prometheus endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               "/prometheus" + native.PromReadURL,
		Handler:            promqlQueryHandler,
		Methods:            native.PromReadHTTPMethods,
		MiddlewareOverride: native.WithRangeQueryParamsAndRangeRewriting,
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               "/prometheus" + native.PromReadInstantURL,
		Handler:            promqlInstantQueryHandler,
		Methods:            native.PromReadInstantHTTPMethods,
		MiddlewareOverride: native.WithInstantQueryParamsAndRangeRewriting,
	}); err != nil {
		return err
	}

	// M3Query endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               "/m3query" + native.PromReadURL,
		Handler:            nativePromReadHandler,
		Methods:            native.PromReadHTTPMethods,
		MiddlewareOverride: native.WithRangeQueryParamsAndRangeRewriting,
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               "/m3query" + native.PromReadInstantURL,
		Handler:            nativePromReadInstantHandler,
		Methods:            native.PromReadInstantHTTPMethods,
		MiddlewareOverride: native.WithInstantQueryParamsAndRangeRewriting,
	}); err != nil {
		return err
	}

	// Prometheus remote read and write endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    remote.PromReadURL,
		Handler: promRemoteReadHandler,
		Methods: remote.PromReadHTTPMethods,
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    remote.PromWriteURL,
		Handler: promRemoteWriteHandler,
		Methods: methods(remote.PromWriteHTTPMethod),
		// Register with no response logging for write calls since so frequent.
		MiddlewareOverride: middleware.WithNoResponseLogging,
	}); err != nil {
		return err
	}

	// InfluxDB write endpoint.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    influxdb.InfluxWriteURL,
		Handler: influxdb.NewInfluxWriterHandler(h.options),
		Methods: methods(influxdb.InfluxWriteHTTPMethod),
		// Register with no response logging for write calls since so frequent.
		MiddlewareOverride: middleware.WithNoResponseLogging,
	}); err != nil {
		return err
	}

	// Native M3 search and write endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    handler.SearchURL,
		Handler: handler.NewSearchHandler(h.options),
		Methods: methods(handler.SearchHTTPMethod),
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    m3json.WriteJSONURL,
		Handler: m3json.NewWriteJSONHandler(h.options),
		Methods: methods(m3json.JSONWriteHTTPMethod),
	}); err != nil {
		return err
	}

	// Readiness endpoint.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    handler.ReadyURL,
		Handler: handler.NewReadyHandler(h.options),
		Methods: methods(handler.ReadyHTTPMethod),
	}); err != nil {
		return err
	}

	// Tag completion endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               native.CompleteTagsURL,
		Handler:            native.NewCompleteTagsHandler(h.options),
		Methods:            methods(native.CompleteTagsHTTPMethod),
		MiddlewareOverride: native.WithQueryParams,
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               remote.TagValuesURL,
		Handler:            remote.NewTagValuesHandler(h.options),
		Methods:            methods(remote.TagValuesHTTPMethod),
		MiddlewareOverride: native.WithQueryParams,
	}); err != nil {
		return err
	}

	// List tag endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               native.ListTagsURL,
		Handler:            native.NewListTagsHandler(h.options),
		Methods:            native.ListTagsHTTPMethods,
		MiddlewareOverride: native.WithQueryParams,
	}); err != nil {
		return err
	}

	// Query parse endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    native.PromParseURL,
		Handler: native.NewPromParseHandler(h.options),
		Methods: methods(native.PromParseHTTPMethod),
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    native.PromThresholdURL,
		Handler: native.NewPromThresholdHandler(h.options),
		Methods: methods(native.PromThresholdHTTPMethod),
	}); err != nil {
		return err
	}

	// Series match endpoints.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:               route.SeriesMatchURL,
		Handler:            remote.NewPromSeriesMatchHandler(h.options),
		Methods:            remote.PromSeriesMatchHTTPMethods,
		MiddlewareOverride: native.WithQueryParams,
	}); err != nil {
		return err
	}

	// Graphite routable endpoints.
	h.options.GraphiteRenderRouter().Setup(options.GraphiteRenderRouterOptions{
		RenderHandler: graphite.NewRenderHandler(h.options).ServeHTTP,
	})
	h.options.GraphiteFindRouter().Setup(options.GraphiteFindRouterOptions{
		FindHandler: graphite.NewFindHandler(h.options).ServeHTTP,
	})
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    graphite.ReadURL,
		Handler: h.options.GraphiteRenderRouter(),
		Methods: graphite.ReadHTTPMethods,
	}); err != nil {
		return err
	}
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    graphite.FindURL,
		Handler: h.options.GraphiteFindRouter(),
		Methods: graphite.FindHTTPMethods,
	}); err != nil {
		return err
	}

	placementOpts, err := h.placementOpts()
	if err != nil {
		return err
	}

	var (
		serviceOptionDefaults = h.options.ServiceOptionDefaults()
		clusterClient         = h.options.ClusterClient()
		config                = h.options.Config()
	)

	debugWriter, err := xdebug.NewZipWriterWithDefaultSources(
		h.options.CPUProfileDuration(),
		instrumentOpts)
	if err != nil {
		return fmt.Errorf("unable to create debug writer: %v", err)
	}

	// Register debug dump handler.
	if err := h.registry.Register(queryhttp.RegisterOptions{
		Path:    xdebug.DebugURL,
		Handler: debugWriter.HTTPHandler(),
		Methods: methods(xdebug.DebugMethod),
	}); err != nil {
		return err
	}

	if clusterClient != nil {
		err = database.RegisterRoutes(h.registry, clusterClient,
			h.options.Config(), h.options.EmbeddedDBCfg(),
			serviceOptionDefaults, instrumentOpts,
			h.options.NamespaceValidator(), h.options.KVStoreProtoParser())
		if err != nil {
			return err
		}

		routes := placementhandler.MakeRoutes(serviceOptionDefaults, placementOpts)
		for _, route := range routes {
			err := h.registry.RegisterPaths(route.Paths, queryhttp.RegisterPathsOptions{
				Handler: route.Handler,
				Methods: route.Methods,
			})
			if err != nil {
				return err
			}
		}

		err = namespace.RegisterRoutes(h.registry, clusterClient,
			h.options.Clusters(), serviceOptionDefaults, instrumentOpts,
			h.options.NamespaceValidator())
		if err != nil {
			return err
		}

		err = topic.RegisterRoutes(h.registry, clusterClient, config, instrumentOpts)
		if err != nil {
			return err
		}
	}

	if err := h.registerHealthEndpoints(); err != nil {
		return err
	}
	if err := h.registerProfileEndpoints(); err != nil {
		return err
	}
	if err := h.registerRoutesEndpoint(); err != nil {
		return err
	}

	customMiddle := make(map[*mux.Route]middleware.OverrideOptions)
	// Register custom endpoints last to have these conflict with
	// any existing routes.
	for _, custom := range h.customHandlers {
		for _, method := range custom.Methods() {
			var prevHandler http.Handler
			entry, prevRoute := h.registry.PathEntry(custom.Route(), method)
			if prevRoute {
				prevHandler = entry.GetHandler()
			}

			handler, err := custom.Handler(nativeSourceOpts, prevHandler)
			if err != nil {
				return err
			}

			if !prevRoute {
				if err := h.registry.Register(queryhttp.RegisterOptions{
					Path:               custom.Route(),
					Handler:            handler,
					Methods:            methods(method),
					MiddlewareOverride: custom.MiddlewareOverride(),
				}); err != nil {
					return err
				}
			} else {
				customMiddle[entry] = custom.MiddlewareOverride()
				entry.Handler(handler)
			}
		}
	}

	// NB: the double http_handler was accidentally introduced and now we are
	// stuck with it for backwards compatibility.
	middleIOpts := instrumentOpts.SetMetricsScope(
		h.options.InstrumentOpts().MetricsScope().SubScope("http_handler_http_handler"))

	// Apply middleware after the custom handlers have overridden the previous handlers so the middleware functions
	// are dispatched before the custom handler.
	// req -> middleware fns -> custom handler -> previous handler.
	err = h.registry.Walk(func(route *mux.Route, router *mux.Router, ancestors []*mux.Route) error {
		handler := route.GetHandler()
		opts := middleware.Options{
			InstrumentOpts: middleIOpts,
			Route:          route,
			Clock:          clockwork.NewRealClock(),
			Logging:        middleware.NewLoggingOptions(h.middlewareConfig.Logging),
			Metrics: middleware.MetricsOptions{
				Config: h.middlewareConfig.Metrics,
				ParseOptions: promql.NewParseOptions().
					SetRequireStartEndTime(h.options.Config().Query.RequireLabelsEndpointStartEndTime).
					SetNowFn(h.options.NowFn()),
			},
			PrometheusRangeRewrite: middleware.PrometheusRangeRewriteOptions{
				FetchOptionsBuilder:  h.options.FetchOptionsBuilder(),
				ResolutionMultiplier: h.middlewareConfig.Prometheus.ResolutionMultiplier,
				DefaultLookback:      h.options.DefaultLookback(),
				Storage:              h.options.Storage(),
				PrometheusEngineFn:   h.options.PrometheusEngineFn(),
			},
		}
		override := h.registry.MiddlewareOpts(route)
		if override != nil {
			opts = override(opts)
		}
		if customMiddle[route] != nil {
			opts = customMiddle[route](opts)
		}
		middle := h.options.RegisterMiddleware()(opts)

		// iterate through in reverse order so each Middleware fn gets the proper next handler to dispatch. this ensures the
		// Middleware is dispatched in the expected order (first -> last).
		for i := len(middle) - 1; i >= 0; i-- {
			handler = middle[i].Middleware(handler)
		}

		route.Handler(handler)
		return nil
	})

	if err != nil {
		return err
	}

	return nil
}

func (h *Handler) placementOpts() (placementhandler.HandlerOptions, error) {
	return placementhandler.NewHandlerOptions(
		h.options.ClusterClient(),
		h.options.Config().ClusterManagement.Placement,
		h.m3AggServiceOptions(),
		h.options.InstrumentOpts(),
	)
}

func (h *Handler) m3AggServiceOptions() *handleroptions.M3AggServiceOptions {
	clusters := h.options.Clusters()
	if clusters == nil {
		return nil
	}

	maxResolution := time.Duration(0)
	for _, ns := range clusters.ClusterNamespaces() {
		resolution := ns.Options().Attributes().Resolution
		if resolution > maxResolution {
			maxResolution = resolution
		}
	}

	if maxResolution == 0 {
		return nil
	}

	return &handleroptions.M3AggServiceOptions{
		MaxAggregationWindowSize: maxResolution,
	}
}

// Endpoints useful for profiling the service.
func (h *Handler) registerHealthEndpoints() error {
	return h.registry.Register(queryhttp.RegisterOptions{
		Path: healthURL,
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			json.NewEncoder(w).Encode(struct {
				Uptime string `json:"uptime"`
				Now    string `json:"now"`
			}{
				Uptime: time.Since(h.options.CreatedAt()).String(),
				Now:    time.Now().String(),
			})
		}),
		Methods: methods(http.MethodGet),
	})
}

// Endpoints useful for profiling the service.
func (h *Handler) registerProfileEndpoints() error {
	debugHandler := http.NewServeMux()
	xdebug.RegisterPProfHandlers(debugHandler)

	return h.registry.Register(queryhttp.RegisterOptions{
		PathPrefix: "/debug/pprof",
		Handler:    debugHandler,
	})
}

// Endpoints useful for viewing routes directory.
func (h *Handler) registerRoutesEndpoint() error {
	return h.registry.Register(queryhttp.RegisterOptions{
		Path: routesURL,
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			var routes []string
			err := h.registry.Walk(
				func(route *mux.Route, router *mux.Router, ancestors []*mux.Route) error {
					str, err := route.GetPathTemplate()
					if err != nil {
						return err
					}
					routes = append(routes, str)
					return nil
				})
			if err != nil {
				xhttp.WriteError(w, err)
				return
			}
			json.NewEncoder(w).Encode(struct {
				Routes []string `json:"routes"`
			}{
				Routes: routes,
			})
		}),
		Methods: methods(http.MethodGet),
	})
}

func methods(str ...string) []string {
	return str
}


<|start_filename|>src/x/net/http/errors_test.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package xhttp

import (
	"context"
	"errors"
	"fmt"
	"net/http/httptest"
	"testing"

	terrors "github.com/m3db/m3/src/dbnode/network/server/tchannelthrift/errors"
	xerrors "github.com/m3db/m3/src/x/errors"

	"github.com/prometheus/prometheus/promql"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel-go"
)

func TestErrorStatus(t *testing.T) {
	a := terrors.NewInternalError(tchannel.ErrTimeout)

	tests := []struct {
		name           string
		err            error
		expectedStatus int
	}{
		{
			name:           "generic error",
			err:            errors.New("generic error"),
			expectedStatus: 500,
		},
		{
			name:           "invalid params",
			err:            xerrors.NewInvalidParamsError(errors.New("bad param")),
			expectedStatus: 400,
		},
		{
			name:           "deadline exceeded",
			err:            context.DeadlineExceeded,
			expectedStatus: 504,
		},
		{
			name:           "canceled",
			err:            context.Canceled,
			expectedStatus: 499,
		},
		{
			name:           "prom canceled",
			err:            promql.ErrQueryCanceled("canceled"),
			expectedStatus: 499,
		},
		{
			name:           "client timeout",
			err:            terrors.NewTimeoutError(fmt.Errorf("timeout")),
			expectedStatus: 504,
		},
		{
			name:           "tchannel timeout",
			err:            tchannel.ErrTimeout,
			expectedStatus: 504,
		},
		{
			name:           "http error",
			err:            NewError(errors.New("some error"), 504),
			expectedStatus: 504,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assert.Error(t, a)
			recorder := httptest.NewRecorder()
			WriteError(recorder, tt.err)
			assert.Equal(t, tt.expectedStatus, recorder.Code)
		})
	}
}

func TestErrorRewrite(t *testing.T) {
	tests := []struct {
		name           string
		err            error
		expectedBody   string
		expectedStatus int
	}{
		{
			name:           "error that should not be rewritten",
			err:            errors.New("random error"),
			expectedBody:   `{"status":"error","error":"random error"}`,
			expectedStatus: 500,
		},
		{
			name:           "error that should be rewritten",
			err:            xerrors.NewInvalidParamsError(errors.New("to be rewritten")),
			expectedBody:   `{"status":"error","error":"rewritten error"}`,
			expectedStatus: 500,
		},
	}

	invalidParamsRewriteFn := func(err error) error {
		if xerrors.IsInvalidParams(err) {
			return errors.New("rewritten error")
		}
		return err
	}

	SetErrorRewriteFn(invalidParamsRewriteFn)

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			recorder := httptest.NewRecorder()
			WriteError(recorder, tt.err)
			assert.JSONEq(t, tt.expectedBody, recorder.Body.String())
			assert.Equal(t, tt.expectedStatus, recorder.Code)
		})
	}
}

func TestIsClientError(t *testing.T) {
	tests := []struct {
		err      error
		expected bool
	}{
		{NewError(fmt.Errorf("xhttp.Error(400)"), 400), true},
		{NewError(fmt.Errorf("xhttp.Error(499)"), 499), true},
		{xerrors.NewInvalidParamsError(fmt.Errorf("InvalidParamsError")), true},
		{xerrors.NewRetryableError(xerrors.NewInvalidParamsError(
			fmt.Errorf("InvalidParamsError insde RetyrableError"))), true},

		{NewError(fmt.Errorf("xhttp.Error(399)"), 399), false},
		{NewError(fmt.Errorf("xhttp.Error(500)"), 500), false},
		{xerrors.NewRetryableError(fmt.Errorf("any error inside RetryableError")), false},
		{fmt.Errorf("any error"), false},
	}

	for _, tt := range tests {
		t.Run(tt.err.Error(), func(t *testing.T) {
			require.Equal(t, tt.expected, IsClientError(tt.err))
		})
	}
}


<|start_filename|>src/aggregator/aggregator/aggregator_test.go<|end_filename|>
// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package aggregator

import (
	"errors"
	"math"
	"sort"
	"testing"
	"time"

	"github.com/m3db/m3/src/aggregator/aggregator/handler"
	"github.com/m3db/m3/src/aggregator/aggregator/handler/writer"
	"github.com/m3db/m3/src/aggregator/client"
	schema "github.com/m3db/m3/src/aggregator/generated/proto/flush"
	"github.com/m3db/m3/src/cluster/generated/proto/placementpb"
	"github.com/m3db/m3/src/cluster/kv"
	"github.com/m3db/m3/src/cluster/kv/mem"
	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/shard"
	"github.com/m3db/m3/src/metrics/aggregation"
	"github.com/m3db/m3/src/metrics/metadata"
	"github.com/m3db/m3/src/metrics/metric"
	"github.com/m3db/m3/src/metrics/metric/aggregated"
	"github.com/m3db/m3/src/metrics/metric/unaggregated"
	"github.com/m3db/m3/src/metrics/pipeline"
	"github.com/m3db/m3/src/metrics/pipeline/applied"
	"github.com/m3db/m3/src/metrics/policy"
	"github.com/m3db/m3/src/metrics/transformation"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/instrument"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/gogo/protobuf/types"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
	"go.uber.org/zap"
)

const (
	testShardSetID       = uint32(1)
	testInstanceID       = "localhost:0"
	testPlacementKey     = "placement"
	testNumShards        = 4
	testPlacementCutover = 1234
)

var (
	testUntimedMetric = unaggregated.MetricUnion{
		Type:       metric.CounterType,
		ID:         []byte("foo"),
		CounterVal: 1234,
	}
	testTimedMetric = aggregated.Metric{
		Type:      metric.CounterType,
		ID:        []byte("testTimed"),
		TimeNanos: 12345,
		Value:     1000,
	}
	testForwardedMetric = aggregated.ForwardedMetric{
		Type:      metric.CounterType,
		ID:        []byte("testForwarded"),
		TimeNanos: 12345,
		Values:    []float64{76109, 23891},
	}
	testPassthroughMetric = aggregated.Metric{
		Type:      metric.CounterType,
		ID:        []byte("testPassthrough"),
		TimeNanos: 12345,
		Value:     1000,
	}
	testInvalidMetric = unaggregated.MetricUnion{
		Type: metric.UnknownType,
		ID:   []byte("testInvalid"),
	}
	testStagedMetadatas = metadata.StagedMetadatas{
		{
			CutoverNanos: 123,
			Tombstoned:   false,
			Metadata: metadata.Metadata{
				Pipelines: []metadata.PipelineMetadata{
					{
						AggregationID: aggregation.DefaultID,
						StoragePolicies: []policy.StoragePolicy{
							policy.NewStoragePolicy(10*time.Second, xtime.Second, 6*time.Hour),
							policy.NewStoragePolicy(time.Minute, xtime.Minute, 2*24*time.Hour),
						},
					},
				},
			},
		},
		{
			CutoverNanos: 456,
			Tombstoned:   false,
			Metadata: metadata.Metadata{
				Pipelines: []metadata.PipelineMetadata{
					{},
				},
			},
		},
	}
	testTimedMetadata = metadata.TimedMetadata{
		AggregationID: aggregation.MustCompressTypes(aggregation.Sum),
		StoragePolicy: policy.NewStoragePolicy(time.Minute, xtime.Minute, 12*time.Hour),
	}
	testForwardMetadata = metadata.ForwardMetadata{
		AggregationID: aggregation.DefaultID,
		StoragePolicy: policy.NewStoragePolicy(time.Minute, xtime.Minute, 12*time.Hour),
		Pipeline: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo"),
					AggregationID: aggregation.MustCompressTypes(aggregation.Count),
				},
			},
		}),
		SourceID:          1234,
		NumForwardedTimes: 3,
	}
	testPassthroughStroagePolicy = policy.NewStoragePolicy(time.Minute, xtime.Minute, 12*time.Hour)
)

func TestAggregatorOpenAlreadyOpen(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	agg.state = aggregatorOpen
	require.Equal(t, errAggregatorAlreadyOpenOrClosed, agg.Open())
}

func TestAggregatorOpenPlacementManagerOpenError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	errPlacementManagerOpen := errors.New("error opening placement manager")
	placementManager := NewMockPlacementManager(ctrl)
	placementManager.EXPECT().Open().Return(errPlacementManagerOpen)

	agg, _ := testAggregator(t, ctrl)
	agg.placementManager = placementManager
	require.Equal(t, errPlacementManagerOpen, agg.Open())
}

func TestAggregatorOpenPlacementError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	errPlacement := errors.New("error getting placement")
	placementManager := NewMockPlacementManager(ctrl)
	placementManager.EXPECT().Open().Return(nil)
	placementManager.EXPECT().Placement().Return(nil, errPlacement)

	agg, _ := testAggregator(t, ctrl)
	agg.placementManager = placementManager
	require.Equal(t, errPlacement, agg.Open())
}

func TestAggregatorOpenInstanceFromError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	testPlacement := placement.NewPlacement().SetCutoverNanos(5678)
	errInstanceFrom := errors.New("error getting instance from placement")
	placementManager := NewMockPlacementManager(ctrl)
	placementManager.EXPECT().Open().Return(nil)
	placementManager.EXPECT().Placement().Return(testPlacement, nil)
	placementManager.EXPECT().InstanceFrom(testPlacement).Return(nil, errInstanceFrom)

	agg, _ := testAggregator(t, ctrl)
	agg.placementManager = placementManager
	require.Equal(t, errInstanceFrom, agg.Open())
}

func TestAggregatorOpenInstanceNotInPlacement(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	placementManager := NewMockPlacementManager(ctrl)

	agg, _ := testAggregator(t, ctrl)
	agg.placementManager = placementManager

	testPlacement := placement.NewPlacement().SetCutoverNanos(5678)

	placementManager.EXPECT().Open().Return(nil)
	placementManager.EXPECT().C().Return(make(chan struct{})).AnyTimes()
	placementManager.EXPECT().InstanceID().Return(agg.opts.PlacementManager().InstanceID())
	placementManager.EXPECT().Placement().Return(testPlacement, nil).AnyTimes()
	placementManager.EXPECT().InstanceFrom(testPlacement).Return(nil, ErrInstanceNotFoundInPlacement)

	require.NoError(t, agg.Open())
	require.Equal(t, uint32(0), agg.shardSetID)
	require.False(t, agg.shardSetOpen)
	require.Equal(t, 0, len(agg.shardIDs))
	require.Nil(t, agg.shards)
	require.Equal(t, agg.currPlacement, testPlacement)
	require.Equal(t, aggregatorOpen, agg.state)
}

func TestAggregatorOpenSuccess(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	require.Equal(t, testShardSetID, agg.shardSetID)
	require.True(t, agg.shardSetOpen)
	require.Equal(t, aggregatorOpen, agg.state)
	require.Equal(t, []uint32{0, 1, 2, 3}, agg.shardIDs)
	for i := 0; i < testNumShards; i++ {
		require.NotNil(t, agg.shards[i])
	}
	require.NotNil(t, agg.currPlacement)
	require.Equal(t, int64(testPlacementCutover), agg.currPlacement.CutoverNanos())
}

func TestAggregatorUpdateStagedMetadatas(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	cases := []struct {
		name       string
		addToReset bool
	}{
		{
			name:       "enabled",
			addToReset: true,
		},
		{
			name:       "disabled",
			addToReset: false,
		},
	}

	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			agg, _ := testAggregator(t, ctrl)
			agg.opts = agg.opts.SetAddToReset(tc.addToReset)
			require.NoError(t, agg.Open())
			agg.shardFn = func([]byte, uint32) uint32 { return 1 }
			sms := metadata.StagedMetadatas{
				{
					Metadata: metadata.Metadata{
						Pipelines: []metadata.PipelineMetadata{
							{
								Pipeline: applied.NewPipeline([]applied.OpUnion{
									{
										Type:           pipeline.TransformationOpType,
										Transformation: pipeline.TransformationOp{Type: transformation.Add},
									},
								}),
							},
						},
					},
				},
			}
			err := agg.AddUntimed(testUntimedMetric, sms)
			require.NoError(t, err)
			require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
			for _, e := range agg.shards[1].metricMap.entries {
				actual := e.Value.(hashedEntry).entry.aggregations[0].key.pipeline
				if tc.addToReset {
					require.Equal(t, transformation.Reset, actual.Operations[0].Transformation.Type)
				} else {
					require.Equal(t, transformation.Add, actual.Operations[0].Transformation.Type)
				}
			}
		})
	}
}

func TestAggregatorInstanceNotFoundThenFoundThenNotFound(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	placements := []*placementpb.PlacementSnapshots{
		{
			Snapshots: []*placementpb.Placement{
				{
					CutoverTime: 0,
				},
			},
		},
		testStagedPlacementProtoWithNumShards(t, testInstanceID, testShardSetID, 4),
		{
			Snapshots: []*placementpb.Placement{
				{
					CutoverTime: testPlacementCutover + 1000,
				},
			},
		},
	}

	// Instance not in the first placement.
	agg, store := testAggregatorWithCustomPlacements(t, ctrl, placements[0])
	require.NoError(t, agg.Open())
	require.Equal(t, uint32(0), agg.shardSetID)
	require.False(t, agg.shardSetOpen)
	require.Equal(t, 0, len(agg.shardIDs))
	require.Nil(t, agg.shards)
	require.Equal(t, int64(0), agg.currPlacement.CutoverNanos())
	require.Equal(t, aggregatorOpen, agg.state)

	// Update placement to add the instance and wait for it to propagate.
	_, err := store.Set(testPlacementKey, placements[1])
	require.NoError(t, err)
	for {
		p, err := agg.placementManager.Placement()
		require.NoError(t, err)
		if p.CutoverNanos() == testPlacementCutover {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Instance is now in the placement.
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	require.NoError(t, agg.AddUntimed(testUntimedMetric, testStagedMetadatas))
	require.Equal(t, testShardSetID, agg.shardSetID)
	require.True(t, agg.shardSetOpen)
	require.Equal(t, aggregatorOpen, agg.state)
	require.Equal(t, []uint32{0, 1, 2, 3}, agg.shardIDs)
	for i := 0; i < testNumShards; i++ {
		require.NotNil(t, agg.shards[i])
	}
	require.Equal(t, int64(testPlacementCutover), agg.currPlacement.CutoverNanos())

	// Update placement to remove the instance and wait for it to propagate.
	_, err = store.Set(testPlacementKey, placements[2])
	require.NoError(t, err)
	for {
		p, err := agg.placementManager.Placement()
		require.NoError(t, err)
		if p.CutoverNanos() == testPlacementCutover+1000 {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Instance is now removed from the placement.
	require.Error(t, agg.AddUntimed(testUntimedMetric, testStagedMetadatas))
	require.Equal(t, uint32(0), agg.shardSetID)
	require.False(t, agg.shardSetOpen)
	require.Equal(t, 0, len(agg.shardIDs))
	require.Nil(t, agg.shards)
	require.Equal(t, int64(testPlacementCutover+1000), agg.currPlacement.CutoverNanos())
	require.Equal(t, aggregatorOpen, agg.state)
}

func TestAggregatorAddUntimedInvalidMetricType(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	err := agg.AddUntimed(testInvalidMetric, testStagedMetadatas)
	require.Equal(t, errInvalidMetricType, err)
}

func TestAggregatorAddUntimedShardNotOwned(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	err := agg.AddUntimed(testUntimedMetric, testStagedMetadatas)
	require.Equal(t, errShardNotOwned, err)
}

func TestAggregatorAddUntimedNotResponsibleForShard(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return testNumShards }
	require.Equal(t, errShardNotOwned, agg.AddUntimed(testUntimedMetric, testStagedMetadatas))
}

func TestAggregatorAddUntimedSuccessNoPlacementUpdate(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	err := agg.AddUntimed(testUntimedMetric, testStagedMetadatas)
	require.NoError(t, err)
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
}

func TestAggregatorAddUntimedSuccessWithPlacementUpdate(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, store := testAggregator(t, ctrl)
	now := time.Unix(0, 12345)
	nowFn := func() time.Time { return now }
	agg.opts = agg.opts.
		SetClockOptions(agg.opts.ClockOptions().SetNowFn(nowFn)).
		SetBufferDurationBeforeShardCutover(time.Duration(500)).
		SetBufferDurationAfterShardCutoff(time.Duration(1000))
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	require.NoError(t, agg.Open())
	require.Equal(t, int64(testPlacementCutover), agg.currPlacement.CutoverNanos())

	existingShard := agg.shards[3]

	newShardAssignment := []shard.Shard{
		shard.NewShard(0).SetState(shard.Initializing).SetCutoverNanos(5000).SetCutoffNanos(20000),
		shard.NewShard(1).SetState(shard.Initializing).SetCutoverNanos(5500).SetCutoffNanos(25000),
		shard.NewShard(2).SetState(shard.Initializing).SetCutoverNanos(6000).SetCutoffNanos(30000),
		shard.NewShard(4).SetState(shard.Initializing).SetCutoverNanos(6500).SetCutoffNanos(math.MaxInt64),
	}
	newPlacementCutoverNanos := int64(testPlacementCutover - 1)
	newStagedPlacementProto := testStagedPlacementProtoWithCustomShards(t, testInstanceID, testShardSetID, newShardAssignment, newPlacementCutoverNanos)
	_, err := store.Set(testPlacementKey, newStagedPlacementProto)
	require.NoError(t, err)

	// Wait for the placement to be updated.
	for {
		placement, err := agg.placementManager.Placement()
		require.NoError(t, err)
		if placement.CutoverNanos() == newPlacementCutoverNanos {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}

	err = agg.AddUntimed(testUntimedMetric, testStagedMetadatas)
	require.NoError(t, err)
	require.Equal(t, 5, len(agg.shards))

	expectedShards := []struct {
		isNil         bool
		earliestNanos int64
		latestNanos   int64
	}{
		{isNil: false, earliestNanos: 4500, latestNanos: 21000},
		{isNil: false, earliestNanos: 5000, latestNanos: 26000},
		{isNil: false, earliestNanos: 5500, latestNanos: 31000},
		{isNil: true},
		{isNil: false, earliestNanos: 6000, latestNanos: math.MaxInt64},
	}
	for i, expected := range expectedShards {
		if expected.isNil {
			require.Nil(t, agg.shards[i])
		} else {
			require.NotNil(t, agg.shards[i])
			require.Equal(t, expected.earliestNanos, agg.shards[i].earliestWritableNanos)
			require.Equal(t, expected.latestNanos, agg.shards[i].latestWriteableNanos)
		}
	}
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
	require.Equal(t, newPlacementCutoverNanos, agg.currPlacement.CutoverNanos())
	for {
		existingShard.RLock()
		closed := existingShard.closed
		existingShard.RUnlock()

		if closed {
			break
		}

		time.Sleep(100 * time.Millisecond)
	}
}

func TestAggregatorAddUntimedWithShardRedirect(t *testing.T) {
	testAddWithShardRedirect(t, func(agg *aggregator) error {
		return agg.AddUntimed(testUntimedMetric, testStagedMetadatas)
	})
}

func TestAggregatorAddUntimedWithShardRedirectToNotOwned(t *testing.T) {
	testAddWithShardRedirectToNotOwned(t, func(agg *aggregator) error {
		return agg.AddUntimed(testUntimedMetric, testStagedMetadatas)
	})
}

func TestAggregatorAddTimedShardNotOwned(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	err := agg.AddTimed(testTimedMetric, testTimedMetadata)
	require.Equal(t, errShardNotOwned, err)
}

func TestAggregatorAddTimedNotResponsibleForShard(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return testNumShards }
	require.Equal(t, errShardNotOwned, agg.AddTimed(testTimedMetric, testTimedMetadata))
}

func TestAggregatorAddTimedSuccessNoPlacementUpdate(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	now := time.Unix(0, 12345)
	nowFn := func() time.Time { return now }
	agg.nowFn = nowFn

	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	err := agg.AddTimed(testTimedMetric, testTimedMetadata)
	require.NoError(t, err)
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
}

func TestAggregatorAddTimedSuccessWithPlacementUpdate(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, store := testAggregator(t, ctrl)
	now := time.Unix(0, 12345)
	nowFn := func() time.Time { return now }
	agg.opts = agg.opts.
		SetClockOptions(agg.opts.ClockOptions().SetNowFn(nowFn)).
		SetBufferDurationBeforeShardCutover(time.Duration(500)).
		SetBufferDurationAfterShardCutoff(time.Duration(1000))
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	require.NoError(t, agg.Open())
	require.Equal(t, int64(testPlacementCutover), agg.currPlacement.CutoverNanos())

	existingShard := agg.shards[3]

	newShardAssignment := []shard.Shard{
		shard.NewShard(0).SetState(shard.Initializing).SetCutoverNanos(5000).SetCutoffNanos(20000),
		shard.NewShard(1).SetState(shard.Initializing).SetCutoverNanos(5500).SetCutoffNanos(25000),
		shard.NewShard(2).SetState(shard.Initializing).SetCutoverNanos(6000).SetCutoffNanos(30000),
		shard.NewShard(4).SetState(shard.Initializing).SetCutoverNanos(6500).SetCutoffNanos(math.MaxInt64),
	}
	newPlacementCutoverNanos := int64(testPlacementCutover - 1)
	newStagedPlacementProto := testStagedPlacementProtoWithCustomShards(t, testInstanceID, testShardSetID, newShardAssignment, newPlacementCutoverNanos)
	_, err := store.Set(testPlacementKey, newStagedPlacementProto)
	require.NoError(t, err)

	// Wait for the placement to be updated.
	for {
		placement, err := agg.placementManager.Placement()
		require.NoError(t, err)
		if placement.CutoverNanos() == newPlacementCutoverNanos {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}

	err = agg.AddTimed(testTimedMetric, testTimedMetadata)
	require.NoError(t, err)
	require.Equal(t, 5, len(agg.shards))

	expectedShards := []struct {
		isNil         bool
		earliestNanos int64
		latestNanos   int64
	}{
		{isNil: false, earliestNanos: 4500, latestNanos: 21000},
		{isNil: false, earliestNanos: 5000, latestNanos: 26000},
		{isNil: false, earliestNanos: 5500, latestNanos: 31000},
		{isNil: true},
		{isNil: false, earliestNanos: 6000, latestNanos: math.MaxInt64},
	}
	for i, expected := range expectedShards {
		if expected.isNil {
			require.Nil(t, agg.shards[i])
		} else {
			require.NotNil(t, agg.shards[i])
			require.Equal(t, expected.earliestNanos, agg.shards[i].earliestWritableNanos)
			require.Equal(t, expected.latestNanos, agg.shards[i].latestWriteableNanos)
		}
	}
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
	require.Equal(t, newPlacementCutoverNanos, agg.currPlacement.CutoverNanos())

	for {
		existingShard.RLock()
		closed := existingShard.closed
		existingShard.RUnlock()
		if closed {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func TestAggregatorAddTimedWithShardRedirect(t *testing.T) {
	testAddWithShardRedirect(t, func(agg *aggregator) error {
		return agg.AddTimed(testTimedMetric, testTimedMetadata)
	})
}

func TestAggregatorAddTimedWithShardRedirectToNotOwned(t *testing.T) {
	testAddWithShardRedirectToNotOwned(t, func(agg *aggregator) error {
		return agg.AddTimed(testTimedMetric, testTimedMetadata)
	})
}

func TestAggregatorAddForwardedShardNotOwned(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	err := agg.AddForwarded(testForwardedMetric, testForwardMetadata)
	require.Equal(t, errShardNotOwned, err)
}

func TestAggregatorAddForwardedNotResponsibleForShard(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return testNumShards }
	require.Equal(t, errShardNotOwned, agg.AddForwarded(testForwardedMetric, testForwardMetadata))
}

func TestAggregatorAddForwardedSuccessNoPlacementUpdate(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	now := time.Unix(0, 12345)
	nowFn := func() time.Time { return now }
	agg.nowFn = nowFn

	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	err := agg.AddForwarded(testForwardedMetric, testForwardMetadata)
	require.NoError(t, err)
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
}

func TestAggregatorAddForwardedSuccessWithPlacementUpdate(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, store := testAggregator(t, ctrl)
	now := time.Unix(0, 12345)
	nowFn := func() time.Time { return now }
	agg.opts = agg.opts.
		SetClockOptions(agg.opts.ClockOptions().SetNowFn(nowFn)).
		SetBufferDurationBeforeShardCutover(time.Duration(500)).
		SetBufferDurationAfterShardCutoff(time.Duration(1000))
	agg.shardFn = func([]byte, uint32) uint32 { return 1 }
	require.NoError(t, agg.Open())
	require.Equal(t, int64(testPlacementCutover), agg.currPlacement.CutoverNanos())

	existingShard := agg.shards[3]

	newShardAssignment := []shard.Shard{
		shard.NewShard(0).SetState(shard.Initializing).SetCutoverNanos(5000).SetCutoffNanos(20000),
		shard.NewShard(1).SetState(shard.Initializing).SetCutoverNanos(5500).SetCutoffNanos(25000),
		shard.NewShard(2).SetState(shard.Initializing).SetCutoverNanos(6000).SetCutoffNanos(30000),
		shard.NewShard(4).SetState(shard.Initializing).SetCutoverNanos(6500).SetCutoffNanos(math.MaxInt64),
	}
	newPlacementCutoverNanos := int64(testPlacementCutover - 1)
	newStagedPlacementProto := testStagedPlacementProtoWithCustomShards(t, testInstanceID, testShardSetID, newShardAssignment, newPlacementCutoverNanos)
	_, err := store.Set(testPlacementKey, newStagedPlacementProto)
	require.NoError(t, err)

	// Wait for the placement to be updated.
	for {
		placement, err := agg.placementManager.Placement()
		require.NoError(t, err)
		if placement.CutoverNanos() == newPlacementCutoverNanos {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}

	err = agg.AddForwarded(testForwardedMetric, testForwardMetadata)
	require.NoError(t, err)
	require.Equal(t, 5, len(agg.shards))

	expectedShards := []struct {
		isNil         bool
		earliestNanos int64
		latestNanos   int64
	}{
		{isNil: false, earliestNanos: 4500, latestNanos: 21000},
		{isNil: false, earliestNanos: 5000, latestNanos: 26000},
		{isNil: false, earliestNanos: 5500, latestNanos: 31000},
		{isNil: true},
		{isNil: false, earliestNanos: 6000, latestNanos: math.MaxInt64},
	}
	for i, expected := range expectedShards {
		if expected.isNil {
			require.Nil(t, agg.shards[i])
		} else {
			require.NotNil(t, agg.shards[i])
			require.Equal(t, expected.earliestNanos, agg.shards[i].earliestWritableNanos)
			require.Equal(t, expected.latestNanos, agg.shards[i].latestWriteableNanos)
		}
	}
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
	require.Equal(t, newPlacementCutoverNanos, agg.currPlacement.CutoverNanos())

	for {
		existingShard.RLock()
		closed := existingShard.closed
		existingShard.RUnlock()
		if closed {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func TestAggregatorAddForwardedWithShardRedirect(t *testing.T) {
	testAddWithShardRedirect(t, func(agg *aggregator) error {
		return agg.AddForwarded(testForwardedMetric, testForwardMetadata)
	})
}

func TestAggregatorAddForwardedWithShardRedirectToNotOwned(t *testing.T) {
	testAddWithShardRedirectToNotOwned(t, func(agg *aggregator) error {
		return agg.AddForwarded(testForwardedMetric, testForwardMetadata)
	})
}

func TestAggregatorResignError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	errTestResign := errors.New("test resign")
	electionMgr := NewMockElectionManager(ctrl)
	electionMgr.EXPECT().Resign(gomock.Any()).Return(errTestResign)
	agg, _ := testAggregator(t, ctrl)
	agg.electionManager = electionMgr

	require.Equal(t, errTestResign, agg.Resign())
}

func TestAggregatorResignSuccess(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	electionMgr := NewMockElectionManager(ctrl)
	electionMgr.EXPECT().Resign(gomock.Any()).Return(nil)
	agg, _ := testAggregator(t, ctrl)
	agg.electionManager = electionMgr
	require.NoError(t, agg.Resign())
}

func TestAggregatorAddPassthroughNotOpen(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	err := agg.AddPassthrough(testPassthroughMetric, testPassthroughStroagePolicy)
	require.Equal(t, errAggregatorNotOpenOrClosed, err)
}

func TestAggregatorAddPassthroughSuccess(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	err := agg.AddPassthrough(testPassthroughMetric, testPassthroughStroagePolicy)
	require.NoError(t, err)
}

func TestAggregatorStatus(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	flushStatus := FlushStatus{
		ElectionState: LeaderState,
		CanLead:       true,
	}
	flushManager := NewMockFlushManager(ctrl)
	flushManager.EXPECT().Status().Return(flushStatus)
	agg, _ := testAggregator(t, ctrl)
	agg.flushManager = flushManager
	require.Equal(t, RuntimeStatus{FlushStatus: flushStatus}, agg.Status())
}

func TestAggregatorCloseAlreadyClosed(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.Equal(t, errAggregatorNotOpenOrClosed, agg.Close())
}

func TestAggregatorCloseSuccess(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	require.NoError(t, agg.Open())
	require.NoError(t, agg.Close())
	require.Equal(t, aggregatorClosed, agg.state)
}

func TestAggregatorTick(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	flushTimesManager := NewMockFlushTimesManager(ctrl)
	flushTimesManager.EXPECT().Reset().Return(nil).AnyTimes()
	flushTimesManager.EXPECT().Open(gomock.Any()).Return(nil).AnyTimes()
	flushTimesManager.EXPECT().Get().Return(nil, nil).AnyTimes()
	flushTimesManager.EXPECT().Close().Return(nil).AnyTimes()

	agg, _ := testAggregator(t, ctrl)
	agg.flushTimesManager = flushTimesManager
	require.NoError(t, agg.Open())

	// Forcing a tick.
	agg.tickInternal()

	require.NoError(t, agg.Close())
}

func TestAggregatorShardSetNotOpenNilInstance(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	agg.shardSetOpen = false
	agg.Lock()
	defer agg.Unlock()
	require.NoError(t, agg.updateShardSetIDWithLock(nil))
}

func TestAggregatorShardSetNotOpenValidInstance(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var (
		flushTimesManagerOpenID *uint32
		electionManagerOpenID   *uint32
		testInstance            = placement.NewInstance().SetShardSetID(testShardSetID)
	)
	electionMgr := NewMockElectionManager(ctrl)
	electionMgr.EXPECT().
		Open(testShardSetID).
		DoAndReturn(func(shardSetID uint32) error {
			electionManagerOpenID = &shardSetID
			return nil
		})
	flushTimesManager := NewMockFlushTimesManager(ctrl)
	flushTimesManager.EXPECT().
		Open(testShardSetID).
		DoAndReturn(func(shardSetID uint32) error {
			flushTimesManagerOpenID = &shardSetID
			return nil
		})
	agg, _ := testAggregator(t, ctrl)
	agg.shardSetOpen = false
	agg.flushTimesManager = flushTimesManager
	agg.electionManager = electionMgr

	agg.Lock()
	defer agg.Unlock()
	require.NoError(t, agg.updateShardSetIDWithLock(testInstance))
	require.Equal(t, testShardSetID, agg.shardSetID)
	require.True(t, agg.shardSetOpen)
	require.Equal(t, testShardSetID, *flushTimesManagerOpenID)
	require.Equal(t, testShardSetID, *electionManagerOpenID)
}

func TestAggregatorShardSetOpenShardSetIDUnchanged(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	testInstance := placement.NewInstance().SetShardSetID(testShardSetID)
	agg, _ := testAggregator(t, ctrl)
	agg.shardSetOpen = true
	agg.shardSetID = testShardSetID
	agg.Lock()
	defer agg.Unlock()
	require.NoError(t, agg.updateShardSetIDWithLock(testInstance))
	require.Equal(t, testShardSetID, agg.shardSetID)
	require.True(t, agg.shardSetOpen)
}

func TestAggregatorShardSetOpenNilInstance(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	agg.shardSetOpen = true
	agg.shardSetID = testShardSetID
	agg.Lock()
	defer agg.Unlock()
	require.NoError(t, agg.updateShardSetIDWithLock(nil))
	require.Equal(t, uint32(0), agg.shardSetID)
	require.False(t, agg.shardSetOpen)
}

func TestAggregatorShardSetOpenValidInstance(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var (
		flushTimesManagerOpenID *uint32
		electionManagerOpenID   *uint32
		newShardSetID           = uint32(2)
		testInstance            = placement.NewInstance().SetShardSetID(newShardSetID)
	)
	electionMgr := NewMockElectionManager(ctrl)
	electionMgr.EXPECT().Reset().Return(nil)
	electionMgr.EXPECT().
		Open(newShardSetID).
		DoAndReturn(func(shardSetID uint32) error {
			electionManagerOpenID = &shardSetID
			return nil
		})
	electionMgr.EXPECT().Close().Return(nil)
	flushTimesManager := NewMockFlushTimesManager(ctrl)
	flushTimesManager.EXPECT().Reset().Return(nil)
	flushTimesManager.EXPECT().
		Open(newShardSetID).
		DoAndReturn(func(shardSetID uint32) error {
			flushTimesManagerOpenID = &shardSetID
			return nil
		})
	flushTimesManager.EXPECT().Close().Return(nil)

	agg, _ := testAggregator(t, ctrl)
	agg.shardSetOpen = true
	agg.shardSetID = testShardSetID
	agg.flushTimesManager = flushTimesManager
	agg.electionManager = electionMgr

	agg.Lock()
	defer agg.Unlock()
	require.NoError(t, agg.updateShardSetIDWithLock(testInstance))
	require.Equal(t, newShardSetID, agg.shardSetID)
	require.True(t, agg.shardSetOpen)
	require.Equal(t, newShardSetID, *flushTimesManagerOpenID)
	require.Equal(t, newShardSetID, *electionManagerOpenID)
}

func TestAggregatorOwnedShards(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregator(t, ctrl)
	now := time.Unix(0, 12345)
	nowFn := func() time.Time { return now }
	agg.nowFn = nowFn
	agg.opts = agg.opts.SetClockOptions(agg.opts.ClockOptions().SetNowFn(nowFn))

	inputs := []struct {
		cutoffNanos         int64
		latestWritableNanos int64
		flushedNanos        int64
	}{
		{
			cutoffNanos:         math.MaxInt64,
			latestWritableNanos: math.MaxInt64,
			flushedNanos:        12340,
		},
		{
			cutoffNanos:         1234,
			latestWritableNanos: 5678,
			flushedNanos:        1234,
		},
		{
			cutoffNanos:         56789,
			latestWritableNanos: 67890,
			flushedNanos:        12345,
		},
		{
			cutoffNanos:         8901,
			latestWritableNanos: 34567,
			flushedNanos:        7200,
		},
	}
	flushTimes := make(map[uint32]*schema.ShardFlushTimes, len(inputs))
	agg.shardIDs = make([]uint32, 0, len(inputs))
	agg.shards = make([]*aggregatorShard, 0, len(inputs))
	for i, input := range inputs {
		shardID := uint32(i)
		shard := newAggregatorShard(shardID, agg.opts)
		shard.cutoffNanos = input.cutoffNanos
		shard.latestWriteableNanos = input.latestWritableNanos
		flushTimes[shardID] = &schema.ShardFlushTimes{
			StandardByResolution: map[int64]int64{
				int64(time.Second): input.flushedNanos,
			},
		}
		agg.shardIDs = append(agg.shardIDs, shardID)
		agg.shards = append(agg.shards, shard)
	}
	flushTimesManager := NewMockFlushTimesManager(ctrl)
	flushTimesManager.EXPECT().
		Get().
		Return(&schema.ShardSetFlushTimes{ByShard: flushTimes}, nil)
	agg.flushTimesManager = flushTimesManager

	expectedOwned := []*aggregatorShard{agg.shards[0], agg.shards[3], agg.shards[2]}
	expectedToClose := []*aggregatorShard{agg.shards[1]}
	owned, toClose := agg.ownedShards()
	require.Equal(t, expectedOwned, owned)
	require.Equal(t, expectedToClose, toClose)

	expectedShardIDs := []uint32{0, 3, 2}
	shardIDs := make([]uint32, len(agg.shardIDs))
	copy(shardIDs, agg.shardIDs)
	sort.Sort(uint32Ascending(shardIDs))
	require.Equal(t, expectedShardIDs, agg.shardIDs)

	for i := 0; i < 4; i++ {
		if i == 1 {
			require.Nil(t, agg.shards[i])
		} else {
			require.NotNil(t, agg.shards[i])
		}
	}
}

func TestAggregatorAddUntimedMetrics(t *testing.T) {
	s := tally.NewTestScope("testScope", nil)
	m := newAggregatorAddUntimedMetrics(s, instrument.TimerOptions{StandardSampleRate: 0.001})
	m.ReportSuccess()
	m.SuccessLatencyStopwatch().Stop()
	log := zap.NewNop()
	for _, state := range []ElectionState{LeaderState, FollowerState} {
		m.ReportError(errInvalidMetricType, state, log)
		m.ReportError(errShardNotOwned, state, log)
		m.ReportError(errAggregatorShardNotWriteable, state, log)
		m.ReportError(errWriteNewMetricRateLimitExceeded, state, log)
		m.ReportError(errWriteValueRateLimitExceeded, state, log)
		m.ReportError(xerrors.NewRenamedError(errArrivedTooLate, errors.New("errorrr")), state, log)
		m.ReportError(errTooFarInTheFuture, state, log)
		m.ReportError(errTooFarInThePast, state, log)
		m.ReportError(errAggregationClosed, state, log)
		m.ReportError(errors.New("foo"), state, log)
	}

	snapshot := s.Snapshot()
	counters, timers, gauges := snapshot.Counters(), snapshot.Timers(), snapshot.Gauges()

	// Validate we count successes and errors correctly.
	expectedCounters := []string{
		"testScope.success+",
		"testScope.errors+reason=invalid-metric-types,role=leader",
		"testScope.errors+reason=invalid-metric-types,role=non-leader",
		"testScope.errors+reason=shard-not-owned,role=leader",
		"testScope.errors+reason=shard-not-owned,role=non-leader",
		"testScope.errors+reason=shard-not-writeable,role=leader",
		"testScope.errors+reason=shard-not-writeable,role=non-leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=too-far-in-the-future,role=leader",
		"testScope.errors+reason=too-far-in-the-future,role=non-leader",
		"testScope.errors+reason=too-far-in-the-past,role=leader",
		"testScope.errors+reason=too-far-in-the-past,role=non-leader",
		"testScope.errors+reason=arrived-too-late,role=leader",
		"testScope.errors+reason=arrived-too-late,role=non-leader",
		"testScope.errors+reason=aggregation-closed,role=leader",
		"testScope.errors+reason=aggregation-closed,role=non-leader",
		"testScope.errors+reason=not-categorized,role=leader",
		"testScope.errors+reason=not-categorized,role=non-leader",
	}
	require.Len(t, counters, len(expectedCounters))
	for _, id := range expectedCounters {
		c, exists := counters[id]
		require.True(t, exists)
		require.Equal(t, int64(1), c.Value())
	}

	// Validate we record times correctly.
	require.Equal(t, 1, len(timers))

	for _, id := range []string{
		"testScope.success-latency+",
	} {
		_, exists := timers[id]
		require.True(t, exists)
	}

	// Validate we do not have any gauges.
	require.Equal(t, 0, len(gauges))
}

func TestAggregatorAddTimedMetrics(t *testing.T) {
	s := tally.NewTestScope("testScope", nil)
	m := newAggregatorAddTimedMetrics(s, instrument.TimerOptions{})
	m.ReportSuccess()
	m.SuccessLatencyStopwatch().Stop()
	log := zap.NewNop()
	for _, state := range []ElectionState{LeaderState, FollowerState} {
		m.ReportError(errShardNotOwned, state, log)
		m.ReportError(errAggregatorShardNotWriteable, state, log)
		m.ReportError(errWriteNewMetricRateLimitExceeded, state, log)
		m.ReportError(errWriteValueRateLimitExceeded, state, log)
		m.ReportError(errTooFarInTheFuture, state, log)
		m.ReportError(errTooFarInThePast, state, log)
		m.ReportError(xerrors.NewRenamedError(errArrivedTooLate, errors.New("errorrr")), state, log)
		m.ReportError(errAggregationClosed, state, log)
		m.ReportError(errors.New("foo"), state, log)
	}

	snapshot := s.Snapshot()
	counters, timers, gauges := snapshot.Counters(), snapshot.Timers(), snapshot.Gauges()

	// Validate we count successes and errors correctly.
	expectedCounters := []string{
		"testScope.success+",
		"testScope.errors+reason=shard-not-owned,role=leader",
		"testScope.errors+reason=shard-not-owned,role=non-leader",
		"testScope.errors+reason=shard-not-writeable,role=leader",
		"testScope.errors+reason=shard-not-writeable,role=non-leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=too-far-in-the-future,role=leader",
		"testScope.errors+reason=too-far-in-the-future,role=non-leader",
		"testScope.errors+reason=too-far-in-the-past,role=leader",
		"testScope.errors+reason=too-far-in-the-past,role=non-leader",
		"testScope.errors+reason=arrived-too-late,role=leader",
		"testScope.errors+reason=arrived-too-late,role=non-leader",
		"testScope.errors+reason=aggregation-closed,role=leader",
		"testScope.errors+reason=aggregation-closed,role=non-leader",
		"testScope.errors+reason=not-categorized,role=leader",
		"testScope.errors+reason=not-categorized,role=non-leader",
	}
	require.Len(t, counters, len(expectedCounters))
	for _, id := range expectedCounters {
		c, exists := counters[id]
		require.True(t, exists)
		require.Equal(t, int64(1), c.Value())
	}

	// Validate we record times correctly.
	require.Equal(t, 1, len(timers))

	for _, id := range []string{
		"testScope.success-latency+",
	} {
		_, exists := timers[id]
		require.True(t, exists)
	}

	// Validate we do not have any gauges.
	require.Equal(t, 0, len(gauges))
}

func TestAggregatorAddPassthroughMetrics(t *testing.T) {
	s := tally.NewTestScope("testScope", nil)
	m := newAggregatorAddPassthroughMetrics(s, instrument.TimerOptions{})
	m.ReportSuccess()
	m.SuccessLatencyStopwatch().Stop()
	m.ReportFollowerNoop()
	log := zap.NewNop()
	for _, state := range []ElectionState{LeaderState, FollowerState} {
		m.ReportError(errShardNotOwned, state, log)
		m.ReportError(errAggregatorShardNotWriteable, state, log)
		m.ReportError(errWriteNewMetricRateLimitExceeded, state, log)
		m.ReportError(errWriteValueRateLimitExceeded, state, log)
		m.ReportError(xerrors.NewRenamedError(errArrivedTooLate, errors.New("errorrr")), state, log)
		m.ReportError(errAggregationClosed, state, log)
		m.ReportError(errors.New("foo"), state, log)
	}

	snapshot := s.Snapshot()
	counters, timers, gauges := snapshot.Counters(), snapshot.Timers(), snapshot.Gauges()

	// Validate we count successes and errors correctly.
	expectedCounters := []string{
		"testScope.success+",
		"testScope.follower-noop+",
		"testScope.errors+reason=shard-not-owned,role=leader",
		"testScope.errors+reason=shard-not-owned,role=non-leader",
		"testScope.errors+reason=shard-not-writeable,role=leader",
		"testScope.errors+reason=shard-not-writeable,role=non-leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=arrived-too-late,role=leader",
		"testScope.errors+reason=arrived-too-late,role=non-leader",
		"testScope.errors+reason=aggregation-closed,role=leader",
		"testScope.errors+reason=aggregation-closed,role=non-leader",
		"testScope.errors+reason=not-categorized,role=leader",
		"testScope.errors+reason=not-categorized,role=non-leader",
	}
	require.Len(t, counters, len(expectedCounters))
	for _, id := range expectedCounters {
		c, exists := counters[id]
		require.True(t, exists)
		require.Equal(t, int64(1), c.Value())
	}

	// Validate we record times correctly.
	require.Equal(t, 1, len(timers))

	for _, id := range []string{
		"testScope.success-latency+",
	} {
		_, exists := timers[id]
		require.True(t, exists)
	}

	// Validate we do not have any gauges.
	require.Equal(t, 0, len(gauges))
}

func TestAggregatorAddForwardedMetrics(t *testing.T) {
	s := tally.NewTestScope("testScope", nil)
	delayFunc := func(resolution time.Duration, numForwardedTimes int) time.Duration {
		return time.Second
	}
	m := newAggregatorAddForwardedMetrics(s, instrument.TimerOptions{}, delayFunc)
	m.ReportSuccess()
	m.SuccessLatencyStopwatch().Stop()
	m.ReportForwardingLatency(time.Second, 1, 100*time.Millisecond)
	m.ReportForwardingLatency(time.Second, 2, 100*time.Millisecond)
	log := zap.NewNop()
	for _, state := range []ElectionState{LeaderState, FollowerState} {
		m.ReportError(errShardNotOwned, state, log)
		m.ReportError(errAggregatorShardNotWriteable, state, log)
		m.ReportError(errWriteNewMetricRateLimitExceeded, state, log)
		m.ReportError(errWriteValueRateLimitExceeded, state, log)
		m.ReportError(xerrors.NewRenamedError(errArrivedTooLate, errors.New("errorrr")), state, log)
		m.ReportError(errAggregationClosed, state, log)
		m.ReportError(errors.New("foo"), state, log)
	}

	var (
		snapshot   = s.Snapshot()
		counters   = snapshot.Counters()
		timers     = snapshot.Timers()
		gauges     = snapshot.Gauges()
		histograms = snapshot.Histograms()
	)

	// Validate we count successes and errors correctly.
	expectedCounters := []string{
		"testScope.success+",
		"testScope.errors+reason=shard-not-owned,role=leader",
		"testScope.errors+reason=shard-not-owned,role=non-leader",
		"testScope.errors+reason=shard-not-writeable,role=leader",
		"testScope.errors+reason=shard-not-writeable,role=non-leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=value-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=leader",
		"testScope.errors+reason=new-metric-rate-limit-exceeded,role=non-leader",
		"testScope.errors+reason=arrived-too-late,role=leader",
		"testScope.errors+reason=arrived-too-late,role=non-leader",
		"testScope.errors+reason=aggregation-closed,role=leader",
		"testScope.errors+reason=aggregation-closed,role=non-leader",
		"testScope.errors+reason=not-categorized,role=leader",
		"testScope.errors+reason=not-categorized,role=non-leader",
	}
	require.Len(t, counters, len(expectedCounters))
	for _, id := range expectedCounters {
		c, exists := counters[id]
		require.True(t, exists)
		require.Equal(t, int64(1), c.Value())
	}

	// Validate we record times correctly.
	require.Equal(t, 1, len(timers))

	for _, id := range []string{
		"testScope.success-latency+",
	} {
		_, exists := timers[id]
		require.True(t, exists)
	}

	// Validate we do not have any gauges.
	require.Equal(t, 0, len(gauges))

	// Validate we measure histograms correctly.
	expectedHistograms := []string{
		"testScope.forwarding-latency+bucket-version=2,num-forwarded-times=1,resolution=1s",
		"testScope.forwarding-latency+bucket-version=2,num-forwarded-times=2,resolution=1s",
	}
	require.Len(t, histograms, len(expectedHistograms))
	for _, id := range expectedHistograms {
		h, exists := histograms[id]
		require.True(t, exists, "missing series id %v", id)
		h.Durations()[100*time.Millisecond] = 1
	}
}

func TestAggregatorUpdateShardsIgnoreCutoffCutover(t *testing.T) {
	testAggregatorUpdateShards(t, true)
}

func TestAggregatorUpdateShardsRespectCutoffCutover(t *testing.T) {
	testAggregatorUpdateShards(t, false)
}

func testAggregatorUpdateShards(t *testing.T, ignoreCutoffCutover bool) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var (
		now     = time.Now().Truncate(time.Hour)
		cutover = now.Add(-time.Hour)
		cutoff  = now.Add(time.Hour)
	)

	shard := shard.NewShard(0).
		SetState(shard.Initializing).
		SetCutoverNanos(cutover.UnixNano()).
		SetCutoffNanos(cutoff.UnixNano())

	placement, shards := testPlacementWithCustomShards(testInstanceID, 1, shard)

	agg, _ := testAggregator(t, ctrl)
	opts := agg.opts

	agg.opts = opts.SetWritesIgnoreCutoffCutover(ignoreCutoffCutover)
	agg.updateShardsWithLock(placement, shards)

	expectedEarliest, expectedLatest := int64(0), int64(math.MaxInt64)
	if !ignoreCutoffCutover {
		expectedEarliest = cutover.Add(-opts.BufferDurationBeforeShardCutover()).UnixNano()
		expectedLatest = cutoff.Add(opts.BufferDurationAfterShardCutoff()).UnixNano()
	}

	aggShard := agg.shards[0]
	assert.Equal(t, expectedEarliest, aggShard.earliestWritableNanos)
	assert.Equal(t, expectedLatest, aggShard.latestWriteableNanos)
}

func testAddWithShardRedirect(t *testing.T, addFn func(*aggregator) error) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregatorWithShardRedirect(t, ctrl, 3, 1)
	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return 3 }
	err := addFn(agg)
	require.NoError(t, err)
	require.Equal(t, 1, len(agg.shards[1].metricMap.entries))
}

func testAddWithShardRedirectToNotOwned(t *testing.T, addFn func(*aggregator) error) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	agg, _ := testAggregatorWithShardRedirect(t, ctrl, 3, 100)
	require.NoError(t, agg.Open())
	agg.shardFn = func([]byte, uint32) uint32 { return 3 }
	err := addFn(agg)
	require.Equal(t, errShardNotOwned, err)
}

func testAggregator(t *testing.T, ctrl *gomock.Controller) (*aggregator, kv.Store) {
	proto := testStagedPlacementProtoWithNumShards(t, testInstanceID, testShardSetID, testNumShards)
	return testAggregatorWithCustomPlacements(t, ctrl, proto)
}

func testAggregatorWithShardRedirect(
	t *testing.T,
	ctrl *gomock.Controller,
	fromShard, toShard uint32,
) (*aggregator, kv.Store) {
	proto := testStagedPlacementProtoWithNumShards(t, testInstanceID, testShardSetID, testNumShards)
	shards := proto.Snapshots[0].Instances[testInstanceID].Shards
	shards[fromShard].RedirectToShardId = &types.UInt32Value{Value: toShard}

	return testAggregatorWithCustomPlacements(t, ctrl, proto)
}

func testAggregatorWithCustomPlacements(
	t *testing.T,
	ctrl *gomock.Controller,
	proto *placementpb.PlacementSnapshots,
) (*aggregator, kv.Store) {
	watcherOpts, store := testWatcherOptsWithPlacementProto(t, testPlacementKey, proto)
	placementManagerOpts := NewPlacementManagerOptions().
		SetInstanceID(testInstanceID).
		SetWatcherOptions(watcherOpts)
	placementManager := NewPlacementManager(placementManagerOpts)
	opts := testOptions(ctrl).
		SetEntryCheckInterval(0).
		SetPlacementManager(placementManager)
	return NewAggregator(opts).(*aggregator), store
}

// nolint: unparam
func testWatcherOptsWithPlacementProto(
	t *testing.T,
	placementKey string,
	proto *placementpb.PlacementSnapshots,
) (placement.WatcherOptions, kv.Store) {
	t.Helper()
	store := mem.NewStore()
	_, err := store.SetIfNotExists(placementKey, proto)
	require.NoError(t, err)
	placementWatcherOpts := placement.NewWatcherOptions().
		SetStagedPlacementKey(placementKey).
		SetStagedPlacementStore(store)
	return placementWatcherOpts, store
}

// nolint: unparam
func testStagedPlacementProtoWithNumShards(
	t *testing.T,
	instanceID string,
	shardSetID uint32,
	numShards int,
) *placementpb.PlacementSnapshots {
	shardSet := make([]shard.Shard, numShards)
	for i := 0; i < numShards; i++ {
		shardSet[i] = shard.NewShard(uint32(i)).
			SetState(shard.Initializing).
			SetCutoverNanos(0).
			SetCutoffNanos(math.MaxInt64)
	}
	return testStagedPlacementProtoWithCustomShards(t, instanceID, shardSetID, shardSet, testPlacementCutover)
}

func testPlacementWithCustomShards(
	instanceID string,
	shardSetID uint32,
	customShards ...shard.Shard,
) (placement.Placement, shard.Shards) {
	shards := shard.NewShards(customShards)
	instance := placement.NewInstance().
		SetID(instanceID).
		SetShards(shards).
		SetShardSetID(shardSetID)
	placement := placement.NewPlacement().
		SetInstances([]placement.Instance{instance}).
		SetShards(shards.AllIDs())

	return placement, shards
}

func testStagedPlacementProtoWithCustomShards(
	t *testing.T,
	instanceID string,
	shardSetID uint32,
	shardSet []shard.Shard,
	placementCutoverNanos int64,
) *placementpb.PlacementSnapshots {
	testPlacement, _ := testPlacementWithCustomShards(instanceID, shardSetID, shardSet...)
	testPlacement.SetCutoverNanos(placementCutoverNanos)
	testStagedPlacement, err := placement.NewPlacementsFromLatest(testPlacement)
	require.NoError(t, err)
	stagedPlacementProto, err := testStagedPlacement.Proto()
	require.NoError(t, err)
	return stagedPlacementProto
}

func testOptions(ctrl *gomock.Controller) Options {
	placementManager := NewMockPlacementManager(ctrl)
	placementManager.EXPECT().Close().Return(nil).AnyTimes()

	flushTimesManager := NewMockFlushTimesManager(ctrl)
	flushTimesManager.EXPECT().Reset().Return(nil).AnyTimes()
	flushTimesManager.EXPECT().Open(gomock.Any()).Return(nil).AnyTimes()
	flushTimesManager.EXPECT().Close().Return(nil).AnyTimes()

	electionMgr := NewMockElectionManager(ctrl)
	electionMgr.EXPECT().Reset().Return(nil).AnyTimes()
	electionMgr.EXPECT().Open(gomock.Any()).Return(nil).AnyTimes()
	electionMgr.EXPECT().Close().Return(nil).AnyTimes()
	electionMgr.EXPECT().ElectionState().Return(LeaderState).AnyTimes()

	flushManager := NewMockFlushManager(ctrl)
	flushManager.EXPECT().Reset().Return(nil).AnyTimes()
	flushManager.EXPECT().Open().Return(nil).AnyTimes()
	flushManager.EXPECT().Register(gomock.Any()).Return(nil).AnyTimes()
	flushManager.EXPECT().Unregister(gomock.Any()).Return(nil).AnyTimes()
	flushManager.EXPECT().Close().Return(nil).AnyTimes()

	w := writer.NewMockWriter(ctrl)
	w.EXPECT().Write(gomock.Any()).Return(nil).AnyTimes()
	w.EXPECT().Flush().Return(nil).AnyTimes()
	w.EXPECT().Close().Return(nil).AnyTimes()

	h := handler.NewMockHandler(ctrl)
	h.EXPECT().NewWriter(gomock.Any()).Return(w, nil).AnyTimes()
	h.EXPECT().Close().AnyTimes()

	pw := writer.NewMockWriter(ctrl)
	pw.EXPECT().Write(gomock.Any()).Return(nil).AnyTimes()
	pw.EXPECT().Flush().Return(nil).AnyTimes()
	pw.EXPECT().Close().Return(nil).AnyTimes()

	cl := client.NewMockAdminClient(ctrl)
	cl.EXPECT().Flush().Return(nil).AnyTimes()
	cl.EXPECT().Close().AnyTimes()

	infiniteAllowedDelayFn := func(time.Duration, int) time.Duration {
		return math.MaxInt64
	}
	infiniteBufferForPastTimedMetricFn := func(time.Duration) time.Duration {
		return math.MaxInt64
	}
	return newTestOptions().
		SetPlacementManager(placementManager).
		SetFlushTimesManager(flushTimesManager).
		SetElectionManager(electionMgr).
		SetFlushManager(flushManager).
		SetFlushHandler(h).
		SetPassthroughWriter(pw).
		SetAdminClient(cl).
		SetMaxAllowedForwardingDelayFn(infiniteAllowedDelayFn).
		SetBufferForFutureTimedMetric(math.MaxInt64).
		SetBufferForPastTimedMetricFn(infiniteBufferForPastTimedMetricFn)
}

type uint32Ascending []uint32

func (a uint32Ascending) Len() int           { return len(a) }
func (a uint32Ascending) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a uint32Ascending) Less(i, j int) bool { return a[i] < a[j] }


<|start_filename|>src/query/api/v1/middleware/middleware_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package middleware

import (
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/m3db/m3/src/query/util/logging"
	"github.com/m3db/m3/src/x/instrument"

	"github.com/gorilla/mux"
	"github.com/opentracing/opentracing-go/mocktracer"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
)

func TestTracing(t *testing.T) {
	core, recorded := observer.New(zapcore.InfoLevel)
	iOpts := instrument.NewOptions().SetLogger(zap.New(core))
	mtr := mocktracer.New()
	r := mux.NewRouter()
	r.Use(Tracing(mtr, iOpts))
	r.HandleFunc(testRoute, func(w http.ResponseWriter, r *http.Request) {
		logging.WithContext(r.Context(), iOpts).Info("test")
	})

	req := httptest.NewRequest("GET", testRoute, nil)
	res := httptest.NewRecorder()
	r.ServeHTTP(res, req)

	spans := mtr.FinishedSpans()
	require.Len(t, spans, 1)
	require.Equal(t, fmt.Sprintf("GET %s", testRoute), spans[0].OperationName)
	require.Len(t, recorded.All(), 1)
	entry := recorded.All()[0]
	require.Equal(t, "test", entry.Message)
	fields := entry.ContextMap()
	require.Len(t, fields, 2)
	require.NotEqual(t, "", fields["trace_id"])
	require.NotEqual(t, "", fields["span_id"])
}

func TestCompression(t *testing.T) {
	router := mux.NewRouter()
	setupTestRouteRouter(router)

	router.Use(Compression())

	req := httptest.NewRequest("GET", testRoute, nil)
	req.Header.Add("Accept-Encoding", "gzip")
	res := httptest.NewRecorder()
	router.ServeHTTP(res, req)

	enc, found := res.Header()["Content-Encoding"]
	require.True(t, found)
	require.Equal(t, 1, len(enc))
	assert.Equal(t, "gzip", enc[0])

	cr, err := gzip.NewReader(res.Body)
	require.NoError(t, err)
	body, err := ioutil.ReadAll(cr)
	require.NoError(t, err)
	assert.Equal(t, "hello!", string(body))
}

func TestCors(t *testing.T) {
	router := mux.NewRouter()
	setupTestRouteRouter(router)

	router.Use(Cors())

	req := httptest.NewRequest("GET", testRoute, nil)
	res := httptest.NewRecorder()
	router.ServeHTTP(res, req)

	assert.Equal(t, "hello!", res.Body.String())
	assert.Equal(t, "*", res.Header().Get("Access-Control-Allow-Origin"))
}

const testRoute = "/foobar"

func setupTestRouteRouter(r *mux.Router) {
	r.HandleFunc(testRoute, func(writer http.ResponseWriter, r *http.Request) {
		writer.WriteHeader(http.StatusOK)
		_, _ = writer.Write([]byte("hello!"))
	})
}


<|start_filename|>src/integration/aggregator/aggregator.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// Package aggregator contains integration tests for aggregators.
package aggregator

import (
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/generated/thrift/rpc"
	"github.com/m3db/m3/src/integration/resources"
	"github.com/m3db/m3/src/query/generated/proto/prompb"
	"github.com/m3db/m3/src/x/headers"

	"github.com/stretchr/testify/require"
)

const (
	// TestAggregatorDBNodeConfig is the test config for the dbnode.
	TestAggregatorDBNodeConfig = `
db: {}
coordinator: {}
`

	// TestAggregatorCoordinatorConfig is the test config for the coordinator.
	TestAggregatorCoordinatorConfig = `
listenAddress: 0.0.0.0:7202
metrics:
  scope:
    prefix: "coordinator"
  prometheus:
    handlerPath: /metrics
    listenAddress: 0.0.0.0:7303
  sanitization: prometheus
  samplingRate: 1.0
  extended: none
carbon:
  ingester:
    listenAddress: "0.0.0.0:7204"
    rules:
      - pattern: .*
        aggregation:
          type: mean
        policies:
          - resolution: 5s
            retention: 6h
downsample:
  rules:
    rollupRules:
      - name: "requests per second by status code"
        filter: "__name__:http_requests app:* status_code:* endpoint:*"
        transforms:
          - transform:
              type: "PerSecond"
          - rollup:
              metricName: "http_requests_by_status_code"
              groupBy: ["app", "status_code", "endpoint"]
              aggregations: ["Sum"]
        storagePolicies:
          - resolution: 5s
            retention: 6h
  remoteAggregator:
    client:
      type: m3msg
      m3msg:
        producer:
          writer:
            topicName: aggregator_ingest
            topicServiceOverride:
              zone: embedded
              environment: default_env
            placement:
              isStaged: true
            placementServiceOverride:
              namespaces:
                placement: /placement
            connection:
              numConnections: 4
            messagePool:
              size: 16384
              watermark:
                low: 0.2
                high: 0.5
ingest:
  ingester:
    workerPoolSize: 10000
    opPool:
      size: 10000
    retry:
      maxRetries: 3
      jitter: true
    logSampleRate: 0.01
  m3msg:
    server:
      listenAddress: "0.0.0.0:7507"
      retry:
        maxBackoff: 10s
        jitter: true
storeMetricsType: true
`

	// TestAggregatorAggregatorConfig is the test config for the aggregators.
	TestAggregatorAggregatorConfig = `
`

	// defaultCarbonPort is the default port of coordinator to receive carbon metrics.
	defaultCarbonPort = 7204
)

var (
	errEmptyResult = errors.New("empty query result")
	errQueryResult = errors.New("wrong query result")
)

// RunTest contains the logic for running the aggregator test.
func RunTest(t *testing.T, m3 resources.M3Resources) {
	t.Run("test_aggregated_graphite_metric", func(t *testing.T) {
		testAggregatedGraphiteMetric(t, m3)
	})

	t.Run("test_rollup_rule", func(t *testing.T) {
		testRollupRule(t, m3)
	})

	t.Run("test_metric_type_survives_aggregation", func(t *testing.T) {
		testMetricTypeSurvivesAggregation(t, m3)
	})
}

// testAggregatedGraphiteMetric tests the write and read of aggregated graphtie metrics.
func testAggregatedGraphiteMetric(t *testing.T, m3 resources.M3Resources) {
	var (
		carbonName         = "foo.bar.baz"
		carbonTarget       = "foo.bar.*"
		carbonLow          = float64(40)
		carbonHigh         = float64(44)
		expectedCarbonMean = float64(42)
	)

	doneCh := make(chan struct{})
	defer func() {
		doneCh <- struct{}{}
		close(doneCh)
	}()
	go func() {
		for {
			select {
			case <-doneCh:
				return
			default:
				require.NoError(t, m3.Coordinator().WriteCarbon(defaultCarbonPort, carbonName, carbonLow, time.Now()))
				require.NoError(t, m3.Coordinator().WriteCarbon(defaultCarbonPort, carbonName, carbonHigh, time.Now()))
				time.Sleep(1 * time.Second)
			}
		}
	}()

	require.NoError(t, resources.RetryWithMaxTime(func() error {
		return verifyGraphiteQuery(m3, carbonTarget, expectedCarbonMean)
	}, 2*time.Minute))
}

func verifyGraphiteQuery(m3 resources.M3Resources, target string, expected float64) error {
	datapoints, err := m3.Coordinator().GraphiteQuery(resources.GraphiteQueryRequest{
		Target: target,
		From:   time.Now().Add(-1000 * time.Second),
		Until:  time.Now(),
	})
	if err != nil {
		return err
	}
	nonNullDPs := filterNull(datapoints)
	if len(nonNullDPs) == 0 {
		return errEmptyResult
	}
	if v := *nonNullDPs[0].Value; v != expected {
		return fmt.Errorf("wrong datapoint result: expected=%f, actual=%f", expected, v)
	}
	return nil
}

func filterNull(datapoints []resources.Datapoint) []resources.Datapoint {
	nonNull := make([]resources.Datapoint, 0, len(datapoints))
	for _, dp := range datapoints {
		if dp.Value != nil {
			nonNull = append(nonNull, dp)
		}
	}
	return nonNull
}

// testRollupRule tests metrics aggregated with a rollup rule.
func testRollupRule(t *testing.T, m3 resources.M3Resources) {
	var (
		numDatapoints = 5
		resolutionSec = 5
		nowTime       = time.Now()
		initWriteTime = nowTime.Truncate(time.Duration(resolutionSec) * time.Second)
		metricName    = "http_requests"

		initVal1 = 42
		valRate1 = 22
		valInc1  = valRate1 * resolutionSec
		tags1    = map[string]string{
			"app":         "nginx_edge",
			"status_code": "500",
			"endpoint":    "/foo/bar",
		}

		initVal2 = 84
		valRate2 = 4
		valInc2  = valRate2 * resolutionSec
		tags2    = map[string]string{
			"app":         "nginx_edge",
			"status_code": "500",
			"endpoint":    "/foo/baz",
		}
	)

	for i := 0; i < numDatapoints; i++ {
		err := m3.Coordinator().WriteProm(
			metricName,
			tags1,
			[]prompb.Sample{{
				Value:     float64(initVal1 + i*valInc1),
				Timestamp: initWriteTime.Add(time.Duration(i*resolutionSec)*time.Second).Unix() * 1000,
			}},
			resources.Headers{headers.PromTypeHeader: []string{"counter"}},
		)
		require.NoError(t, err)
	}

	for i := 0; i < numDatapoints; i++ {
		err := m3.Coordinator().WriteProm(
			metricName,
			tags2,
			[]prompb.Sample{{
				Value:     float64(initVal2 + i*valInc2),
				Timestamp: initWriteTime.Add(time.Duration(i*resolutionSec)*time.Second).Unix() * 1000,
			}},
			resources.Headers{headers.PromTypeHeader: []string{"gauge"}},
		)
		require.NoError(t, err)
	}

	require.NoError(t, resources.RetryWithMaxTime(func() error {
		return verifyPromQuery(
			m3,
			`http_requests_by_status_code{endpoint="/foo/bar"}`,
			float64(valRate1),
		)
	}, 2*time.Minute))

	require.NoError(t, resources.RetryWithMaxTime(func() error {
		return verifyPromQuery(
			m3,
			`http_requests_by_status_code{endpoint="/foo/baz"}`,
			float64(valRate2),
		)
	}, 2*time.Minute))
}

func verifyPromQuery(
	m3 resources.M3Resources,
	queryStr string,
	expected float64,
) error {
	results, err := m3.Coordinator().RangeQuery(
		resources.RangeQueryRequest{
			Query: queryStr,
			Start: time.Now().Add(-1 * time.Hour),
			End:   time.Now().Add(1 * time.Hour),
			Step:  30 * time.Second,
		},
		map[string][]string{
			"M3-Metrics-Type":   {"aggregated"},
			"M3-Storage-Policy": {"5s:6h"},
		},
	)
	if err != nil {
		return err
	}
	if len(results) == 0 {
		return errEmptyResult
	}
	if len(results) > 1 {
		return errors.New("more results than expected")
	}
	if v := float64(results[0].Values[0].Value); v != expected {
		return fmt.Errorf("wrong datapoint result: expected=%f, actual=%f", expected, v)
	}
	return nil
}

// testMetricTypeSurvivesAggregation verifies that the metric type information
// is stored in db after the aggregation.
func testMetricTypeSurvivesAggregation(t *testing.T, m3 resources.M3Resources) {
	nowTime := time.Now()
	value := 42
	metricName := "metric_type_test"

	require.NoError(t, m3.Coordinator().WriteProm(
		metricName,
		map[string]string{
			"label0": "label0",
			"label1": "label1",
			"label2": "label2",
		},
		[]prompb.Sample{{
			Value:     float64(value),
			Timestamp: nowTime.Unix() * 1000,
		}},
		resources.Headers{headers.PromTypeHeader: []string{"counter"}},
	))

	node := m3.Nodes()[0]
	require.NoError(t, resources.Retry(func() error {
		res, err := node.Fetch(&rpc.FetchRequest{
			NameSpace:  "aggregated",
			ID:         `{__name__="metric_type_test",label0="label0",label1="label1",label2="label2"}`,
			RangeStart: nowTime.Add(-1 * time.Hour).Unix(),
			RangeEnd:   nowTime.Add(time.Hour).Unix(),
		})
		if err != nil {
			return err
		}
		if len(res.Datapoints) == 0 {
			return errEmptyResult
		}
		if len(res.Datapoints[0].Annotation) == 0 {
			return errQueryResult
		}
		if res.Datapoints[0].Value != float64(value) {
			return errQueryResult
		}
		return nil
	}))
}


<|start_filename|>src/query/api/v1/middleware/source_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package middleware

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/m3db/m3/src/query/source"
	"github.com/m3db/m3/src/query/util/logging"
	"github.com/m3db/m3/src/x/headers"
	"github.com/m3db/m3/src/x/instrument"

	"github.com/gorilla/mux"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
)

type testSource struct {
	name string
}

var testDeserialize = func(bytes []byte) (interface{}, error) {
	return testSource{string(bytes)}, nil
}

func TestMiddleware(t *testing.T) {
	cases := []struct {
		name         string
		sourceHeader string
		expected     testSource
		expectedLog  string
		deserializer source.Deserializer
		invalidErr   bool
	}{
		{
			name:         "happy path",
			sourceHeader: "foobar",
			expected:     testSource{"foobar"},
			expectedLog:  "foobar",
		},
		{
			name:         "no source header",
			sourceHeader: "",
			expected:     testSource{""},
		},
		{
			name:         "deserialize error",
			sourceHeader: "foobar",
			invalidErr:   true,
			deserializer: func(bytes []byte) (interface{}, error) {
				return nil, errors.New("boom")
			},
		},
	}

	for _, tc := range cases {
		tc := tc
		core, recorded := observer.New(zapcore.InfoLevel)
		l := zap.New(core)
		iOpts := instrument.NewOptions().SetLogger(l)
		t.Run(tc.name, func(t *testing.T) {
			r := mux.NewRouter()
			r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
				l = logging.WithContext(r.Context(), iOpts)
				l.Info("test")
				typed, ok := source.FromContext(r.Context())
				if tc.expected.name == "" {
					require.False(t, ok)
					require.Nil(t, typed)
				} else {
					require.True(t, ok)
					require.Equal(t, tc.expected, typed.(testSource))
				}
			})
			if tc.deserializer == nil {
				tc.deserializer = testDeserialize
			}
			r.Use(Source(Options{
				InstrumentOpts: iOpts,
				Source: SourceOptions{
					Deserializer: tc.deserializer,
				},
			}))
			s := httptest.NewServer(r)
			defer s.Close()

			req, err := http.NewRequestWithContext(context.Background(), "GET", s.URL, nil)
			require.NoError(t, err)
			req.Header.Set(headers.SourceHeader, tc.sourceHeader)
			resp, err := s.Client().Do(req)
			require.NoError(t, err)
			require.NoError(t, resp.Body.Close())
			if tc.invalidErr {
				require.Equal(t, http.StatusBadRequest, resp.StatusCode)
			} else {
				require.Equal(t, http.StatusOK, resp.StatusCode)
				testMsgs := recorded.FilterMessage("test").All()
				require.Len(t, testMsgs, 1)
				entry := testMsgs[0]
				require.Equal(t, "test", entry.Message)
				fields := entry.ContextMap()
				if tc.expectedLog != "" {
					require.Len(t, fields, 1)
					require.Equal(t, tc.expectedLog, fields["source"])
				}
			}
		})
	}
}


<|start_filename|>src/dbnode/storage/limits/permits/permit.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package permits

import (
	"github.com/m3db/m3/src/x/instrument"

	"go.uber.org/atomic"
	"go.uber.org/zap"
)

// permit tracks the quota used by a caller and provides basic sanity checks that a caller
// correctly uses the permit.
type permit struct {
	// immutable state
	quota int64
	iOpts instrument.Options

	// mutable state
	quotaUsed int64
	// refCount is used to check if a caller incorrectly double releases/acquires a permit. the value should
	// always be 0 (nobody holds the permit) or 1 (somebody holds the permit).
	refCount atomic.Int32
}

// NewPermit constructs a new Permit with the provided quota.
func NewPermit(quota int64, iOpts instrument.Options) Permit {
	return &permit{
		quota: quota,
		iOpts: iOpts,
	}
}

func (p *permit) PostRelease() {
	if p.iOpts != nil && p.refCount.Dec() != 0 {
		instrument.EmitAndLogInvariantViolation(p.iOpts, func(l *zap.Logger) {
			l.Error("permit released more than once")
		})
	}
}

func (p *permit) PreAcquire() {
	if p.iOpts != nil && p.refCount.Inc() != 1 {
		instrument.EmitAndLogInvariantViolation(p.iOpts, func(l *zap.Logger) {
			l.Error("permit acquired more than once")
		})
	}
	p.quotaUsed = 0
}

// AllowedQuota is the amount of quota the caller can use with this Permit.
func (p *permit) AllowedQuota() int64 {
	return p.quota
}

// QuotaRemaining is the amount of remaining quota for this Permit. Can be negative if the caller used more quota
// than they were allowed.
func (p *permit) QuotaRemaining() int64 {
	return p.quota - p.quotaUsed
}

// Use adds the quota to the total used quota.
func (p *permit) Use(quota int64) {
	p.quotaUsed += quota
}


<|start_filename|>Dockerfile<|end_filename|>
FROM golang:1.18-stretch

RUN apt-get update && apt-get install -y lsof


<|start_filename|>src/aggregator/integration/placement.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

import (
	"math"
	"testing"

	clusterclient "github.com/m3db/m3/src/cluster/client"
	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/services"
	"github.com/m3db/m3/src/cluster/shard"

	"github.com/stretchr/testify/require"
)

type placementInstanceConfig struct {
	instanceID          string
	shardSetID          uint32
	shardStartInclusive uint32
	shardEndExclusive   uint32
}

func (c *placementInstanceConfig) newPlacementInstance() placement.Instance {
	numShards := c.shardEndExclusive - c.shardStartInclusive
	shardSet := make([]shard.Shard, 0, numShards)
	for shardID := c.shardStartInclusive; shardID < c.shardEndExclusive; shardID++ {
		shard := shard.NewShard(shardID).
			SetState(shard.Available).
			SetCutoverNanos(0).
			SetCutoffNanos(math.MaxInt64)
		shardSet = append(shardSet, shard)
	}
	shards := shard.NewShards(shardSet)
	return placement.NewInstance().
		SetID(c.instanceID).
		SetShards(shards).
		SetShardSetID(c.shardSetID).
		SetEndpoint(c.instanceID).
		SetWeight(1)
}

func newPlacement(numShards int, instances []placement.Instance) placement.Placement {
	shards := make([]uint32, numShards)
	for i := 0; i < numShards; i++ {
		shards[i] = uint32(i)
	}
	var maxShardSetID uint32
	for _, inst := range instances {
		if maxShardSetID < inst.ShardSetID() {
			maxShardSetID = inst.ShardSetID()
		}
	}
	return placement.NewPlacement().
		SetInstances(instances).
		SetShards(shards).
		SetIsSharded(true).
		SetMaxShardSetID(maxShardSetID).
		SetReplicaFactor(1)
}

func setPlacement(
	t *testing.T,
	key string,
	clusterClient clusterclient.Client,
	pl placement.Placement,
) {
	svcs, err := clusterClient.Services(nil)
	require.NoError(t, err)
	ps, err := svcs.PlacementService(services.NewServiceID().SetName(defaultServiceName), placement.NewOptions())
	require.NoError(t, err)
	_, err = ps.Set(pl)
	require.NoError(t, err)

	store, err := clusterClient.KV()
	require.NoError(t, err)

	stagedPlacement, err := placement.NewPlacementsFromLatest(pl)
	require.NoError(t, err)
	stagedPlacementProto, err := stagedPlacement.Proto()
	require.NoError(t, err)
	_, err = store.Set(key, stagedPlacementProto)
	require.NoError(t, err)
}


<|start_filename|>src/dbnode/integration/index_block_orphaned_entry_test.go<|end_filename|>
//go:build integration
// +build integration

//
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

import (
	"fmt"
	"math/rand"
	"runtime"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/client"
	"github.com/m3db/m3/src/dbnode/namespace"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/dbnode/storage"
	"github.com/m3db/m3/src/dbnode/storage/index/compaction"
	xclock "github.com/m3db/m3/src/x/clock"
	"github.com/m3db/m3/src/x/ident"
	xsync "github.com/m3db/m3/src/x/sync"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
)

const (
	numTestSeries     = 5
	concurrentWorkers = 25
	writesPerWorker   = 5
	blockSize         = 2 * time.Hour
)

func TestIndexBlockOrphanedEntry(t *testing.T) {
	nsOpts := namespace.NewOptions().
		SetRetentionOptions(DefaultIntegrationTestRetentionOpts).
		SetIndexOptions(namespace.NewIndexOptions().SetEnabled(true))

	setup := generateTestSetup(t, nsOpts)
	defer setup.Close()

	// Start the server
	log := setup.StorageOpts().InstrumentOptions().Logger()
	require.NoError(t, setup.StartServer())

	// Stop the server
	defer func() {
		assert.NoError(t, setup.StopServer())
		log.Debug("server is now down")
	}()

	client := setup.M3DBClient()
	session, err := client.DefaultSession()
	require.NoError(t, err)

	// Write concurrent metrics to generate multiple entries for the same series
	ids := make([]ident.ID, 0, numTestSeries)
	for i := 0; i < numTestSeries; i++ {
		fooID := ident.StringID(fmt.Sprintf("foo.%v", i))
		ids = append(ids, fooID)

		writeConcurrentMetrics(t, setup, session, fooID)
	}

	// Write metrics for a different series to push current foreground segment
	// to the background. After this, all documents for foo.X exist in background segments
	barID := ident.StringID("bar")
	writeConcurrentMetrics(t, setup, session, barID)

	// Fast-forward to a block rotation
	newBlock := xtime.Now().Truncate(blockSize).Add(blockSize)
	newCurrentTime := newBlock.Add(30 * time.Minute) // Add extra to account for buffer past
	setup.SetNowFn(newCurrentTime)

	// Wait for flush
	log.Info("waiting for block rotation to complete")
	nsID := setup.Namespaces()[0].ID()
	found := xclock.WaitUntil(func() bool {
		filesets, err := fs.IndexFileSetsAt(setup.FilePathPrefix(), nsID, newBlock.Add(-blockSize))
		require.NoError(t, err)
		return len(filesets) == 1
	}, 60*time.Second)
	require.True(t, found)

	// Do post-block rotation writes
	for _, id := range ids {
		writeMetric(t, session, nsID, id, newCurrentTime, 999.0)
	}
	writeMetric(t, session, nsID, barID, newCurrentTime, 999.0)

	// Foreground segments should be in the background again which means updated index entry
	// is now behind the orphaned entry so index reads should fail.
	log.Info("waiting for metrics to be indexed")
	var (
		missing string
		ok      bool
	)
	found = xclock.WaitUntil(func() bool {
		for _, id := range ids {
			ok, err = isIndexedCheckedWithTime(
				t, session, nsID, id, genTags(id), newCurrentTime,
			)
			if !ok || err != nil {
				missing = id.String()
				return false
			}
		}
		return true
	}, 30*time.Second)
	assert.True(t, found, fmt.Sprintf("series %s never indexed\n", missing))
	assert.NoError(t, err)
}

func writeConcurrentMetrics(
	t *testing.T,
	setup TestSetup,
	session client.Session,
	seriesID ident.ID,
) {
	var wg sync.WaitGroup
	nowFn := setup.DB().Options().ClockOptions().NowFn()

	workerPool := xsync.NewWorkerPool(concurrentWorkers)
	workerPool.Init()

	mdID := setup.Namespaces()[0].ID()
	for i := 0; i < concurrentWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			for j := 0; j < writesPerWorker; j++ {
				j := j
				wg.Add(1)
				workerPool.Go(func() {
					defer wg.Done()
					writeMetric(t, session, mdID, seriesID, xtime.ToUnixNano(nowFn()), float64(j))
				})
			}
		}()
	}

	wg.Wait()
}

func genTags(seriesID ident.ID) ident.TagsIterator {
	return ident.NewTagsIterator(ident.NewTags(ident.StringTag("tagName", seriesID.String())))
}

func writeMetric(
	t *testing.T,
	session client.Session,
	nsID ident.ID,
	seriesID ident.ID,
	timestamp xtime.UnixNano,
	value float64,
) {
	err := session.WriteTagged(nsID, seriesID, genTags(seriesID),
		timestamp, value, xtime.Second, nil)
	require.NoError(t, err)
}

func generateTestSetup(t *testing.T, nsOpts namespace.Options) TestSetup {
	md, err := namespace.NewMetadata(testNamespaces[0], nsOpts)
	require.NoError(t, err)

	testOpts := NewTestOptions(t).
		SetNamespaces([]namespace.Metadata{md}).
		SetWriteNewSeriesAsync(true)
	testSetup, err := NewTestSetup(t, testOpts, nil,
		func(s storage.Options) storage.Options {
			s = s.SetCoreFn(func() int {
				return rand.Intn(4) //nolint:gosec
			})
			compactionOpts := s.IndexOptions().ForegroundCompactionPlannerOptions()
			compactionOpts.Levels = []compaction.Level{
				{
					MinSizeInclusive: 0,
					MaxSizeExclusive: 1,
				},
			}
			return s.SetIndexOptions(
				s.IndexOptions().SetForegroundCompactionPlannerOptions(compactionOpts))
		})
	require.NoError(t, err)

	return testSetup
}

func TestIndexBlockOrphanedIndexValuesUpdatedAcrossTimes(t *testing.T) {
	tests := []struct {
		name     string
		numIDs   int
		interval time.Duration
	}{
		{
			name:     "4 series every 100 nanos",
			numIDs:   4,
			interval: 100,
		},
		{
			name:     "4 series every block",
			numIDs:   4,
			interval: blockSize,
		},
		{
			name:     "12 series every 100 nanos",
			numIDs:   12,
			interval: 100,
		},
		{
			name:     "12 series every block",
			numIDs:   12,
			interval: blockSize,
		},
		{
			name:     "120 series every 100 nanos",
			numIDs:   120,
			interval: 100,
		},
		{
			name:     "120 series every block",
			numIDs:   120,
			interval: blockSize,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			testIndexBlockOrphanedIndexValuesUpdatedAcrossTimes(t, tt.numIDs, tt.interval)
		})
	}
}

func testIndexBlockOrphanedIndexValuesUpdatedAcrossTimes(
	t *testing.T, numIDs int, writeInterval time.Duration,
) {
	// Write a metric concurrently for multiple index blocks to generate
	// multiple entries for the same series
	var (
		concurrentWriteMax = runtime.NumCPU() / 2
		ids                = make([]ident.ID, 0, numIDs)
		writerCh           = make(chan func(), concurrentWriteMax)

		nsID = testNamespaces[0]

		writesPerWorker = 5
		writeTimes      = make([]xtime.UnixNano, 0, writesPerWorker)
		retention       = blockSize * time.Duration(1+writesPerWorker)

		seed = time.Now().UnixNano()
		rng  = rand.New(rand.NewSource(seed)) // nolint:gosec
	)

	retOpts := DefaultIntegrationTestRetentionOpts.SetRetentionPeriod(retention)
	nsOpts := namespace.NewOptions().
		SetRetentionOptions(retOpts).
		SetIndexOptions(namespace.NewIndexOptions().SetEnabled(true)).
		SetColdWritesEnabled(true)

	setup := generateTestSetup(t, nsOpts)
	defer setup.Close()

	// Start the server
	log := setup.StorageOpts().InstrumentOptions().Logger()
	log.Info("running test with seed", zap.Int("seed", int(seed)))
	require.NoError(t, setup.StartServer())

	// Stop the server
	defer func() {
		assert.NoError(t, setup.StopServer())
		log.Debug("server is now down")
	}()

	client := setup.M3DBClient()
	session, err := client.DefaultSession()
	require.NoError(t, err)

	var (
		nowFn = setup.DB().Options().ClockOptions().NowFn()
		// NB: write in the middle of a block to avoid block boundaries.
		now = nowFn().Truncate(blockSize / 2)
	)

	for i := 0; i < writesPerWorker; i++ {
		writeTime := xtime.ToUnixNano(now.Add(time.Duration(i) * -writeInterval))
		writeTimes = append(writeTimes, writeTime)
	}

	fns := make([]func(), 0, numIDs*writesPerWorker)
	for i := 0; i < numIDs; i++ {
		fooID := ident.StringID(fmt.Sprintf("foo.%v", i))
		ids = append(ids, fooID)
		fns = append(fns, writeConcurrentMetricsAcrossTime(t, setup, session, writeTimes, fooID)...)
	}

	rng.Shuffle(len(fns), func(i, j int) { fns[i], fns[j] = fns[j], fns[i] })
	var wg sync.WaitGroup
	for i := 0; i < concurrentWriteMax; i++ {
		wg.Add(1)
		go func() {
			for writeFn := range writerCh {
				writeFn()
			}

			wg.Done()
		}()
	}

	for _, fn := range fns {
		writerCh <- fn
	}

	close(writerCh)
	wg.Wait()

	queryIDs := func() {
		notFoundIds := make(notFoundIDs, 0, len(ids)*len(writeTimes))
		for _, id := range ids {
			for _, writeTime := range writeTimes {
				notFoundIds = append(notFoundIds, notFoundID{id: id, runAt: writeTime})
			}
		}

		found := xclock.WaitUntil(func() bool {
			filteredIds := notFoundIds[:0]
			for _, id := range notFoundIds {
				ok, err := isIndexedCheckedWithTime(
					t, session, nsID, id.id, genTags(id.id), id.runAt,
				)
				if !ok || err != nil {
					filteredIds = append(filteredIds, id)
				}
			}

			if len(filteredIds) == 0 {
				return true
			}

			notFoundIds = filteredIds
			return false
		}, time.Second*30)

		require.True(t, found, fmt.Sprintf("series %s never indexed\n", notFoundIds))
	}

	// Ensure all IDs are eventually queryable, even when only in the foreground
	// segments.
	queryIDs()

	// Write metrics for a different series to push current foreground segment
	// to the background. After this, all documents for foo.X exist in background segments
	barID := ident.StringID("bar")
	writeConcurrentMetrics(t, setup, session, barID)

	queryIDs()

	// Fast-forward to a block rotation
	newBlock := xtime.Now().Truncate(blockSize).Add(blockSize)
	newCurrentTime := newBlock.Add(30 * time.Minute) // Add extra to account for buffer past
	setup.SetNowFn(newCurrentTime)

	// Wait for flush
	log.Info("waiting for block rotation to complete")
	found := xclock.WaitUntil(func() bool {
		filesets, err := fs.IndexFileSetsAt(setup.FilePathPrefix(), nsID, newBlock.Add(-blockSize))
		require.NoError(t, err)
		return len(filesets) == 1
	}, 30*time.Second)
	require.True(t, found)

	queryIDs()
}

type notFoundID struct {
	id    ident.ID
	runAt xtime.UnixNano
}

func (i notFoundID) String() string {
	return fmt.Sprintf("{%s: %s}", i.id.String(), i.runAt.String())
}

type notFoundIDs []notFoundID

func (ids notFoundIDs) String() string {
	strs := make([]string, 0, len(ids))
	for _, id := range ids {
		strs = append(strs, id.String())
	}

	return fmt.Sprintf("[%s]", strings.Join(strs, ", "))
}

// writeConcurrentMetricsAcrossTime writes a datapoint for the given series at
// each `writeTime` simultaneously.
func writeConcurrentMetricsAcrossTime(
	t *testing.T,
	setup TestSetup,
	session client.Session,
	writeTimes []xtime.UnixNano,
	seriesID ident.ID,
) []func() {
	workerPool := xsync.NewWorkerPool(concurrentWorkers)
	workerPool.Init()

	mdID := setup.Namespaces()[0].ID()
	fns := make([]func(), 0, len(writeTimes))

	for j, writeTime := range writeTimes {
		j, writeTime := j, writeTime
		fns = append(fns, func() {
			writeMetric(t, session, mdID, seriesID, writeTime, float64(j))
		})
	}

	return fns
}


<|start_filename|>src/dbnode/storage/limits/permits/fixed_permits.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package permits

import (
	"github.com/m3db/m3/src/x/context"
	"github.com/m3db/m3/src/x/instrument"

	"go.uber.org/zap"
)

type fixedPermits struct {
	permits chan Permit
	iOpts   instrument.Options
}

type fixedPermitsManager struct {
	fp fixedPermits
}

var (
	_ Permits = &fixedPermits{}
	_ Manager = &fixedPermitsManager{}
)

// NewFixedPermitsManager returns a permits manager that uses a fixed size of permits.
func NewFixedPermitsManager(size int, quotaPerPermit int64, iOpts instrument.Options) Manager {
	fp := fixedPermits{permits: make(chan Permit, size), iOpts: iOpts}
	for i := 0; i < size; i++ {
		fp.permits <- NewPermit(quotaPerPermit, iOpts)
	}
	return &fixedPermitsManager{fp: fp}
}

func (f *fixedPermitsManager) NewPermits(_ context.Context) (Permits, error) {
	return &f.fp, nil
}

func (f *fixedPermits) Acquire(ctx context.Context) (AcquireResult, error) {
	// don't acquire a permit if ctx is already done.
	select {
	case <-ctx.GoContext().Done():
		return AcquireResult{}, ctx.GoContext().Err()
	default:
	}

	select {
	case <-ctx.GoContext().Done():
		return AcquireResult{}, ctx.GoContext().Err()
	case p := <-f.permits:
		p.PreAcquire()
		return AcquireResult{
			Permit: p,
		}, nil
	}
}

func (f *fixedPermits) TryAcquire(ctx context.Context) (Permit, error) {
	// don't acquire a permit if ctx is already done.
	select {
	case <-ctx.GoContext().Done():
		return nil, ctx.GoContext().Err()
	default:
	}

	select {
	case p := <-f.permits:
		p.PreAcquire()
		return p, nil
	default:
		return nil, nil
	}
}

func (f *fixedPermits) Release(permit Permit) {
	permit.PostRelease()

	select {
	case f.permits <- permit:
	default:
		instrument.EmitAndLogInvariantViolation(f.iOpts, func(l *zap.Logger) {
			l.Error("more permits released than acquired")
		})
	}
}

func (f *fixedPermits) Close() {
}


<|start_filename|>src/aggregator/integration/setup.go<|end_filename|>
// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

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

	"github.com/m3db/m3/src/aggregator/aggregator"
	"github.com/m3db/m3/src/aggregator/aggregator/handler"
	"github.com/m3db/m3/src/aggregator/aggregator/handler/writer"
	aggclient "github.com/m3db/m3/src/aggregator/client"
	"github.com/m3db/m3/src/aggregator/runtime"
	httpserver "github.com/m3db/m3/src/aggregator/server/http"
	m3msgserver "github.com/m3db/m3/src/aggregator/server/m3msg"
	rawtcpserver "github.com/m3db/m3/src/aggregator/server/rawtcp"
	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/services"
	"github.com/m3db/m3/src/cmd/services/m3aggregator/serve"
	"github.com/m3db/m3/src/metrics/metric/aggregated"
	"github.com/m3db/m3/src/metrics/policy"
	"github.com/m3db/m3/src/msg/consumer"
	"github.com/m3db/m3/src/msg/producer"
	"github.com/m3db/m3/src/msg/producer/buffer"
	msgwriter "github.com/m3db/m3/src/msg/producer/writer"
	"github.com/m3db/m3/src/x/instrument"
	xio "github.com/m3db/m3/src/x/io"
	"github.com/m3db/m3/src/x/retry"
	xserver "github.com/m3db/m3/src/x/server"
	xsync "github.com/m3db/m3/src/x/sync"

	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
)

var (
	errServerStartTimedOut   = errors.New("server took too long to start")
	errLeaderElectionTimeout = errors.New("took too long to become leader")
)

type testServerSetups []*testServerSetup

type testServerSetup struct {
	opts             testServerOptions
	m3msgAddr        string
	rawTCPAddr       string
	httpAddr         string
	clientOptions    aggclient.Options
	m3msgServerOpts  m3msgserver.Options
	rawTCPServerOpts rawtcpserver.Options
	httpServerOpts   httpserver.Options
	aggregator       aggregator.Aggregator
	aggregatorOpts   aggregator.Options
	handler          handler.Handler
	electionKey      string
	leaderValue      string
	leaderService    services.LeaderService
	electionCluster  *testCluster
	workerPool       xsync.WorkerPool
	results          map[resultKey]aggregated.MetricWithStoragePolicy
	resultLock       *sync.Mutex

	// Signals.
	doneCh   chan struct{}
	closedCh chan struct{}
}

func newTestServerSetup(t *testing.T, opts testServerOptions) *testServerSetup {
	if opts == nil {
		opts = newTestServerOptions(t)
	}

	// TODO: based on environment variable, use M3MSG aggregator as default
	// server and client, run both legacy and M3MSG tests by setting it to
	// different type in the Makefile.

	// Set up worker pool.
	workerPool := xsync.NewWorkerPool(opts.WorkerPoolSize())
	workerPool.Init()

	// Create the server options.
	rwOpts := xio.NewOptions()
	rawTCPServerOpts := rawtcpserver.NewOptions().SetRWOptions(rwOpts)
	m3msgServerOpts := m3msgserver.NewOptions().
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetServerOptions(xserver.NewOptions()).
		SetConsumerOptions(consumer.NewOptions())
	httpServerOpts := httpserver.NewOptions().
		// use a new mux per test to avoid collisions registering the same handlers between tests.
		SetMux(http.NewServeMux())

	// Creating the aggregator options.
	clockOpts := opts.ClockOptions()
	aggregatorOpts := aggregator.NewOptions(clockOpts).
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetAggregationTypesOptions(opts.AggregationTypesOptions()).
		SetEntryCheckInterval(opts.EntryCheckInterval()).
		SetMaxAllowedForwardingDelayFn(opts.MaxAllowedForwardingDelayFn()).
		SetBufferForPastTimedMetric(opts.BufferForPastTimedMetric()).
		SetBufferForPastTimedMetricFn(func(resolution time.Duration) time.Duration {
			return resolution + opts.BufferForPastTimedMetric()
		}).
		SetDiscardNaNAggregatedValues(opts.DiscardNaNAggregatedValues()).
		SetEntryTTL(opts.EntryTTL())

	// Set up placement manager.
	kvStore, err := opts.ClusterClient().KV()
	require.NoError(t, err)
	placementWatcherOpts := placement.NewWatcherOptions().
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetStagedPlacementKey(opts.PlacementKVKey()).
		SetStagedPlacementStore(kvStore)
	placementManagerOpts := aggregator.NewPlacementManagerOptions().
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetInstanceID(opts.InstanceID()).
		SetWatcherOptions(placementWatcherOpts)
	placementManager := aggregator.NewPlacementManager(placementManagerOpts)
	aggregatorOpts = aggregatorOpts.
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetShardFn(opts.ShardFn()).
		SetPlacementManager(placementManager)

	// Set up flush times manager.
	flushTimesManagerOpts := aggregator.NewFlushTimesManagerOptions().
		SetClockOptions(clockOpts).
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetFlushTimesKeyFmt(opts.FlushTimesKeyFmt()).
		SetFlushTimesStore(kvStore)
	flushTimesManager := aggregator.NewFlushTimesManager(flushTimesManagerOpts)
	aggregatorOpts = aggregatorOpts.SetFlushTimesManager(flushTimesManager)

	// Set up election manager.
	leaderValue := opts.InstanceID()
	campaignOpts, err := services.NewCampaignOptions()
	require.NoError(t, err)
	campaignOpts = campaignOpts.SetLeaderValue(leaderValue)
	electionKey := fmt.Sprintf(opts.ElectionKeyFmt(), opts.ShardSetID())
	electionCluster := opts.ElectionCluster()
	if electionCluster == nil {
		electionCluster = newTestCluster(t)
	}
	leaderService := electionCluster.LeaderService()
	electionManagerOpts := aggregator.NewElectionManagerOptions().
		SetClockOptions(clockOpts).
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetCampaignOptions(campaignOpts).
		SetElectionKeyFmt(opts.ElectionKeyFmt()).
		SetLeaderService(leaderService).
		SetPlacementManager(placementManager).
		SetFlushTimesManager(flushTimesManager)
	electionManager := aggregator.NewElectionManager(electionManagerOpts)
	aggregatorOpts = aggregatorOpts.SetElectionManager(electionManager)

	// Set up flush manager.
	flushManagerOpts := aggregator.NewFlushManagerOptions().
		SetClockOptions(clockOpts).
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetPlacementManager(placementManager).
		SetFlushTimesManager(flushTimesManager).
		SetElectionManager(electionManager).
		SetJitterEnabled(opts.JitterEnabled()).
		SetMaxJitterFn(opts.MaxJitterFn()).
		SetBufferForPastTimedMetric(aggregatorOpts.BufferForPastTimedMetric())
	flushManager := aggregator.NewFlushManager(flushManagerOpts)
	aggregatorOpts = aggregatorOpts.SetFlushManager(flushManager)

	// Set up admin client.
	m3msgOpts := aggclient.NewM3MsgOptions()
	if opts.AggregatorClientType() == aggclient.M3MsgAggregatorClient {
		producer, err := newM3MsgProducer(opts)
		require.NoError(t, err)
		m3msgOpts = m3msgOpts.SetProducer(producer)
	}

	clientOpts := aggclient.NewOptions().
		SetClockOptions(clockOpts).
		SetConnectionOptions(opts.ClientConnectionOptions()).
		SetShardFn(opts.ShardFn()).
		SetWatcherOptions(placementWatcherOpts).
		SetRWOptions(rwOpts).
		SetM3MsgOptions(m3msgOpts).
		SetAggregatorClientType(opts.AggregatorClientType())
	c, err := aggclient.NewClient(clientOpts)
	require.NoError(t, err)
	adminClient, ok := c.(aggclient.AdminClient)
	require.True(t, ok)
	require.NoError(t, adminClient.Init())
	aggregatorOpts = aggregatorOpts.SetAdminClient(adminClient)

	testClientOpts := clientOpts.SetAggregatorClientType(opts.AggregatorClientType())

	// Set up the handler.
	var (
		results    map[resultKey]aggregated.MetricWithStoragePolicy
		resultLock sync.Mutex
	)
	results = make(map[resultKey]aggregated.MetricWithStoragePolicy)
	handler := &capturingHandler{results: results, resultLock: &resultLock}
	pw, err := handler.NewWriter(tally.NoopScope)
	if err != nil {
		panic(err.Error())
	}
	aggregatorOpts = aggregatorOpts.SetFlushHandler(handler).SetPassthroughWriter(pw)

	// Set up entry pool.
	runtimeOpts := runtime.NewOptions()
	entryPool := aggregator.NewEntryPool(nil)
	entryPool.Init(func() *aggregator.Entry {
		return aggregator.NewEntry(nil, runtimeOpts, aggregatorOpts)
	})
	aggregatorOpts = aggregatorOpts.SetEntryPool(entryPool)

	// Set up elem pools.
	counterElemPool := aggregator.NewCounterElemPool(nil)
	aggregatorOpts = aggregatorOpts.SetCounterElemPool(counterElemPool)
	elemOpts := aggregator.NewElemOptions(aggregatorOpts)
	counterElemPool.Init(func() *aggregator.CounterElem {
		return aggregator.MustNewCounterElem(aggregator.ElemData{}, elemOpts)
	})

	timerElemPool := aggregator.NewTimerElemPool(nil)
	aggregatorOpts = aggregatorOpts.SetTimerElemPool(timerElemPool)
	timerElemPool.Init(func() *aggregator.TimerElem {
		return aggregator.MustNewTimerElem(aggregator.ElemData{}, elemOpts)
	})

	gaugeElemPool := aggregator.NewGaugeElemPool(nil)
	aggregatorOpts = aggregatorOpts.SetGaugeElemPool(gaugeElemPool)
	gaugeElemPool.Init(func() *aggregator.GaugeElem {
		return aggregator.MustNewGaugeElem(aggregator.ElemData{}, elemOpts)
	})

	return &testServerSetup{
		opts:             opts,
		rawTCPAddr:       opts.RawTCPAddr(),
		httpAddr:         opts.HTTPAddr(),
		m3msgAddr:        opts.M3MsgAddr(),
		clientOptions:    testClientOpts,
		rawTCPServerOpts: rawTCPServerOpts,
		m3msgServerOpts:  m3msgServerOpts,
		httpServerOpts:   httpServerOpts,
		aggregatorOpts:   aggregatorOpts,
		handler:          handler,
		electionKey:      electionKey,
		leaderValue:      leaderValue,
		leaderService:    leaderService,
		electionCluster:  electionCluster,
		workerPool:       workerPool,
		results:          results,
		resultLock:       &resultLock,
		doneCh:           make(chan struct{}),
		closedCh:         make(chan struct{}),
	}
}

func (ts *testServerSetup) newClient(t *testing.T) *client {
	clientType := ts.opts.AggregatorClientType()
	clientOpts := ts.clientOptions.
		SetAggregatorClientType(clientType)

	if clientType == aggclient.M3MsgAggregatorClient {
		producer, err := newM3MsgProducer(ts.opts)
		require.NoError(t, err)
		m3msgOpts := aggclient.NewM3MsgOptions().SetProducer(producer)
		clientOpts = clientOpts.SetM3MsgOptions(m3msgOpts)
	}

	testClient, err := aggclient.NewClient(clientOpts)
	require.NoError(t, err)
	testAdminClient, ok := testClient.(aggclient.AdminClient)
	require.True(t, ok)
	return newClient(testAdminClient)
}

func (ts *testServerSetup) getStatusResponse(path string, response interface{}) error {
	resp, err := http.Get("http://" + ts.httpAddr + path) //nolint
	if err != nil {
		return err
	}

	defer resp.Body.Close() //nolint:errcheck
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("got a non-200 status code: %v", resp.StatusCode)
	}
	return json.Unmarshal(b, response)
}

func (ts *testServerSetup) waitUntilServerIsUp() error {
	isUp := func() bool {
		var resp httpserver.Response
		if err := ts.getStatusResponse(httpserver.HealthPath, &resp); err != nil {
			return false
		}

		if resp.State == "OK" {
			return true
		}

		return false
	}

	if waitUntil(isUp, ts.opts.ServerStateChangeTimeout()) {
		return nil
	}

	return errServerStartTimedOut
}

func (ts *testServerSetup) startServer() error {
	errCh := make(chan error, 1)

	// Creating the aggregator.
	ts.aggregator = aggregator.NewAggregator(ts.aggregatorOpts)
	if err := ts.aggregator.Open(); err != nil {
		return err
	}

	instrumentOpts := instrument.NewOptions()
	serverOpts := serve.NewOptions(instrumentOpts).
		SetM3MsgAddr(ts.m3msgAddr).
		SetM3MsgServerOpts(ts.m3msgServerOpts).
		SetRawTCPAddr(ts.rawTCPAddr).
		SetRawTCPServerOpts(ts.rawTCPServerOpts).
		SetHTTPAddr(ts.httpAddr).
		SetHTTPServerOpts(ts.httpServerOpts).
		SetRWOptions(xio.NewOptions())

	go func() {
		if err := serve.Serve(
			ts.aggregator,
			ts.doneCh,
			serverOpts,
		); err != nil {
			select {
			case errCh <- err:
			default:
			}
		}
		close(ts.closedCh)
	}()

	go func() {
		select {
		case errCh <- ts.waitUntilServerIsUp():
		default:
		}
	}()

	return <-errCh
}

func (ts *testServerSetup) waitUntilLeader() error {
	isLeader := func() bool {
		var resp httpserver.StatusResponse
		if err := ts.getStatusResponse(httpserver.StatusPath, &resp); err != nil {
			return false
		}

		if resp.Status.FlushStatus.ElectionState == aggregator.LeaderState {
			return true
		}
		return false
	}

	if waitUntil(isLeader, ts.opts.ElectionStateChangeTimeout()) {
		return nil
	}

	return errLeaderElectionTimeout
}

// return the metric value for the provided params.
func (ts *testServerSetup) value(timeNanos int64, metricID string, sp policy.StoragePolicy) float64 {
	ts.resultLock.Lock()
	defer ts.resultLock.Unlock()
	return ts.results[resultKey{
		timeNanos:     timeNanos,
		metricID:      metricID,
		storagePolicy: sp,
	}].Value
}

// remove the value from the results if it matches the provided value.
func (ts *testServerSetup) removeIf(timeNanos int64, metricID string, sp policy.StoragePolicy, value float64) bool {
	ts.resultLock.Lock()
	defer ts.resultLock.Unlock()

	key := resultKey{
		timeNanos:     timeNanos,
		metricID:      metricID,
		storagePolicy: sp,
	}
	m, ok := ts.results[key]
	if !ok {
		return false
	}
	if m.Value != value {
		return false
	}
	delete(ts.results, key)
	return true
}

type resultKey struct {
	timeNanos     int64
	metricID      string
	storagePolicy policy.StoragePolicy
}

func (ts *testServerSetup) sortedResults() []aggregated.MetricWithStoragePolicy {
	ts.resultLock.Lock()
	defer ts.resultLock.Unlock()
	metrics := make([]aggregated.MetricWithStoragePolicy, 0, len(ts.results))
	for _, r := range ts.results {
		metrics = append(metrics, r)
	}
	sort.Sort(byTimeIDPolicyAscending(metrics))
	return metrics
}

func (ts *testServerSetup) stopServer() error {
	if err := ts.aggregator.Close(); err != nil {
		return err
	}
	close(ts.doneCh)

	// Wait for graceful server shutdown.
	<-ts.closedCh
	return nil
}

func (ts *testServerSetup) close() {
	ts.electionCluster.Close()
}

func (tss testServerSetups) newClient(t *testing.T) *client {
	require.NotEmpty(t, tss)
	// NB: the client can be constructed from any of the setups. The client does the routing and
	// sends the writes to the server which holds related shard.
	return tss[0].newClient(t)
}

func newM3MsgProducer(opts testServerOptions) (producer.Producer, error) {
	svcs, err := opts.ClusterClient().Services(nil)
	if err != nil {
		return nil, err
	}

	bufferOpts := buffer.NewOptions().
		// NB: the default values of cleanup retry options causes very slow m3msg client shutdowns
		// in some of the tests. The values below were set to avoid that.
		SetCleanupRetryOptions(retry.NewOptions().SetInitialBackoff(100 * time.Millisecond).SetMaxRetries(0))
	buffer, err := buffer.NewBuffer(bufferOpts)
	if err != nil {
		return nil, err
	}
	connectionOpts := msgwriter.NewConnectionOptions().
		SetNumConnections(1).
		SetFlushInterval(1 * time.Millisecond)
	writerOpts := msgwriter.NewOptions().
		SetInstrumentOptions(opts.InstrumentOptions()).
		SetTopicName(opts.TopicName()).
		SetTopicService(opts.TopicService()).
		SetServiceDiscovery(svcs).
		SetMessageQueueNewWritesScanInterval(10 * time.Millisecond).
		SetMessageQueueFullScanInterval(100 * time.Millisecond).
		SetConnectionOptions(connectionOpts)
	writer := msgwriter.NewWriter(writerOpts)
	producerOpts := producer.NewOptions().
		SetBuffer(buffer).
		SetWriter(writer)
	producer := producer.NewProducer(producerOpts)
	return producer, nil
}

type capturingWriter struct {
	results    map[resultKey]aggregated.MetricWithStoragePolicy
	resultLock *sync.Mutex
}

func (w *capturingWriter) Write(mp aggregated.ChunkedMetricWithStoragePolicy) error {
	w.resultLock.Lock()
	defer w.resultLock.Unlock()
	var fullID []byte
	fullID = append(fullID, mp.ChunkedID.Prefix...)
	fullID = append(fullID, mp.ChunkedID.Data...)
	fullID = append(fullID, mp.ChunkedID.Suffix...)
	var clonedAnnotation []byte
	clonedAnnotation = append(clonedAnnotation, mp.Annotation...)
	metric := aggregated.Metric{
		ID:         fullID,
		TimeNanos:  mp.TimeNanos,
		Value:      mp.Value,
		Annotation: clonedAnnotation,
	}
	key := resultKey{
		timeNanos:     metric.TimeNanos,
		metricID:      string(fullID),
		storagePolicy: mp.StoragePolicy,
	}
	w.results[key] = aggregated.MetricWithStoragePolicy{
		Metric:        metric,
		StoragePolicy: mp.StoragePolicy,
	}
	return nil
}

func (w *capturingWriter) Flush() error { return nil }
func (w *capturingWriter) Close() error { return nil }

type capturingHandler struct {
	results    map[resultKey]aggregated.MetricWithStoragePolicy
	resultLock *sync.Mutex
}

func (h *capturingHandler) NewWriter(tally.Scope) (writer.Writer, error) {
	return &capturingWriter{results: h.results, resultLock: h.resultLock}, nil
}

func (h *capturingHandler) Close() {}


<|start_filename|>src/integration/resources/inprocess/aggregator.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package inprocess

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"time"

	m3agg "github.com/m3db/m3/src/aggregator/aggregator"
	"github.com/m3db/m3/src/aggregator/server"
	"github.com/m3db/m3/src/aggregator/tools/deploy"
	"github.com/m3db/m3/src/cmd/services/m3aggregator/config"
	"github.com/m3db/m3/src/integration/resources"
	nettest "github.com/m3db/m3/src/integration/resources/net"
	"github.com/m3db/m3/src/x/config/hostid"
	xos "github.com/m3db/m3/src/x/os"

	"github.com/google/uuid"
	"go.uber.org/zap"
	"gopkg.in/yaml.v2"
)

var errAggregatorNotStarted = errors.New("aggregator instance has not started")

// Aggregator is an in-process implementation of resources.Aggregator for use
// in integration tests.
type Aggregator struct {
	cfg     config.Configuration
	logger  *zap.Logger
	tmpDirs []string
	startFn AggregatorStartFn

	started    bool
	httpClient deploy.AggregatorClient

	interruptCh chan<- error
	shutdownCh  <-chan struct{}
}

// AggregatorOptions are options of starting an in-process aggregator.
type AggregatorOptions struct {
	// Logger is the logger to use for the in-process aggregator.
	Logger *zap.Logger
	// StartFn is a custom function that can be used to start the Aggregator.
	StartFn AggregatorStartFn
	// Start indicates whether to start the aggregator instance
	Start bool
	// GeneratePorts will automatically update the config to use open ports
	// if set to true. If false, configuration is used as-is re: ports.
	GeneratePorts bool
	// GenerateHostID will automatically update the host ID specified in
	// the config if set to true. If false, configuration is used as-is re: host ID.
	GenerateHostID bool
}

// NewAggregatorFromYAML creates a new in-process aggregator based on the yaml configuration
// and options provided.
func NewAggregatorFromYAML(yamlCfg string, opts AggregatorOptions) (resources.Aggregator, error) {
	var cfg config.Configuration
	if err := yaml.Unmarshal([]byte(yamlCfg), &cfg); err != nil {
		return nil, err
	}

	return NewAggregator(cfg, opts)
}

// NewAggregator creates a new in-process aggregator based on the configuration
// and options provided.
func NewAggregator(cfg config.Configuration, opts AggregatorOptions) (resources.Aggregator, error) {
	cfg, tmpDirs, err := updateAggregatorConfig(cfg, opts)
	if err != nil {
		return nil, err
	}

	// configure logger
	hostID, err := cfg.AggregatorOrDefault().HostID.Resolve()
	if err != nil {
		return nil, err
	}

	loggingCfg := cfg.LoggingOrDefault()
	if len(loggingCfg.Fields) == 0 {
		loggingCfg.Fields = make(map[string]interface{})
	}
	loggingCfg.Fields["component"] = fmt.Sprintf("m3aggregator:%s", hostID)

	if opts.Logger == nil {
		var err error
		opts.Logger, err = resources.NewLogger()
		if err != nil {
			return nil, err
		}
	}

	agg := &Aggregator{
		cfg:        cfg,
		logger:     opts.Logger,
		tmpDirs:    tmpDirs,
		startFn:    opts.StartFn,
		started:    false,
		httpClient: deploy.NewAggregatorClient(&http.Client{}),
	}

	if opts.Start {
		agg.Start()
	}

	return agg, nil
}

// HostDetails returns the aggregator's host details.
func (a *Aggregator) HostDetails() (*resources.InstanceInfo, error) {
	id, err := a.cfg.AggregatorOrDefault().HostID.Resolve()
	if err != nil {
		return nil, err
	}

	addr, p, err := net.SplitHostPort(a.cfg.HTTPOrDefault().ListenAddress)
	if err != nil {
		return nil, err
	}

	port, err := strconv.Atoi(p)
	if err != nil {
		return nil, err
	}

	m3msgAddr, m3msgP, err := net.SplitHostPort(a.cfg.M3MsgOrDefault().Server.ListenAddress)
	if err != nil {
		return nil, err
	}

	m3msgPort, err := strconv.Atoi(m3msgP)
	if err != nil {
		return nil, err
	}

	return &resources.InstanceInfo{
		ID:           id,
		Env:          a.cfg.KVClientOrDefault().Etcd.Env,
		Zone:         a.cfg.KVClientOrDefault().Etcd.Zone,
		Address:      addr,
		Port:         uint32(port),
		M3msgAddress: m3msgAddr,
		M3msgPort:    uint32(m3msgPort),
	}, nil
}

// Start starts the aggregator instance.
//nolint:dupl
func (a *Aggregator) Start() {
	if a.started {
		a.logger.Debug("aggregator instance has started already")
		return
	}
	a.started = true

	if a.startFn != nil {
		a.interruptCh, a.shutdownCh = a.startFn(&a.cfg)
		return
	}

	interruptCh := make(chan error, 1)
	shutdownCh := make(chan struct{}, 1)

	go func() {
		server.Run(server.RunOptions{
			Config:      a.cfg,
			InterruptCh: interruptCh,
			ShutdownCh:  shutdownCh,
		})
	}()

	a.interruptCh = interruptCh
	a.shutdownCh = shutdownCh
}

// IsHealthy determines whether an instance is healthy.
func (a *Aggregator) IsHealthy() error {
	if !a.started {
		return errAggregatorNotStarted
	}

	return a.httpClient.IsHealthy(a.cfg.HTTPOrDefault().ListenAddress)
}

// Status returns the instance status.
func (a *Aggregator) Status() (m3agg.RuntimeStatus, error) {
	if !a.started {
		return m3agg.RuntimeStatus{}, errAggregatorNotStarted
	}

	return a.httpClient.Status(a.cfg.HTTPOrDefault().ListenAddress)
}

// Resign asks an aggregator instance to give up its current leader role if applicable.
func (a *Aggregator) Resign() error {
	if !a.started {
		return errAggregatorNotStarted
	}

	return a.httpClient.Resign(a.cfg.HTTPOrDefault().ListenAddress)
}

// Close closes the wrapper and releases any held resources, including
// deleting docker containers.
func (a *Aggregator) Close() error {
	if !a.started {
		return errAggregatorNotStarted
	}

	defer func() {
		for _, dir := range a.tmpDirs {
			if err := os.RemoveAll(dir); err != nil {
				a.logger.Error("error removing temp directory", zap.String("dir", dir), zap.Error(err))
			}
		}
		a.started = false
	}()

	select {
	case a.interruptCh <- xos.NewInterruptError("in-process aggregator being shut down"):
	case <-time.After(interruptTimeout):
		return errors.New("timeout sending interrupt. closing without graceful shutdown")
	}

	select {
	case <-a.shutdownCh:
	case <-time.After(shutdownTimeout):
		return errors.New("timeout waiting for shutdown notification. server closing may" +
			" not be completely graceful")
	}

	return nil
}

// Configuration returns a copy of the configuration used to
// start this aggregator.
func (a *Aggregator) Configuration() config.Configuration {
	return a.cfg
}

func updateAggregatorConfig(
	cfg config.Configuration,
	opts AggregatorOptions,
) (config.Configuration, []string, error) {
	var (
		tmpDirs []string
		err     error
	)

	// Replace host ID with a config-based version.
	if opts.GenerateHostID {
		cfg = updateAggregatorHostID(cfg)
	}

	// Replace any ports with open ports
	if opts.GeneratePorts {
		cfg, err = updateAggregatorPorts(cfg)
		if err != nil {
			return config.Configuration{}, nil, err
		}
	}

	// Replace any filepath with a temporary directory
	cfg, tmpDirs, err = updateAggregatorFilepaths(cfg)
	if err != nil {
		return config.Configuration{}, nil, err
	}

	return cfg, tmpDirs, nil
}

func updateAggregatorHostID(cfg config.Configuration) config.Configuration {
	hostID := uuid.New().String()
	aggCfg := cfg.AggregatorOrDefault()
	aggCfg.HostID = &hostid.Configuration{
		Resolver: hostid.ConfigResolver,
		Value:    &hostID,
	}
	cfg.Aggregator = &aggCfg

	return cfg
}

func updateAggregatorPorts(cfg config.Configuration) (config.Configuration, error) {
	httpCfg := cfg.HTTPOrDefault()
	addr, _, err := nettest.GeneratePort(httpCfg.ListenAddress)
	if err != nil {
		return cfg, err
	}
	httpCfg.ListenAddress = addr
	cfg.HTTP = &httpCfg

	metricsCfg := cfg.MetricsOrDefault()
	if metricsCfg.PrometheusReporter != nil && metricsCfg.PrometheusReporter.ListenAddress != "" {
		addr, _, err := nettest.GeneratePort(metricsCfg.PrometheusReporter.ListenAddress)
		if err != nil {
			return cfg, err
		}
		promReporter := *metricsCfg.PrometheusReporter
		promReporter.ListenAddress = addr
		metricsCfg.PrometheusReporter = &promReporter
	}
	cfg.Metrics = &metricsCfg

	m3msgCfg := cfg.M3MsgOrDefault()
	if m3msgAddr := m3msgCfg.Server.ListenAddress; m3msgAddr != "" {
		addr, _, err := nettest.GeneratePort(m3msgAddr)
		if err != nil {
			return cfg, err
		}
		m3msgCfg.Server.ListenAddress = addr
	}
	cfg.M3Msg = &m3msgCfg

	return cfg, nil
}

func updateAggregatorFilepaths(cfg config.Configuration) (config.Configuration, []string, error) {
	tmpDirs := make([]string, 0, 1)

	kvCfg := cfg.KVClientOrDefault()
	if kvCfg.Etcd != nil {
		dir, err := ioutil.TempDir("", "m3agg-*")
		if err != nil {
			return cfg, tmpDirs, err
		}
		tmpDirs = append(tmpDirs, dir)
		kvCfg.Etcd.CacheDir = dir
	}
	cfg.KVClient = &kvCfg

	return cfg, tmpDirs, nil
}


<|start_filename|>src/query/functions/scalar/time.go<|end_filename|>
// Copyright (c) 2019 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package scalar

import (
	"github.com/m3db/m3/src/query/block"
	"github.com/m3db/m3/src/query/executor/transform"
	"github.com/m3db/m3/src/query/models"
	"github.com/m3db/m3/src/query/parser"
	"github.com/m3db/m3/src/query/util/logging"

	"go.uber.org/zap"
)

type timeOp struct {
	tagOptions models.TagOptions
}

func (o timeOp) OpType() string {
	return TimeType
}

func (o timeOp) String() string {
	return "type: time"
}

func (o timeOp) Node(
	controller *transform.Controller,
	opts transform.Options,
) parser.Source {
	return &timeNode{
		controller: controller,
		tagOptions: o.tagOptions,
		opts:       opts,
	}
}

// NewTimeOp creates an operation that yields a time-based source.
func NewTimeOp(tagOptions models.TagOptions) (parser.Params, error) {
	return &timeOp{
		tagOptions: tagOptions,
	}, nil
}

// timeNode is the execution node for time source.
type timeNode struct {
	tagOptions models.TagOptions
	controller *transform.Controller
	opts       transform.Options
}

// Execute runs the time source's pipeline.
func (n *timeNode) Execute(queryCtx *models.QueryContext) error {
	bounds := n.opts.TimeSpec().Bounds()
	meta := block.Metadata{
		Bounds:         bounds,
		Tags:           models.NewTags(0, n.tagOptions),
		ResultMetadata: block.NewResultMetadata(),
	}

	seriesMeta := []block.SeriesMeta{
		{
			Tags: models.NewTags(0, n.tagOptions),
			Name: []byte(TimeType),
		},
	}

	builder := block.NewColumnBlockBuilder(queryCtx, meta, seriesMeta)
	steps := bounds.Steps()
	err := builder.AddCols(steps)
	if err != nil {
		return err
	}

	for i := 0; i < steps; i++ {
		t, err := bounds.TimeForIndex(i)
		if err != nil {
			return err
		}

		timeVal := float64(t.Seconds())
		if err := builder.AppendValue(i, timeVal); err != nil {
			return err
		}
	}

	block := builder.BuildAsType(block.BlockTime)
	if n.opts.Debug() {
		// Ignore any errors
		iter, _ := block.StepIter()
		if iter != nil {
			logging.WithContext(queryCtx.Ctx, n.opts.InstrumentOptions()).
				Info("time node", zap.String("meta", block.Meta().String()))
		}
	}

	if err := n.controller.Process(queryCtx, block); err != nil {
		block.Close()
		// Fail on first error
		return err
	}

	block.Close()
	return nil
}


<|start_filename|>src/dbnode/integration/disk_cleanup_index_corrupted_test.go<|end_filename|>
//go:build integration
// +build integration

// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

import (
	"os"
	"reflect"
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/namespace"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/dbnode/retention"
	xclock "github.com/m3db/m3/src/x/clock"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestDiskCleanupIndexCorrupted(t *testing.T) {
	if testing.Short() {
		t.SkipNow() // Just skip if we're doing a short run
	}

	var (
		rOpts        = retention.NewOptions().SetRetentionPeriod(48 * time.Hour)
		nsBlockSize  = time.Hour
		idxBlockSize = 2 * time.Hour

		nsROpts = rOpts.SetBlockSize(nsBlockSize)
		idxOpts = namespace.NewIndexOptions().SetBlockSize(idxBlockSize).SetEnabled(true)
		nsOpts  = namespace.NewOptions().
			SetCleanupEnabled(true).
			SetRetentionOptions(nsROpts).
			SetIndexOptions(idxOpts)
	)

	ns, err := namespace.NewMetadata(testNamespaces[0], nsOpts)
	require.NoError(t, err)

	opts := NewTestOptions(t).
		SetNamespaces([]namespace.Metadata{ns})
	setup, err := NewTestSetup(t, opts, nil)
	require.NoError(t, err)
	defer setup.Close()

	filePathPrefix := setup.StorageOpts().CommitLogOptions().FilesystemOptions().FilePathPrefix()

	// Now create some fileset files
	var (
		filesetsIdentifiers = make([]fs.FileSetFileIdentifier, 0)
		numVolumes          = 4

		now         = setup.NowFn()().Truncate(idxBlockSize)
		blockStarts = []xtime.UnixNano{
			now.Add(-3 * idxBlockSize),
			now.Add(-2 * idxBlockSize),
			now.Add(-1 * idxBlockSize),
		}
	)
	for _, blockStart := range blockStarts {
		for idx := 0; idx < numVolumes; idx++ {
			filesetsIdentifiers = append(filesetsIdentifiers, fs.FileSetFileIdentifier{
				Namespace:   ns.ID(),
				BlockStart:  blockStart,
				VolumeIndex: idx,
			})
		}
	}
	writeIndexFileSetFiles(t, setup.StorageOpts(), ns, filesetsIdentifiers)

	filesets := fs.ReadIndexInfoFiles(fs.ReadIndexInfoFilesOptions{
		FilePathPrefix:   filePathPrefix,
		Namespace:        ns.ID(),
		ReaderBufferSize: setup.FilesystemOpts().InfoReaderBufferSize(),
	})
	require.Len(t, filesets, len(blockStarts)*numVolumes)

	filesThatShouldBeKept := make([]string, 0)
	keep := func(files []string) {
		filesThatShouldBeKept = append(filesThatShouldBeKept, files...)
	}
	// Corrupt some filesets.
	forBlockStart(blockStarts[0], filesets, func(filesets []fs.ReadIndexInfoFileResult) {
		keep(missingDigest(t, filesets[0])) // most recent volume index for volume type
		corruptedInfo(t, filesets[1])
		missingInfo(t, filesets[2])
		keep(corruptedInfo(t, filesets[3])) // corrupted info files are kept if it's the most recent volume index
	})

	forBlockStart(blockStarts[1], filesets, func(filesets []fs.ReadIndexInfoFileResult) {
		keep(filesets[0].AbsoluteFilePaths)
		missingDigest(t, filesets[1])
		missingDigest(t, filesets[2])
		keep(missingDigest(t, filesets[3])) // most recent volume index for volume type
	})

	forBlockStart(blockStarts[2], filesets, func(filesets []fs.ReadIndexInfoFileResult) {
		missingInfo(t, filesets[0])
		corruptedInfo(t, filesets[1])
		missingDigest(t, filesets[2])
		keep(filesets[3].AbsoluteFilePaths) // most recent volume index for volume type
	})
	sort.Strings(filesThatShouldBeKept)

	// Start the server
	log := setup.StorageOpts().InstrumentOptions().Logger()
	require.NoError(t, setup.StartServer())
	log.Debug("server is now up")

	// Stop the server
	defer func() {
		require.NoError(t, setup.StopServer())
		log.Debug("server is now down")
	}()

	// Check if corrupted files have been deleted
	waitTimeout := 30 * time.Second
	deleted := xclock.WaitUntil(func() bool {
		files, err := fs.IndexFileSetsBefore(filePathPrefix, ns.ID(), now.Add(time.Minute))
		require.NoError(t, err)
		sort.Strings(files)
		return reflect.DeepEqual(files, filesThatShouldBeKept)
	}, waitTimeout)
	if !assert.True(t, deleted) {
		files, err := fs.IndexFileSetsBefore(filePathPrefix, ns.ID(), now.Add(time.Minute))
		require.NoError(t, err)
		sort.Strings(files)
		require.Equal(t, filesThatShouldBeKept, files)
	}
}

func forBlockStart(
	blockStart xtime.UnixNano,
	filesets []fs.ReadIndexInfoFileResult,
	fn func([]fs.ReadIndexInfoFileResult),
) {
	res := make([]fs.ReadIndexInfoFileResult, 0)
	for _, f := range filesets {
		if f.ID.BlockStart.Equal(blockStart) {
			res = append(res, f)
		}
	}
	sort.Slice(res, func(i, j int) bool {
		return res[i].ID.VolumeIndex < res[j].ID.VolumeIndex
	})
	fn(res)
}

func corruptedInfo(t *testing.T, fileset fs.ReadIndexInfoFileResult) []string {
	for _, f := range fileset.AbsoluteFilePaths {
		if strings.Contains(f, "info.db") {
			require.NoError(t, os.Truncate(f, 0))
			return fileset.AbsoluteFilePaths
		}
	}
	require.Fail(t, "could not find info file")
	return nil
}

func missingInfo(t *testing.T, fileset fs.ReadIndexInfoFileResult) []string { //nolint:unparam
	for i, f := range fileset.AbsoluteFilePaths {
		if strings.Contains(f, "info.db") {
			require.NoError(t, os.Remove(f))
			res := make([]string, 0)
			res = append(res, fileset.AbsoluteFilePaths[:i]...)
			res = append(res, fileset.AbsoluteFilePaths[i+1:]...)
			return res
		}
	}
	require.Fail(t, "could not find info file")
	return nil
}

func missingDigest(t *testing.T, fileset fs.ReadIndexInfoFileResult) []string {
	for i, f := range fileset.AbsoluteFilePaths {
		if strings.Contains(f, "digest.db") {
			require.NoError(t, os.Remove(f))
			res := make([]string, 0)
			res = append(res, fileset.AbsoluteFilePaths[:i]...)
			res = append(res, fileset.AbsoluteFilePaths[i+1:]...)
			return res
		}
	}
	require.Fail(t, "could not find digest file")
	return nil
}


<|start_filename|>src/query/api/v1/middleware/metrics_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package middleware

import (
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"github.com/m3db/m3/src/cmd/services/m3query/config"
	"github.com/m3db/m3/src/query/api/v1/route"
	"github.com/m3db/m3/src/query/parser/promql"
	"github.com/m3db/m3/src/x/headers"
	"github.com/m3db/m3/src/x/instrument"
	"github.com/m3db/m3/src/x/tallytest"

	"github.com/gorilla/mux"
	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
)

func TestResponseMetrics(t *testing.T) {
	scope := tally.NewTestScope("", nil)
	iOpts := instrument.NewOptions().SetMetricsScope(scope)

	r := mux.NewRouter()
	route := r.NewRoute()
	opts := Options{
		InstrumentOpts: iOpts,
		Route:          route,
	}

	h := ResponseMetrics(opts).Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	}))
	route.Path("/test").Handler(h)

	server := httptest.NewServer(r)
	defer server.Close()

	resp, err := server.Client().Get(server.URL + "/test?foo=bar") //nolint: noctx
	require.NoError(t, err)
	require.NoError(t, resp.Body.Close())

	snapshot := scope.Snapshot()
	tallytest.AssertCounterValue(t, 1, snapshot, "request", map[string]string{
		"path":   "/test",
		"status": "200",
		"type":   "coordinator",
	})

	hist := snapshot.Histograms()
	require.True(t, len(hist) == 1)
	for _, h := range hist {
		require.Equal(t, "latency", h.Name())
		require.Equal(t, map[string]string{
			"path": "/test",
			"type": "coordinator",
		}, h.Tags())
	}
}

func TestResponseMetricsCustomMetricType(t *testing.T) {
	scope := tally.NewTestScope("", nil)
	iOpts := instrument.NewOptions().SetMetricsScope(scope)

	r := mux.NewRouter()
	route := r.NewRoute()
	opts := Options{
		InstrumentOpts: iOpts,
		Route:          route,
	}

	h := ResponseMetrics(opts).Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.Header.Set(headers.CustomResponseMetricsType, "foo")
		w.WriteHeader(200)
	}))
	route.Path("/test").Handler(h)

	server := httptest.NewServer(r)
	defer server.Close()

	resp, err := server.Client().Get(server.URL + "/test?foo=bar") //nolint: noctx
	require.NoError(t, err)
	require.NoError(t, resp.Body.Close())

	snapshot := scope.Snapshot()
	tallytest.AssertCounterValue(t, 1, snapshot, "request", map[string]string{
		"path":   "/test",
		"status": "200",
		"type":   "foo",
	})

	hist := snapshot.Histograms()
	require.True(t, len(hist) == 1)
	for _, h := range hist {
		require.Equal(t, "latency", h.Name())
		require.Equal(t, map[string]string{
			"path": "/test",
			"type": "foo",
		}, h.Tags())
	}
}

var parseQueryParams ParseQueryParams = func(r *http.Request, _ time.Time) (QueryParams, error) {
	if err := r.ParseForm(); err != nil {
		return QueryParams{}, err
	}
	params := QueryParams{
		Query: r.FormValue("query"),
	}
	if s := r.FormValue("start"); s != "" {
		start, err := strconv.Atoi(r.FormValue("start"))
		if err != nil {
			return QueryParams{}, err
		}
		params.Start = time.Unix(int64(start), 0)
	}

	if s := r.FormValue("end"); s != "" {
		end, err := strconv.Atoi(r.FormValue("end"))
		if err != nil {
			return QueryParams{}, err
		}
		params.End = time.Unix(int64(end), 0)
	}
	return params, nil
}

func TestLargeResponseMetrics(t *testing.T) {
	scope := tally.NewTestScope("", nil)
	iOpts := instrument.NewOptions().SetMetricsScope(scope)

	r := mux.NewRouter()
	route := r.NewRoute()
	opts := Options{
		InstrumentOpts: iOpts,
		Route:          route,
		Metrics: MetricsOptions{
			Config: config.MetricsMiddlewareConfiguration{
				QueryEndpointsClassification: config.QueryClassificationConfig{
					ResultsBuckets:  []int{1, 10},
					DurationBuckets: []time.Duration{1 * time.Minute, 15 * time.Minute},
				},
			},
			ParseQueryParams: parseQueryParams,
		},
	}

	h := ResponseMetrics(opts).Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add(headers.FetchedSeriesCount, "15")
		w.WriteHeader(200)
	}))
	route.Path("/api/v1/query").Handler(h)

	server := httptest.NewServer(r)
	defer server.Close()

	resp, err := server.Client().Get(server.URL + "/api/v1/query?query=rate(up[20m])") //nolint: noctx
	require.NoError(t, err)
	require.NoError(t, resp.Body.Close())

	snapshot := scope.Snapshot()
	tallytest.AssertCounterValue(t, 1, snapshot, "request", map[string]string{
		"path":                 "/api/v1/query",
		"status":               "200",
		"type":                 "coordinator",
		resultsClassification:  "10",
		durationClassification: "15m0s",
	})

	hist := snapshot.Histograms()
	require.True(t, len(hist) == 1)
	for _, h := range hist {
		require.Equal(t, "latency", h.Name())
		require.Equal(t, map[string]string{
			"path":                 "/api/v1/query",
			"type":                 "coordinator",
			resultsClassification:  "10",
			durationClassification: "15m0s",
		}, h.Tags())
	}
}

func TestMultipleLargeResponseMetricsWithLatencyStatus(t *testing.T) {
	testMultipleLargeResponseMetrics(t, true)
}

func TestMultipleLargeResponseMetricsWithoutLatencyStatus(t *testing.T) {
	testMultipleLargeResponseMetrics(t, false)
}

func TestCustomMetricsRepeatedGets(t *testing.T) {
	scope := tally.NewTestScope("", nil)
	iOpts := instrument.NewOptions().SetMetricsScope(scope)

	cm := newCustomMetrics(iOpts)

	_ = cm.getOrCreate("foo")
	_ = cm.getOrCreate("foo")
	require.Equal(t, len(cm.metrics), 1)

	_ = cm.getOrCreate("foo2")
	require.Equal(t, len(cm.metrics), 2)
}

func testMultipleLargeResponseMetrics(t *testing.T, addStatus bool) {
	scope := tally.NewTestScope("", nil)
	iOpts := instrument.NewOptions().SetMetricsScope(scope)

	r := mux.NewRouter()
	route := r.NewRoute()
	opts := Options{
		InstrumentOpts: iOpts,
		Route:          route,
		Metrics: MetricsOptions{
			Config: config.MetricsMiddlewareConfiguration{
				QueryEndpointsClassification: config.QueryClassificationConfig{
					ResultsBuckets:  []int{1, 10},
					DurationBuckets: []time.Duration{1 * time.Minute, 15 * time.Minute},
				},
				AddStatusToLatencies: addStatus,
			},
			ParseQueryParams: parseQueryParams,
		},
	}

	// NB: pass expected seriesCount from qs to the test.
	seriesCount := "series_count"
	responseCode := "response_code"
	h := ResponseMetrics(opts).Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		count := r.URL.Query().Get(seriesCount)
		w.Header().Add(headers.FetchedSeriesCount, count)
		code := r.URL.Query().Get(responseCode)
		c, err := strconv.Atoi(code)
		require.NoError(t, err)
		w.WriteHeader(c)
	}))

	route.Path("/api/v1/query_range").Handler(h)
	server := httptest.NewServer(r)
	defer server.Close()

	urls := []string{
		"query=rate(up[20m])&series_count=15&response_code=200&start=1&end=1",
		"query=rate(up[10m])&series_count=15&response_code=200&start=1&end=1",
		"query=rate(up[10m])&series_count=5&response_code=200&start=1&end=1",
		"query=rate(up[20m])&series_count=15&response_code=300&start=1&end=1",
		// NB: this should be large since the end-start + query duration is 15m.
		"query=rate(up[14m])&series_count=15&response_code=200&start=1621458000&end=1621458060",
	}

	for _, url := range urls {
		resp, err := server.Client().Get(server.URL + "/api/v1/query_range?" + url) //nolint: noctx
		require.NoError(t, err)
		require.NoError(t, resp.Body.Close())
	}

	snapshot := scope.Snapshot()
	tallytest.AssertCounterValue(t, 2, snapshot, "request", map[string]string{
		"path":                 "/api/v1/query_range",
		"status":               "200",
		"type":                 "coordinator",
		resultsClassification:  "10",
		durationClassification: "15m0s",
	})

	tallytest.AssertCounterValue(t, 1, snapshot, "request", map[string]string{
		"path":                 "/api/v1/query_range",
		"status":               "300",
		"type":                 "coordinator",
		resultsClassification:  "unclassified",
		durationClassification: "unclassified",
	})

	tallytest.AssertCounterValue(t, 1, snapshot, "request", map[string]string{
		"path":                 "/api/v1/query_range",
		"status":               "200",
		"type":                 "coordinator",
		resultsClassification:  "1",
		durationClassification: "1m0s",
	})

	tallytest.AssertCounterValue(t, 4, snapshot, "count", map[string]string{
		"type":   "coordinator",
		"status": "classified_duration",
	})
	tallytest.AssertCounterValue(t, 4, snapshot, "count", map[string]string{
		"type":   "coordinator",
		"status": "classified_result",
	})

	var (
		hist       = snapshot.Histograms()
		exHistLen  = 4
		exTagCount = 4
	)

	if addStatus {
		// NB: if status is added, we expect to see three histogram entries,
		// since they will also include the status code in the tags list; otherwise
		// code:200 and code:300 queries are expected to go to the same histogram
		// metric.
		exTagCount++
	}

	require.Equal(t, exHistLen, len(hist))
	buckets := map[string]map[string]int{
		resultsClassification: {
			"1":  0,
			"10": 0,
		},
		durationClassification: {
			"1m0s":  0,
			"15m0s": 0,
		},
	}
	statuses := map[string]int{}
	for _, h := range hist {
		require.Equal(t, "latency", h.Name())

		tags := h.Tags()
		require.Equal(t, exTagCount, len(tags))
		require.Equal(t, "/api/v1/query_range", tags["path"])
		require.Equal(t, metricsTypeTagDefaultValue, tags[metricsTypeTagName])

		buckets[resultsClassification][tags[resultsClassification]]++
		buckets[durationClassification][tags[durationClassification]]++
		if addStatus {
			statuses[tags["status"]]++
		}
	}

	expectedBuckets := map[string]map[string]int{
		resultsClassification: {
			"1":            1,
			"10":           2,
			"unclassified": 1,
		},
		durationClassification: {
			"1m0s":         2,
			"15m0s":        1,
			"unclassified": 1,
		},
	}
	if addStatus {
		require.Equal(t, map[string]int{"200": 3, "300": 1}, statuses)
	}

	require.Equal(t, expectedBuckets, buckets)
}

func TestRequestClassificationByEndpoints(t *testing.T) {
	defaultConfig := config.MetricsMiddlewareConfiguration{
		QueryEndpointsClassification: config.QueryClassificationConfig{
			ResultsBuckets:  []int{1, 10, 100, 1000},
			DurationBuckets: []time.Duration{1 * time.Minute, 10 * time.Minute, 100 * time.Minute},
		},
		LabelEndpointsClassification: config.QueryClassificationConfig{
			ResultsBuckets:  []int{2, 20, 200, 2000},
			DurationBuckets: []time.Duration{2 * time.Minute, 20 * time.Minute, 200 * time.Minute},
		},
	}
	tests := []struct {
		name             string
		path             string
		config           config.MetricsMiddlewareConfiguration
		isQueryEndpoint  bool
		query            string
		fetchedResult    string
		expectedResult   string
		expectedDuration string
	}{
		{
			name:             "query_range",
			path:             route.QueryRangeURL,
			config:           defaultConfig,
			isQueryEndpoint:  true,
			query:            "query=sum(rate(coordinator_http_handler_http_handler_request[1m]))&start=1&end=660",
			fetchedResult:    "25",
			expectedResult:   "10",
			expectedDuration: "10m0s",
		},
		{
			name:             "query_range - no results",
			path:             route.QueryRangeURL,
			config:           defaultConfig,
			isQueryEndpoint:  true,
			query:            "query=sum(rate(coordinator_http_handler_http_handler_request[1m]))&start=1&end=660",
			fetchedResult:    "",
			expectedResult:   "1",
			expectedDuration: "10m0s",
		},
		{
			name:             "query",
			path:             route.QueryURL,
			config:           defaultConfig,
			isQueryEndpoint:  true,
			query:            "query=sum(rate(coordinator_http_handler_http_handler_request[1m]))&time=1630611461",
			fetchedResult:    "1",
			expectedResult:   "1",
			expectedDuration: "1m0s",
		},
		{
			name:             "label_values",
			path:             route.Prefix + "/label/__name__/values",
			config:           defaultConfig,
			isQueryEndpoint:  false,
			query:            "start=0&end=1800",
			fetchedResult:    "30000",
			expectedResult:   "2000",
			expectedDuration: "20m0s",
		},
		{
			name:             "label_values -- max time",
			path:             route.Prefix + "/label/__name__/values",
			config:           defaultConfig,
			isQueryEndpoint:  false,
			fetchedResult:    "30000",
			expectedResult:   "2000",
			expectedDuration: "3h20m0s",
		},
		{
			name:             "label_names",
			path:             route.LabelNamesURL,
			config:           defaultConfig,
			isQueryEndpoint:  false,
			query:            "start=0&end=21600",
			fetchedResult:    "300",
			expectedResult:   "200",
			expectedDuration: "3h20m0s",
		},
		{
			name:             "label_names - no results",
			path:             route.LabelNamesURL,
			config:           defaultConfig,
			isQueryEndpoint:  false,
			query:            "start=0&end=21600",
			fetchedResult:    "",
			expectedResult:   "2",
			expectedDuration: "3h20m0s",
		},
		{
			name:             "non-classifiable endpoint",
			path:             "/foo",
			config:           defaultConfig,
			isQueryEndpoint:  false,
			query:            "start=0&end=21600",
			expectedResult:   "unclassified",
			expectedDuration: "unclassified",
		},
		{
			name:            "disabled",
			path:            "/foo",
			config:          config.MetricsMiddlewareConfiguration{},
			isQueryEndpoint: false,
			query:           "start=0&end=21600",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			scope := tally.NewTestScope("", nil)
			iOpts := instrument.NewOptions().SetMetricsScope(scope)

			r := mux.NewRouter()
			route := r.NewRoute()
			opts := Options{
				InstrumentOpts: iOpts,
				Route:          route,
				Metrics: MetricsOptions{
					Config:           tt.config,
					ParseQueryParams: parseQueryParams,
					ParseOptions: promql.NewParseOptions().
						SetRequireStartEndTime(false).
						SetNowFn(time.Now),
				},
			}

			h := ResponseMetrics(opts).Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if tt.isQueryEndpoint {
					w.Header().Add(headers.FetchedSeriesCount, tt.fetchedResult)
				} else {
					w.Header().Add(headers.FetchedMetadataCount, tt.fetchedResult)
				}
				w.WriteHeader(200)
			}))

			route.Path(tt.path).Handler(h)
			server := httptest.NewServer(r)
			defer server.Close()

			resp, err := server.Client().Get(server.URL + tt.path + "?" + tt.query) //nolint: noctx
			require.NoError(t, err)
			require.NoError(t, resp.Body.Close())

			snapshot := scope.Snapshot()
			// Validate request counter
			tags := map[string]string{
				"path":                 tt.path,
				"status":               "200",
				"type":                 "coordinator",
				resultsClassification:  tt.expectedResult,
				durationClassification: tt.expectedDuration,
			}
			if tt.expectedResult == "" {
				delete(tags, resultsClassification)
			}
			if tt.expectedDuration == "" {
				delete(tags, durationClassification)
			}
			tallytest.AssertCounterValue(t, 1, snapshot, "request", tags)

			// Validate latency histogram
			require.Equal(t, 1, len(snapshot.Histograms()))
			var hist tally.HistogramSnapshot
			for _, hist = range snapshot.Histograms() {
			}

			if !tt.config.AddStatusToLatencies {
				delete(tags, "status")
			}
			require.Equal(t, tags, hist.Tags())

			count := int64(0)
			for _, value := range hist.Durations() {
				count += value
			}
			require.Equal(t, int64(1), count)
		})
	}
}


<|start_filename|>src/cluster/placement/algo/sharded_prop_test.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package algo

import (
	"fmt"
	"math"
	"os"
	"testing"
	"time"

	"github.com/m3db/m3/src/cluster/placement"

	"github.com/leanovate/gopter"
	"github.com/leanovate/gopter/gen"
	"github.com/leanovate/gopter/prop"
)

const minSuccessfulTests = 100

func TestInitialPlacementIsBalancedPropTest(t *testing.T) {
	var (
		parameters = gopter.DefaultTestParameters()
		seed       = time.Now().UnixNano()
		props      = gopter.NewProperties(parameters)
		reporter   = gopter.NewFormatedReporter(true, 160, os.Stdout)
	)

	parameters.MinSuccessfulTests = minSuccessfulTests
	parameters.Rng.Seed(seed)

	props.Property("Initial placement is balanced", prop.ForAll(
		testInitialPlacementIsBalanced,
		gen.IntRange(1, 3),
		gen.IntRange(1, 20),
		gen.IntRange(64, 3072),
	))

	if !props.Run(reporter) {
		t.Errorf("failed with initial seed: %d", seed)
	}
}

func testInitialPlacementIsBalanced(replicaCount, instanceCount, shardCount int) (bool, error) {
	instances := make([]placement.Instance, 0)
	for i := 0; i < replicaCount; i++ {
		for j := 0; j < instanceCount; j++ {
			var (
				instanceID     = fmt.Sprintf("instance-%d-%03d", i, j)
				isolationGroup = fmt.Sprintf("iso-%d", i)
				instance       = placement.NewEmptyInstance(instanceID, isolationGroup, "zone", "endpoint", 1)
			)
			instances = append(instances, instance)
		}
	}

	shardIDs := make([]uint32, shardCount)
	for i := range shardIDs {
		shardIDs[i] = uint32(i)
	}

	algo := newShardedAlgorithm(placement.NewOptions())
	p, err := algo.InitialPlacement(instances, shardIDs, replicaCount)
	if err != nil {
		return false, err
	}

	for _, shardID := range shardIDs {
		if n := len(p.InstancesForShard(shardID)); n != replicaCount {
			return false, fmt.Errorf("shard %d has %d replicas, but replication factor is %d", shardID, n, replicaCount)
		}
	}

	var (
		min = math.MaxInt32
		max = math.MinInt32
	)
	for _, instance := range p.Instances() {
		n := instance.Shards().NumShards()
		if n < min {
			min = n
		}
		if n > max {
			max = n
		}
	}
	if max-min > 1 {
		return false, fmt.Errorf("shard count differs by more than 1, min=%v max=%v", min, max)
	}
	return true, nil
}


<|start_filename|>src/dbnode/encoding/m3tsz/decoder_benchmark_test.go<|end_filename|>
// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package m3tsz

import (
	"encoding/base64"
	"math/rand"
	"testing"

	"github.com/m3db/m3/src/dbnode/encoding"
	"github.com/m3db/m3/src/dbnode/x/xio"

	"github.com/stretchr/testify/require"
)

// BenchmarkM3TSZDecode-12    	   16867	     69272 ns/op
func BenchmarkM3TSZDecode(b *testing.B) {
	var (
		encodingOpts = encoding.NewOptions()
		reader       = xio.NewBytesReader64(nil)
		seriesRun    = prepareSampleSeriesRun(b)
	)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		reader.Reset(seriesRun[i])
		iter := NewReaderIterator(reader, DefaultIntOptimizationEnabled, encodingOpts)
		for iter.Next() {
			_, _, _ = iter.Current()
		}
		require.NoError(b, iter.Err())
	}
}

func prepareSampleSeriesRun(b *testing.B) [][]byte {
	b.Helper()
	var (
		rnd          = rand.New(rand.NewSource(42)) // nolint: gosec
		sampleSeries = make([][]byte, 0, len(sampleSeriesBase64))
		seriesRun    = make([][]byte, 0, b.N)
	)

	for _, b64 := range sampleSeriesBase64 {
		data, err := base64.StdEncoding.DecodeString(b64)
		require.NoError(b, err)
		sampleSeries = append(sampleSeries, data)
	}

	for i := 0; i < b.N; i++ {
		seriesRun = append(seriesRun, sampleSeries[rnd.Intn(len(sampleSeries))])
	}

	return seriesRun
}


<|start_filename|>docker/m3dbnode/Dockerfile<|end_filename|>
# stage 1: build
FROM golang:1.18-alpine3.15 AS builder
LABEL maintainer="The M3DB Authors <<EMAIL>>"

# Install deps
RUN apk add --update git make bash

# Add source code
RUN mkdir -p /go/src/github.com/m3db/m3
ADD . /go/src/github.com/m3db/m3

# Build m3dbnode binary
RUN cd /go/src/github.com/m3db/m3/ && \
    git submodule update --init      && \
    make m3dbnode-linux-amd64

# Stage 2: lightweight "release"
FROM alpine:3.15
LABEL maintainer="The M3DB Authors <<EMAIL>>"

ENV GODEBUG madvdontneed=1

EXPOSE 2379/tcp 2380/tcp 7201/tcp 7203/tcp 9000-9004/tcp

# Provide timezone data to allow TZ environment variable to be set
# for parsing relative times such as "9am" correctly and respect
# the TZ environment variable.
RUN apk add --no-cache tzdata curl jq

COPY --from=builder /go/src/github.com/m3db/m3/src/dbnode/config/m3dbnode-local-etcd.yml /etc/m3dbnode/m3dbnode.yml
COPY --from=builder /go/src/github.com/m3db/m3/bin/m3dbnode \
  /go/src/github.com/m3db/m3/scripts/m3dbnode_bootstrapped.sh \
  /bin/

ENV GODEBUG madvdontneed=1

ENTRYPOINT [ "/bin/m3dbnode" ]
CMD [ "-f", "/etc/m3dbnode/m3dbnode.yml" ]


<|start_filename|>src/query/api/v1/middleware/middleware.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// Package middleware contains HTTP middleware functions.
package middleware

import (
	"net/http"

	"github.com/m3db/m3/src/x/instrument"
	"github.com/m3db/m3/src/x/net/http/cors"

	"github.com/gorilla/mux"
	"github.com/jonboulle/clockwork"
	"github.com/opentracing/opentracing-go"
	"github.com/prometheus/prometheus/util/httputil"
)

// Register is a func to build the set of middleware functions.
type Register func(opts Options) []mux.MiddlewareFunc

// Options is the set of parameters passed to the Register function.
type Options struct {
	// Common options for all middleware functions.
	InstrumentOpts instrument.Options
	Clock          clockwork.Clock
	Route          *mux.Route

	// Specific options for middleware functions.
	Logging                LoggingOptions
	Metrics                MetricsOptions
	Source                 SourceOptions
	PrometheusRangeRewrite PrometheusRangeRewriteOptions
}

// OverrideOptions is a function that returns new Options from the provided Options.
type OverrideOptions func(opts Options) Options

// Default is the default list of middleware functions applied if no middleware functions are set in the
// HandlerOptions.
func Default(opts Options) []mux.MiddlewareFunc {
	// The order of middleware is important. Be very careful when reordering existing middleware.
	return []mux.MiddlewareFunc{
		Cors(),
		// install tracing before logging so the trace_id is available for response logging.
		Tracing(opentracing.GlobalTracer(), opts.InstrumentOpts),
		// install source before logging so the source is available for response logging.
		Source(opts),
		RequestID(opts.InstrumentOpts),
		PrometheusRangeRewrite(opts),
		ResponseLogging(opts),
		ResponseMetrics(opts),
		// install panic handler after any middleware that adds extra useful information to the context logger.
		Panic(opts.InstrumentOpts),
		Compression(),
	}
}

// Cors adds CORS headers will be added to all responses.
func Cors() mux.MiddlewareFunc {
	return func(base http.Handler) http.Handler {
		return &cors.Handler{
			Handler: base,
			Info: &cors.Info{
				"*": true,
			},
		}
	}
}

// Compression adds suitable response compression based on the client's Accept-Encoding headers.
func Compression() mux.MiddlewareFunc {
	return func(base http.Handler) http.Handler {
		return httputil.CompressionHandler{
			Handler: base,
		}
	}
}


<|start_filename|>src/aggregator/aggregator/elem_test.go<|end_filename|>
// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package aggregator

import (
	"math"
	"sync"
	"testing"
	"time"

	raggregation "github.com/m3db/m3/src/aggregator/aggregation"
	"github.com/m3db/m3/src/aggregator/aggregation/quantile/cm"
	maggregation "github.com/m3db/m3/src/metrics/aggregation"
	"github.com/m3db/m3/src/metrics/metadata"
	"github.com/m3db/m3/src/metrics/metric"
	"github.com/m3db/m3/src/metrics/metric/aggregated"
	"github.com/m3db/m3/src/metrics/metric/id"
	"github.com/m3db/m3/src/metrics/metric/unaggregated"
	"github.com/m3db/m3/src/metrics/pipeline"
	"github.com/m3db/m3/src/metrics/pipeline/applied"
	"github.com/m3db/m3/src/metrics/policy"
	"github.com/m3db/m3/src/metrics/transformation"
	"github.com/m3db/m3/src/x/instrument"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
	"github.com/willf/bitset"
	"go.uber.org/atomic"
)

var (
	testCounterID                 = id.RawID("testCounter")
	testBatchTimerID              = id.RawID("testBatchTimer")
	testGaugeID                   = id.RawID("testGauge")
	testAnnot                     = []byte("testAnnotation")
	testStoragePolicy             = policy.NewStoragePolicy(10*time.Second, xtime.Second, 6*time.Hour)
	testAggregationTypes          = maggregation.Types{maggregation.Mean, maggregation.Sum}
	testAggregationTypesExpensive = maggregation.Types{maggregation.SumSq}
	testTimerAggregationTypes     = maggregation.Types{maggregation.SumSq, maggregation.P99}
	testGaugeElemData             = ElemData{
		ID:                testGaugeID,
		StoragePolicy:     testStoragePolicy,
		Pipeline:          testPipeline,
		NumForwardedTimes: testNumForwardedTimes,
	}
	testCounterElemData = ElemData{
		ID:                testCounterID,
		StoragePolicy:     testStoragePolicy,
		Pipeline:          testPipeline,
		NumForwardedTimes: testNumForwardedTimes,
	}
	testTimerElemData = ElemData{
		ID:                testBatchTimerID,
		StoragePolicy:     testStoragePolicy,
		Pipeline:          testPipeline,
		NumForwardedTimes: testNumForwardedTimes,
	}
	testCounter = unaggregated.MetricUnion{
		Type:       metric.CounterType,
		ID:         testCounterID,
		CounterVal: 1234,
	}
	testBatchTimer = unaggregated.MetricUnion{
		Type:          metric.TimerType,
		ID:            testBatchTimerID,
		BatchTimerVal: []float64{1.0, 3.5, 2.2, 6.5, 4.8},
	}
	testGauge = unaggregated.MetricUnion{
		Type:     metric.GaugeType,
		ID:       testGaugeID,
		GaugeVal: 123.456,
	}
	testPipeline = applied.NewPipeline([]applied.OpUnion{
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.Absolute},
		},
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.PerSecond},
		},
		{
			Type: pipeline.RollupOpType,
			Rollup: applied.RollupOp{
				ID:            []byte("foo.bar"),
				AggregationID: maggregation.MustCompressTypes(maggregation.Count),
			},
		},
		{
			Type: pipeline.RollupOpType,
			Rollup: applied.RollupOp{
				ID:            []byte("foo.baz"),
				AggregationID: maggregation.MustCompressTypes(maggregation.Max),
			},
		},
	})
	testNumForwardedTimes = 0
	testOpts              = newTestOptions()
	testTimestamps        = []time.Time{
		time.Unix(216, 0), time.Unix(217, 0), time.Unix(221, 0),
	}
	testAlignedStarts = []int64{
		time.Unix(210, 0).UnixNano(), time.Unix(220, 0).UnixNano(), time.Unix(230, 0).UnixNano(),
	}
	testCounterVals    = []int64{testCounter.CounterVal, testCounter.CounterVal}
	testBatchTimerVals = [][]float64{testBatchTimer.BatchTimerVal, testBatchTimer.BatchTimerVal}
	testGaugeVals      = []float64{testGauge.GaugeVal, testGauge.GaugeVal}
)

func TestCounterResetSetData(t *testing.T) {
	opts := newTestOptions()
	elemData := testCounterElemData
	elemData.NumForwardedTimes = 1
	ce, err := NewCounterElem(elemData, NewElemOptions(opts))
	require.NoError(t, err)
	require.Equal(t, opts.AggregationTypesOptions().DefaultCounterAggregationTypes(), ce.aggTypes)
	require.True(t, ce.useDefaultAggregation)
	require.False(t, ce.aggOpts.HasExpensiveAggregations)
	require.Equal(t, 1, ce.numForwardedTimes)

	// Reset element with a default pipeline.
	elemData.AggTypes = testAggregationTypesExpensive
	elemData.NumForwardedTimes = 2
	elemData.Pipeline = applied.DefaultPipeline
	err = ce.ResetSetData(elemData)
	require.NoError(t, err)
	require.Equal(t, testCounterID, ce.id)
	require.Equal(t, testStoragePolicy, ce.sp)
	require.Equal(t, testAggregationTypesExpensive, ce.aggTypes)
	require.Equal(t, parsedPipeline{}, ce.parsedPipeline)
	require.False(t, ce.tombstoned)
	require.False(t, ce.closed)
	require.False(t, ce.useDefaultAggregation)
	require.True(t, ce.aggOpts.HasExpensiveAggregations)
	require.Empty(t, ce.flushState)
	require.Equal(t, 2, ce.numForwardedTimes)

	// Reset element with a pipeline containing a derivative transformation.
	expectedParsedPipeline := parsedPipeline{
		HasDerivativeTransform: true,
		Transformations: []transformation.Op{
			mustNewOp(t, transformation.Absolute),
			mustNewOp(t, transformation.PerSecond),
		},
		HasRollup: true,
		Rollup: applied.RollupOp{
			ID:            []byte("foo.bar"),
			AggregationID: maggregation.MustCompressTypes(maggregation.Count),
		},
		Remainder: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo.baz"),
					AggregationID: maggregation.MustCompressTypes(maggregation.Max),
				},
			},
		}),
	}
	elemData.NumForwardedTimes = 0
	elemData.Pipeline = testPipeline
	err = ce.ResetSetData(elemData)
	require.NoError(t, err)
	requirePipelinesMatch(t, expectedParsedPipeline, ce.parsedPipeline)
	require.Empty(t, ce.flushState)
}

func TestCounterResetSetDataInvalidAggregationType(t *testing.T) {
	opts := newTestOptions()
	ce := MustNewCounterElem(testCounterElemData, NewElemOptions(opts))
	elemData := testCounterElemData
	elemData.AggTypes = maggregation.Types{maggregation.Last}
	err := ce.ResetSetData(elemData)
	require.Error(t, err)
}

func TestCounterResetSetDataNoRollup(t *testing.T) {
	opts := newTestOptions()
	ce := MustNewCounterElem(testCounterElemData, NewElemOptions(opts))

	pipelineNoRollup := applied.NewPipeline([]applied.OpUnion{
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.Absolute},
		},
	})
	elemData := testCounterElemData
	elemData.Pipeline = pipelineNoRollup
	err := ce.ResetSetData(elemData)
	require.NoError(t, err)
}

func TestCounterElemAddUnion(t *testing.T) {
	e, err := NewCounterElem(testCounterElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a counter metric.
	testCounter.Annotation = []byte{1}
	require.NoError(t, e.AddUnion(testTimestamps[0], testCounter, false))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, testCounter.CounterVal, v.aggregation.Sum())
	require.Equal(t, int64(1), v.aggregation.Count())
	require.Equal(t, int64(0), v.aggregation.SumSq())
	require.Equal(t, []byte{1}, v.aggregation.Annotation())

	// Add the counter metric at slightly different time
	// but still within the same aggregation interval.
	testCounter.Annotation = []byte{}
	require.NoError(t, e.AddUnion(testTimestamps[1], testCounter, false))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, 2*testCounter.CounterVal, v.aggregation.Sum())
	require.Equal(t, int64(2), v.aggregation.Count())
	require.Equal(t, int64(0), v.aggregation.SumSq())
	require.Equal(t, []byte{1}, v.aggregation.Annotation())

	// Add the counter metric in the next aggregation interval.
	testCounter.Annotation = []byte{2}
	require.NoError(t, e.AddUnion(testTimestamps[2], testCounter, false))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	require.Equal(t, testCounter.CounterVal, v1.aggregation.Sum())
	require.Equal(t, int64(2), v.aggregation.Count())
	require.Equal(t, int64(0), v.aggregation.SumSq())
	require.Equal(t, []byte{1}, v.aggregation.Annotation())

	// Adding the counter metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnion(testTimestamps[2], testCounter, false))
}

func TestCounterElemAddUnionWithCustomAggregation(t *testing.T) {
	elemData := testCounterElemData
	elemData.AggTypes = testAggregationTypesExpensive
	e, err := NewCounterElem(elemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a counter metric.
	testCounter.Annotation = []byte{1}
	require.NoError(t, e.AddUnion(testTimestamps[0], testCounter, false))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, testCounter.CounterVal, v.aggregation.Sum())
	require.Equal(t, testCounter.CounterVal, v.aggregation.Max())
	require.Equal(t, testCounter.CounterVal*testCounter.CounterVal, v.aggregation.SumSq())
	require.Equal(t, []byte{1}, v.aggregation.Annotation())

	// Add the counter metric at slightly different time
	// but still within the same aggregation interval.
	testCounter.Annotation = []byte{}
	require.NoError(t, e.AddUnion(testTimestamps[1], testCounter, false))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, 2*testCounter.CounterVal, v.aggregation.Sum())
	require.Equal(t, testCounter.CounterVal, v.aggregation.Max())
	require.Equal(t, []byte{1}, v.aggregation.Annotation())

	// Add the counter metric in the next aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[2], testCounter, false))
	require.Equal(t, 2, len(e.values))
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	a, err = e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, testCounter.CounterVal, v.aggregation.Sum())
	require.Equal(t, testCounter.CounterVal, v.aggregation.Max())

	// Adding the counter metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnion(testTimestamps[2], testCounter, false))
}

func TestCounterElemAddUnique(t *testing.T) {
	e, err := NewCounterElem(testCounterElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a metric.
	source1 := uint32(1234)
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{345}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, int64(345), v.aggregation.Sum())
	require.Equal(t, int64(1), v.aggregation.Count())
	require.Equal(t, int64(0), v.aggregation.SumSq())
	require.Nil(t, v.aggregation.Annotation())
	versions, seen := v.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add another metric at slightly different time but still within the
	// same aggregation interval with a different source.
	source2 := uint32(5678)
	require.NoError(t, e.AddUnique(testTimestamps[1],
		aggregated.ForwardedMetric{Values: []float64{500}, Annotation: testAnnot},
		metadata.ForwardMetadata{SourceID: source2}))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, int64(845), v.aggregation.Sum())
	require.Equal(t, int64(2), v.aggregation.Count())
	require.Equal(t, int64(0), v.aggregation.SumSq())
	require.Equal(t, testAnnot, v.aggregation.Annotation())
	versions, seen = v.sourcesSeen[source2]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add the counter metric in the next aggregation interval.
	require.NoError(t, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{278}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	require.Equal(t, int64(278), v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.Equal(t, int64(0), v1.aggregation.SumSq())
	versions, seen = v.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))
	require.Equal(t, testAnnot, v.aggregation.Annotation())

	// Add the counter metric in the same aggregation interval with the same
	// source results in an error.
	require.Equal(t, errDuplicateForwardingSource, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{278}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	a1, err = e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 = a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, int64(278), v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.Equal(t, int64(0), v1.aggregation.SumSq())
	versions, seen = v1.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Adding the counter metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{100}},
		metadata.ForwardMetadata{SourceID: 1376}))
}

func TestCounterElemAddUniqueWithCustomAggregation(t *testing.T) {
	elemData := testCounterElemData
	elemData.AggTypes = testAggregationTypesExpensive
	e, err := NewCounterElem(elemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a counter metric.
	source1 := uint32(1234)
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{12}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, int64(12), v.aggregation.Sum())
	require.Equal(t, int64(12), v.aggregation.Max())
	require.Equal(t, int64(144), v.aggregation.SumSq())
	versions, seen := v.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add the counter metric at slightly different time
	// but still within the same aggregation interval.
	source2 := uint32(5678)
	require.NoError(t, e.AddUnique(testTimestamps[1],
		aggregated.ForwardedMetric{Values: []float64{14}},
		metadata.ForwardMetadata{SourceID: source2}))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, int64(26), v.aggregation.Sum())
	require.Equal(t, int64(14), v.aggregation.Max())

	// Add the counter metric in the next aggregation interval.
	require.NoError(t, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{20}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, int64(20), v1.aggregation.Sum())
	require.Equal(t, int64(20), v1.aggregation.Max())
	require.Equal(t, int64(400), v1.aggregation.SumSq())

	// Add the counter metric in the same aggregation interval with the same
	// source results in an error.
	require.Equal(t, errDuplicateForwardingSource, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{30}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	a1, err = e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 = a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, int64(20), v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.Equal(t, int64(400), v1.aggregation.SumSq())
	versions, seen = v1.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Adding the counter metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{40}},
		metadata.ForwardMetadata{SourceID: 1376}))
}

func TestCounterElemConsumeDefaultAggregationDefaultPipeline(t *testing.T) {
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions()
	e := testCounterElem(testAlignedStarts[:len(testAlignedStarts)-1], testCounterVals, maggregation.DefaultTypes, applied.DefaultPipeline, opts)

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForCounter(testAlignedStarts[1], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForCounter(testAlignedStarts[2], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestCounterElemConsumeCustomAggregationDefaultPipeline(t *testing.T) {
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions()
	e := testCounterElem(testAlignedStarts[:len(testAlignedStarts)-1], testCounterVals, testAggregationTypes, applied.DefaultPipeline, opts)

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForCounter(testAlignedStarts[1], testStoragePolicy, testAggregationTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForCounter(testAlignedStarts[2], testStoragePolicy, testAggregationTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestCounterElemConsumeCustomAggregationCustomPipeline(t *testing.T) {
	alignedstartAtNanos := []int64{
		time.Unix(210, 0).UnixNano(),
		time.Unix(220, 0).UnixNano(),
		time.Unix(230, 0).UnixNano(),
		time.Unix(240, 0).UnixNano(),
	}
	counterVals := []int64{-123, -456, -589}
	aggregationTypes := maggregation.Types{maggregation.Sum}
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions().SetDiscardNaNAggregatedValues(false)
	e := testCounterElem(alignedstartAtNanos[:3], counterVals, aggregationTypes, testPipeline, opts)

	aggKey := aggregationKey{
		aggregationID: maggregation.MustCompressTypes(maggregation.Count),
		storagePolicy: testStoragePolicy,
		pipeline: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo.baz"),
					AggregationID: maggregation.MustCompressTypes(maggregation.Max),
				},
			},
		}),
		numForwardedTimes: testNumForwardedTimes + 1,
	}
	expectedOnFlushedRes := []testOnForwardedFlushedData{
		{
			aggregationKey: aggKey,
		},
	}

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()

	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 3, len(e.values))

	// Consume one value.
	expectedForwardedRes := []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(220, 0).UnixNano(),
			value:          nan,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()

	require.False(t, e.Consume(alignedstartAtNanos[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 3, len(e.values))
	require.Len(t, e.flushState, 1)
	fState := e.flushState[xtime.UnixNano(alignedstartAtNanos[0])]
	consumedVal := fState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 123.0, consumedVal[0])

	// Consume all values.
	expectedForwardedRes = []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(230, 0).UnixNano(),
			value:          33.3,
		},
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(240, 0).UnixNano(),
			value:          13.3,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()

	// Test a consume time well past the latest value in the buffer to ensure we always keep around
	// one value to reference in the future as a previous.
	ts := alignedstartAtNanos[3] + int64(time.Second*20)
	require.False(t, e.Consume(ts, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	expectedOnFlushedRes[0].expiredTimes = append(expectedOnFlushedRes[0].expiredTimes,
		xtime.ToUnixNano(time.Unix(220, 0)), xtime.ToUnixNano(time.Unix(230, 0)))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 1, len(e.values))
	require.Len(t, e.flushState, 1)
	fState = e.flushState[xtime.UnixNano(alignedstartAtNanos[2])]
	consumedVal = fState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 589.0, consumedVal[0])

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	expectedOnFlushedRes[0].expiredTimes = nil
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestCounterElemClose(t *testing.T) {
	e := testCounterElem(testAlignedStarts[:len(testAlignedStarts)-1], testCounterVals,
		maggregation.DefaultTypes, applied.DefaultPipeline, newTestOptions())
	require.False(t, e.closed)

	// Closing the element.
	e.Close()

	// Closing a second time should have no impact.
	e.Close()

	require.True(t, e.closed)
	require.Nil(t, e.id)
	require.Equal(t, parsedPipeline{}, e.parsedPipeline)
	require.Nil(t, e.writeForwardedMetricFn)
	require.Nil(t, e.onForwardedAggregationWrittenFn)
	require.Nil(t, e.cachedSourceSets)
	require.Equal(t, 0, len(e.values))
	require.Equal(t, 0, len(e.toConsume))
	require.Equal(t, 0, len(e.flushState))
	require.NotNil(t, e.values)
}

func TestCounterFindOrCreateNoSourceSet(t *testing.T) {
	e, err := NewCounterElem(testCounterElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	inputs := []xtime.UnixNano{10, 10, 20, 10, 15}
	expected := make(map[xtime.UnixNano]bool)
	for _, inputNanos := range inputs {
		expected[inputNanos] = true
		res, err := e.findOrCreate(int64(inputNanos), createAggregationOptions{initSourceSet: false})
		require.NoError(t, err)
		require.Equal(t, e.values[inputNanos].lockedAgg, res)
		require.Nil(t, e.values[inputNanos].lockedAgg.sourcesSeen)
		require.Equal(t, len(expected), len(e.values))
		for k := range e.values {
			_, ok := expected[k]
			require.True(t, ok)
		}
	}
}

// nolint: dupl
func TestCounterFindOrCreateWithSourceSet(t *testing.T) {
	e, err := NewCounterElem(testCounterElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)
	e.cachedSourceSets = make([]map[uint32]*bitset.BitSet, 0)

	inputs := []xtime.UnixNano{10, 20}
	expected := make(map[xtime.UnixNano]bool)
	for _, inputNanos := range inputs {
		expected[inputNanos] = true
		res, err := e.findOrCreate(int64(inputNanos), createAggregationOptions{initSourceSet: true})
		require.NoError(t, err)
		require.Equal(t, e.values[inputNanos].lockedAgg, res)
		require.NotNil(t, e.values[inputNanos].lockedAgg.sourcesSeen)
		require.Equal(t, len(expected), len(e.values))
		for k := range e.values {
			_, ok := expected[k]
			require.True(t, ok)
		}
	}
	require.Equal(t, 0, len(e.cachedSourceSets))
}

func TestTimerResetSetData(t *testing.T) {
	opts := newTestOptions()
	te, err := NewTimerElem(testTimerElemData, NewElemOptions(opts))
	require.NoError(t, err)
	require.Nil(t, te.quantilesPool)
	require.NotNil(t, te.quantiles)
	require.True(t, te.aggOpts.HasExpensiveAggregations)
	require.Equal(t, opts.AggregationTypesOptions().DefaultTimerAggregationTypes(), te.aggTypes)
	require.True(t, te.useDefaultAggregation)

	// Reset element with a default pipeline.
	elemData := testTimerElemData
	elemData.AggTypes = maggregation.Types{maggregation.Max, maggregation.P999}
	elemData.Pipeline = applied.DefaultPipeline
	err = te.ResetSetData(elemData)
	require.NoError(t, err)
	require.Equal(t, testBatchTimerID, te.id)
	require.Equal(t, testStoragePolicy, te.sp)
	require.Equal(t, maggregation.Types{maggregation.Max, maggregation.P999}, te.aggTypes)
	require.Equal(t, parsedPipeline{}, te.parsedPipeline)
	require.False(t, te.tombstoned)
	require.False(t, te.closed)
	require.False(t, te.useDefaultAggregation)
	require.False(t, te.aggOpts.HasExpensiveAggregations)
	require.Equal(t, []float64{0.999}, te.quantiles)
	require.NotNil(t, te.quantilesPool)
	require.Empty(t, te.flushState)

	// Reset element with a pipeline containing a derivative transformation.
	expectedParsedPipeline := parsedPipeline{
		HasDerivativeTransform: true,
		Transformations: []transformation.Op{
			mustNewOp(t, transformation.Absolute),
			mustNewOp(t, transformation.PerSecond),
		},
		HasRollup: true,
		Rollup: applied.RollupOp{
			ID:            []byte("foo.bar"),
			AggregationID: maggregation.MustCompressTypes(maggregation.Count),
		},
		Remainder: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo.baz"),
					AggregationID: maggregation.MustCompressTypes(maggregation.Max),
				},
			},
		}),
	}
	elemData.Pipeline = testPipeline
	err = te.ResetSetData(elemData)
	require.NoError(t, err)
	requirePipelinesMatch(t, expectedParsedPipeline, te.parsedPipeline)
	require.Empty(t, te.flushState)
}

func TestTimerResetSetDataInvalidAggregationType(t *testing.T) {
	opts := newTestOptions()
	te := MustNewTimerElem(testTimerElemData, NewElemOptions(opts))
	elemData := testTimerElemData
	elemData.AggTypes = maggregation.Types{maggregation.Last}
	err := te.ResetSetData(elemData)
	require.Error(t, err)
}

func TestTimerResetSetDataNoRollup(t *testing.T) {
	opts := newTestOptions()
	te := MustNewTimerElem(testTimerElemData, NewElemOptions(opts))

	pipelineNoRollup := applied.NewPipeline([]applied.OpUnion{
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.Absolute},
		},
	})
	elemData := testTimerElemData
	elemData.Pipeline = pipelineNoRollup
	err := te.ResetSetData(elemData)
	require.NoError(t, err)
}

func TestTimerElemAddUnion(t *testing.T) {
	e, err := NewTimerElem(testTimerElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a timer metric.
	require.NoError(t, e.AddUnion(testTimestamps[0], testBatchTimer, false))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	timer := v.aggregation
	require.Equal(t, int64(5), timer.Count())
	require.Equal(t, 18.0, timer.Sum())
	require.Equal(t, 3.5, timer.Quantile(0.5))

	require.Equal(t, 6.5, timer.Quantile(0.95))
	require.Equal(t, 6.5, timer.Quantile(0.99))

	// Add the timer metric at slightly different time
	// but still within the same aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[1], testBatchTimer, false))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	timer = v.aggregation
	require.Equal(t, int64(10), timer.Count())
	require.Equal(t, 36.0, timer.Sum())
	require.Equal(t, 3.5, timer.Quantile(0.5))
	require.Equal(t, 6.5, timer.Quantile(0.95))
	require.Equal(t, 6.5, timer.Quantile(0.99))

	// Add the timer metric in the next aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[2], testBatchTimer, false))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	timer = v1.aggregation
	require.Equal(t, int64(5), timer.Count())
	require.Equal(t, 18.0, timer.Sum())
	require.Equal(t, 3.5, timer.Quantile(0.5))
	require.Equal(t, 6.5, timer.Quantile(0.95))
	require.Equal(t, 6.5, timer.Quantile(0.99))

	// Adding the timer metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnion(testTimestamps[2], testBatchTimer, false))
}

func TestTimerElemAddUnique(t *testing.T) {
	e, err := NewTimerElem(testTimerElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a metric.
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{11.1}},
		metadata.ForwardMetadata{SourceID: 1}))
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{12.2}},
		metadata.ForwardMetadata{SourceID: 2}))
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{13.3}},
		metadata.ForwardMetadata{SourceID: 3}))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	timer := v.aggregation
	require.Equal(t, int64(3), timer.Count())
	require.InEpsilon(t, 36.6, timer.Sum(), 1e-10)
	require.Equal(t, 12.2, timer.Quantile(0.5))

	// Add another metric at slightly different time but still within the
	// same aggregation interval with a different source.
	require.NoError(t, e.AddUnique(testTimestamps[1],
		aggregated.ForwardedMetric{Values: []float64{14.4}},
		metadata.ForwardMetadata{SourceID: 4}))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	timer = v.aggregation
	require.Equal(t, int64(4), timer.Count())
	require.InEpsilon(t, 51, timer.Sum(), 1e-10)

	// Add the metric in the next aggregation interval.
	require.NoError(t, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{20.0}},
		metadata.ForwardMetadata{SourceID: 1}))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, 20.0, v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.Equal(t, 20.0, v1.aggregation.Sum())

	// Add the metric in the same aggregation interval with the same
	// source results in an error.
	require.Equal(t, errDuplicateForwardingSource, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{30.0}},
		metadata.ForwardMetadata{SourceID: 1}))
	require.Equal(t, 2, len(e.values))
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, 20.0, v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.InEpsilon(t, 400.0, v1.aggregation.SumSq(), 1e-10)
	versions, seen := v1.sourcesSeen[1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Adding the timer metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{100}},
		metadata.ForwardMetadata{SourceID: 3}))
}

func TestTimerElemConsumeDefaultAggregationDefaultPipeline(t *testing.T) {
	// Set up stream options.
	streamOpts := cm.NewOptions()
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos

	// Verify the pool is big enough to supply all the streams.
	opts := newTestOptions().SetStreamOptions(streamOpts)
	e := testTimerElem(testAlignedStarts[:len(testAlignedStarts)-1], testBatchTimerVals, maggregation.DefaultTypes, applied.DefaultPipeline, opts)

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForTimer(testAlignedStarts[1], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForTimer(testAlignedStarts[2], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestTimerElemConsumeCustomAggregationDefaultPipeline(t *testing.T) {
	// Set up stream options.
	streamOpts := cm.NewOptions()
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos

	// Verify the pool is big enough to supply all the streams.
	opts := newTestOptions().SetStreamOptions(streamOpts)
	e := testTimerElem(testAlignedStarts[:len(testAlignedStarts)-1], testBatchTimerVals, testTimerAggregationTypes, applied.DefaultPipeline, opts)

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()

	require.False(t, e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForTimer(testAlignedStarts[1], testStoragePolicy, testTimerAggregationTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForTimer(testAlignedStarts[2], testStoragePolicy, testTimerAggregationTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestTimerElemConsumeCustomAggregationCustomPipeline(t *testing.T) {
	alignedstartAtNanos := []int64{
		time.Unix(210, 0).UnixNano(),
		time.Unix(220, 0).UnixNano(),
		time.Unix(230, 0).UnixNano(),
		time.Unix(240, 0).UnixNano(),
	}
	timerVals := [][]float64{
		{123, 1245},
		{456},
		{589, 1120},
	}
	aggregationTypes := maggregation.Types{maggregation.Min}
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions().SetDiscardNaNAggregatedValues(false)
	e := testTimerElem(alignedstartAtNanos[:3], timerVals, aggregationTypes, testPipeline, opts)

	aggKey := aggregationKey{
		aggregationID: maggregation.MustCompressTypes(maggregation.Count),
		storagePolicy: testStoragePolicy,
		pipeline: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo.baz"),
					AggregationID: maggregation.MustCompressTypes(maggregation.Max),
				},
			},
		}),
		numForwardedTimes: testNumForwardedTimes + 1,
	}
	expectedOnFlushedRes := []testOnForwardedFlushedData{
		{
			aggregationKey: aggKey,
		},
	}

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 3, len(e.values))

	// Consume one value.
	expectedForwardedRes := []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(220, 0).UnixNano(),
			value:          nan,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 3, len(e.values))
	require.Len(t, e.flushState, 1)
	fState := e.flushState[xtime.UnixNano(alignedstartAtNanos[0])]
	consumedVal := fState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 123.0, consumedVal[0])

	// Consume all values.
	expectedForwardedRes = []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(230, 0).UnixNano(),
			value:          33.3,
		},
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(240, 0).UnixNano(),
			value:          13.3,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	expectedOnFlushedRes[0].expiredTimes = append(expectedOnFlushedRes[0].expiredTimes,
		xtime.ToUnixNano(time.Unix(220, 0)), xtime.ToUnixNano(time.Unix(230, 0)))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 1, len(e.values))
	require.Len(t, e.flushState, 1)
	fState = e.flushState[xtime.UnixNano(alignedstartAtNanos[2])]
	consumedVal = fState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 589.0, consumedVal[0])

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	expectedOnFlushedRes[0].expiredTimes = nil
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestTimerElemClose(t *testing.T) {
	// Set up stream options.
	streamOpts := cm.NewOptions()

	// Verify the pool is big enough to supply all the streams.
	opts := newTestOptions().SetStreamOptions(streamOpts)
	e := testTimerElem(testAlignedStarts[:len(testAlignedStarts)-1], testBatchTimerVals, maggregation.DefaultTypes, applied.DefaultPipeline, opts)

	require.False(t, e.closed)

	// Closing the element.
	e.Close()

	// Closing a second time should have no impact.
	e.Close()

	require.True(t, e.closed)
	require.Nil(t, e.id)
	require.Equal(t, parsedPipeline{}, e.parsedPipeline)
	require.Nil(t, e.writeForwardedMetricFn)
	require.Nil(t, e.onForwardedAggregationWrittenFn)
	require.Nil(t, e.cachedSourceSets)
	require.Equal(t, 0, len(e.values))
	require.Equal(t, 0, len(e.toConsume))
	require.Equal(t, 0, len(e.flushState))
	require.NotNil(t, e.values)
}

func TestTimerFindOrCreateNoSourceSet(t *testing.T) {
	e, err := NewTimerElem(testTimerElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	inputs := []xtime.UnixNano{10, 10, 20, 10, 15}
	expected := make(map[xtime.UnixNano]bool)
	for _, inputNanos := range inputs {
		expected[inputNanos] = true
		res, err := e.findOrCreate(int64(inputNanos), createAggregationOptions{initSourceSet: false})
		require.NoError(t, err)
		require.Equal(t, e.values[inputNanos].lockedAgg, res)
	}
}

func TestTimerFindOrCreateWithSourceSet(t *testing.T) {
	e, err := NewTimerElem(testTimerElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)
	e.cachedSourceSets = make([]map[uint32]*bitset.BitSet, 0)

	inputs := []xtime.UnixNano{10, 20}
	expected := make(map[xtime.UnixNano]bool)
	for _, inputNanos := range inputs {
		expected[inputNanos] = true
		res, err := e.findOrCreate(int64(inputNanos), createAggregationOptions{initSourceSet: true})
		require.NoError(t, err)
		require.Equal(t, e.values[inputNanos].lockedAgg, res)
		require.NotNil(t, e.values[inputNanos].lockedAgg.sourcesSeen)
		require.Equal(t, len(expected), len(e.values))
		for k := range e.values {
			_, ok := expected[k]
			require.True(t, ok)
		}
	}
	require.Equal(t, 0, len(e.cachedSourceSets))
}

func TestGaugeResetSetData(t *testing.T) {
	opts := newTestOptions()
	ge, err := NewGaugeElem(testGaugeElemData, NewElemOptions(opts))
	require.NoError(t, err)
	require.Equal(t, opts.AggregationTypesOptions().DefaultGaugeAggregationTypes(), ge.aggTypes)
	require.True(t, ge.useDefaultAggregation)
	require.False(t, ge.aggOpts.HasExpensiveAggregations)

	// Reset element with a default pipeline.
	elemData := testGaugeElemData
	elemData.Pipeline = applied.DefaultPipeline
	elemData.AggTypes = testAggregationTypesExpensive
	err = ge.ResetSetData(elemData)
	require.NoError(t, err)
	require.Equal(t, testGaugeID, ge.id)
	require.Equal(t, testStoragePolicy, ge.sp)
	require.Equal(t, testAggregationTypesExpensive, ge.aggTypes)
	require.Equal(t, parsedPipeline{}, ge.parsedPipeline)
	require.False(t, ge.tombstoned)
	require.False(t, ge.closed)
	require.False(t, ge.useDefaultAggregation)
	require.True(t, ge.aggOpts.HasExpensiveAggregations)
	require.Empty(t, ge.flushState)

	// Reset element with a pipeline containing a derivative transformation.
	expectedParsedPipeline := parsedPipeline{
		HasDerivativeTransform: true,
		Transformations: []transformation.Op{
			mustNewOp(t, transformation.Absolute),
			mustNewOp(t, transformation.PerSecond),
		},
		HasRollup: true,
		Rollup: applied.RollupOp{
			ID:            []byte("foo.bar"),
			AggregationID: maggregation.MustCompressTypes(maggregation.Count),
		},
		Remainder: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo.baz"),
					AggregationID: maggregation.MustCompressTypes(maggregation.Max),
				},
			},
		}),
	}
	elemData.Pipeline = testPipeline
	err = ge.ResetSetData(elemData)
	require.NoError(t, err)
	requirePipelinesMatch(t, expectedParsedPipeline, ge.parsedPipeline)
	require.Empty(t, ge.flushState)
}

func TestGaugeElemAddUnion(t *testing.T) {
	e, err := NewGaugeElem(testGaugeElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a gauge metric.
	require.NoError(t, e.AddUnion(testTimestamps[0], testGauge, false))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Last())
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Sum())
	require.Equal(t, 0.0, v.aggregation.SumSq())

	// Add the gauge metric at slightly different time
	// but still within the same aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[1], testGauge, false))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Last())
	require.Equal(t, 2*testGauge.GaugeVal, v.aggregation.Sum())
	require.Equal(t, 0.0, v.aggregation.SumSq())

	// Add the gauge metric in the next aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[2], testGauge, false))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, testGauge.GaugeVal, v1.aggregation.Last())
	require.Equal(t, testGauge.GaugeVal, v1.aggregation.Sum())
	require.Equal(t, 0.0, v1.aggregation.SumSq())

	// Adding the gauge metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnion(testTimestamps[2], testGauge, false))
}

func TestGaugeElemAddUnionWithCustomAggregation(t *testing.T) {
	elemData := testGaugeElemData
	elemData.AggTypes = testAggregationTypesExpensive
	e, err := NewGaugeElem(elemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a gauge metric.
	require.NoError(t, e.AddUnion(testTimestamps[0], testGauge, false))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Last())
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Sum())
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Mean())
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Sum())
	require.Equal(t, testGauge.GaugeVal*testGauge.GaugeVal, v.aggregation.SumSq())

	// Add the gauge metric at slightly different time
	// but still within the same aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[1], testGauge, false))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Last())
	require.Equal(t, testGauge.GaugeVal, v.aggregation.Max())
	require.Equal(t, 2*testGauge.GaugeVal, v.aggregation.Sum())
	require.Equal(t, 2*testGauge.GaugeVal*testGauge.GaugeVal, v.aggregation.SumSq())

	// Add the gauge metric in the next aggregation interval.
	require.NoError(t, e.AddUnion(testTimestamps[2], testGauge, false))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, testGauge.GaugeVal, v1.aggregation.Last())
	require.Equal(t, testGauge.GaugeVal, v1.aggregation.Max())

	// Adding the gauge metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnion(testTimestamps[2], testGauge, false))
}

func TestGaugeElemAddUnique(t *testing.T) {
	e, err := NewGaugeElem(testGaugeElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a metric.
	source1 := uint32(1234)
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{12.3, 34.5}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, 46.8, v.aggregation.Sum())
	require.Equal(t, int64(2), v.aggregation.Count())
	require.Equal(t, 0.0, v.aggregation.SumSq())
	require.Nil(t, v.aggregation.Annotation())
	versions, seen := v.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add another metric at slightly different time but still within the
	// same aggregation interval with a different source.
	source2 := uint32(5678)
	require.NoError(t, e.AddUnique(testTimestamps[1],
		aggregated.ForwardedMetric{Values: []float64{50}, Annotation: testAnnot},
		metadata.ForwardMetadata{SourceID: source2}))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.Equal(t, 96.8, v.aggregation.Sum())
	require.Equal(t, int64(3), v.aggregation.Count())
	require.Equal(t, 0.0, v.aggregation.SumSq())
	require.Equal(t, testAnnot, v.aggregation.Annotation())
	versions, seen = v.sourcesSeen[source2]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add the metric in the next aggregation interval.
	require.NoError(t, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{27.8}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, testAnnot, v.aggregation.Annotation())
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, 27.8, v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.Equal(t, 0.0, v1.aggregation.SumSq())
	versions, seen = v1.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add the gauge metric in the same aggregation interval with the same
	// source results in an error.
	require.Equal(t, errDuplicateForwardingSource, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{27.8}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, 27.8, v1.aggregation.Sum())
	require.Equal(t, int64(1), v1.aggregation.Count())
	require.Equal(t, 0.0, v1.aggregation.SumSq())
	versions, seen = v1.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Update the value from the previous aggregation interval
	// ensure out of order versions are handled correctly.
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{11.0, 33.0}, PrevValues: []float64{12.0, 34.0}, Version: 2},
		metadata.ForwardMetadata{SourceID: source1}))
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{12.0, 34.0}, PrevValues: []float64{12.3, 34.5}, Version: 1},
		metadata.ForwardMetadata{SourceID: source1}))
	// sum changes from 96.8 -> 94.0
	require.Equal(t, 94.0, v.aggregation.Sum())
	require.Equal(t, int64(3), v.aggregation.Count())
	require.Equal(t, 0.0, v.aggregation.SumSq())
	require.Equal(t, testAnnot, v.aggregation.Annotation())
	versions, seen = v.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(1))
	require.True(t, versions.Test(2))

	// Adding the gauge metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{10.0}},
		metadata.ForwardMetadata{SourceID: 3}))
}

func TestGaugeElemAddUniqueWithCustomAggregation(t *testing.T) {
	elemData := testGaugeElemData
	elemData.AggTypes = testAggregationTypesExpensive
	e, err := NewGaugeElem(elemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	// Add a gauge metric.
	source1 := uint32(1234)
	require.NoError(t, e.AddUnique(testTimestamps[0],
		aggregated.ForwardedMetric{Values: []float64{1.2}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 1, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	require.Equal(t, 1.2, v.aggregation.Sum())
	require.Equal(t, 1.2, v.aggregation.Max())
	require.Equal(t, 1.44, v.aggregation.SumSq())
	versions, seen := v.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Add the gauge metric at slightly different time
	// but still within the same aggregation interval.
	source2 := uint32(5678)
	require.NoError(t, e.AddUnique(testTimestamps[1],
		aggregated.ForwardedMetric{Values: []float64{1.4}},
		metadata.ForwardMetadata{SourceID: source2}))
	require.Equal(t, 1, len(e.values))
	a, err = e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v = a.lockedAgg
	require.InEpsilon(t, 2.6, v.aggregation.Sum(), 1e-10)
	require.Equal(t, 1.4, v.aggregation.Max())

	// Add the gauge metric in the next aggregation interval.
	require.NoError(t, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{2.0}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	require.NoError(t, err)
	v1 := a1.lockedAgg
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, 2.0, v1.aggregation.Sum())
	require.Equal(t, 2.0, v1.aggregation.Max())
	require.Equal(t, 4.0, v1.aggregation.SumSq())

	// Add the gauge metric in the same aggregation interval with the same
	// source results in an error.
	require.Equal(t, errDuplicateForwardingSource, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{3.0}},
		metadata.ForwardMetadata{SourceID: source1}))
	require.Equal(t, 2, len(e.values))
	for i := 0; i < len(e.values); i++ {
		require.Equal(t, xtime.UnixNano(testAlignedStarts[i]), e.values[e.dirty[i]].startAt)
	}
	require.Equal(t, 2.0, v1.aggregation.Sum())
	require.Equal(t, 2.0, v1.aggregation.Max())
	require.Equal(t, 4.0, v1.aggregation.SumSq())
	versions, seen = v1.sourcesSeen[source1]
	require.True(t, seen)
	require.True(t, versions.Test(0))

	// Adding the gauge metric to a closed element results in an error.
	e.closed = true
	require.Equal(t, errElemClosed, e.AddUnique(testTimestamps[2],
		aggregated.ForwardedMetric{Values: []float64{4.0}},
		metadata.ForwardMetadata{SourceID: 3}))
}

func TestGaugeElemConsumeDefaultAggregationDefaultPipeline(t *testing.T) {
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions()
	e := testGaugeElem(testAlignedStarts[:len(testAlignedStarts)-1], testGaugeVals, maggregation.DefaultTypes, applied.DefaultPipeline, opts)

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForGauge(testAlignedStarts[1], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForGauge(testAlignedStarts[2], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestGaugeElemConsumeResendBuffer(t *testing.T) {
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions()
	pastBuffer := time.Second * 20
	opts = opts.
		SetInstrumentOptions(instrument.NewOptions()).
		SetBufferForPastTimedMetricFn(func(resolution time.Duration) time.Duration {
			return resolution + pastBuffer
		})
	pipeline := applied.NewPipeline([]applied.OpUnion{{
		Type: pipeline.TransformationOpType,
		Transformation: pipeline.TransformationOp{
			Type: transformation.Increase,
		},
	}})

	elemData := testGaugeData
	elemData.Pipeline = pipeline
	e := testGaugeElemWithData(t, testAlignedStarts[:len(testAlignedStarts)-1], testGaugeVals, elemData, opts,
		true)

	// Consume one value.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t,
		expectedLocalMetricsForGauge(testAlignedStarts[1], testStoragePolicy, maggregation.DefaultTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	v1 := a1.lockedAgg
	require.NoError(t, err)
	require.False(t, v.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[0])].flushed)
	require.True(t, v1.dirty)
	require.False(t, e.flushState[xtime.UnixNano(testAlignedStarts[1])].flushed)

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t,
		expectedLocalMetricsForGaugeWithVal(testAlignedStarts[2], 0.0, testStoragePolicy,
			maggregation.DefaultTypes),
		*localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))
	require.False(t, v.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[0])].flushed)
	require.False(t, v1.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[1])].flushed)

	// Update the first value after flushing
	updatedVal := testGaugeVals[0] - 1.0
	mu := unaggregated.MetricUnion{
		GaugeVal: updatedVal,
	}
	require.NoError(t, e.AddUnion(time.Unix(0, testAlignedStarts[0]), mu, true))
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	// the update cascades, resulting in 2 flushed values.
	exp := expectedLocalMetricsForGaugeWithVal(testAlignedStarts[1], updatedVal, testStoragePolicy,
		maggregation.DefaultTypes)
	exp = append(exp,
		expectedLocalMetricsForGaugeWithVal(testAlignedStarts[2], 1.0, testStoragePolicy,
			maggregation.DefaultTypes)...)
	require.Equal(t, exp, *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))
	require.False(t, v.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[0])].flushed)
	require.False(t, v1.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[1])].flushed)

	// expire the values past the buffer.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	pastBufferTime := pastBuffer + testStoragePolicy.Resolution().Window
	require.False(t, e.Consume(time.Unix(0, testAlignedStarts[1]).Add(pastBufferTime).UnixNano(),
		isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(time.Unix(0, testAlignedStarts[2]).Add(pastBufferTime).UnixNano(),
		isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestGaugeElemConsumeCustomAggregationDefaultPipeline(t *testing.T) {
	opts := newTestOptions()
	e := testGaugeElem(testAlignedStarts[:len(testAlignedStarts)-1], testGaugeVals, testAggregationTypes, applied.DefaultPipeline, opts)
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForGauge(testAlignedStarts[1], testStoragePolicy, testAggregationTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, expectedLocalMetricsForGauge(testAlignedStarts[2], testStoragePolicy, testAggregationTypes), *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(testAlignedStarts[2], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestGaugeElemConsumeCustomAggregationCustomPipeline(t *testing.T) {
	alignedstartAtNanos := []int64{
		time.Unix(210, 0).UnixNano(),
		time.Unix(220, 0).UnixNano(),
		time.Unix(230, 0).UnixNano(),
		time.Unix(240, 0).UnixNano(),
	}
	gaugeVals := []float64{-123.0, -456.0, -589.0}
	aggregationTypes := maggregation.Types{maggregation.Last}
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions().SetDiscardNaNAggregatedValues(false)
	e := testGaugeElem(alignedstartAtNanos[:3], gaugeVals, aggregationTypes, testPipeline, opts)

	aggKey := aggregationKey{
		aggregationID: maggregation.MustCompressTypes(maggregation.Count),
		storagePolicy: testStoragePolicy,
		pipeline: applied.NewPipeline([]applied.OpUnion{
			{
				Type: pipeline.RollupOpType,
				Rollup: applied.RollupOp{
					ID:            []byte("foo.baz"),
					AggregationID: maggregation.MustCompressTypes(maggregation.Max),
				},
			},
		}),
		numForwardedTimes: testNumForwardedTimes + 1,
	}
	expectedOnFlushedRes := []testOnForwardedFlushedData{
		{
			aggregationKey: aggKey,
		},
	}

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 3, len(e.values))

	// Consume one value.
	expectedForwardedRes := []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(220, 0).UnixNano(),
			value:          nan,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 3, len(e.values))
	require.Len(t, e.flushState, 1)
	flushedState := e.flushState[xtime.UnixNano(alignedstartAtNanos[0])]
	consumedVal := flushedState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 123.0, consumedVal[0])

	// Consume all values.
	expectedForwardedRes = []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(230, 0).UnixNano(),
			value:          33.3,
		},
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(240, 0).UnixNano(),
			value:          13.3,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	expectedOnFlushedRes[0].expiredTimes = append(expectedOnFlushedRes[0].expiredTimes,
		xtime.ToUnixNano(time.Unix(220, 0)), xtime.ToUnixNano(time.Unix(230, 0)))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 1, len(e.values))
	require.Len(t, e.flushState, 1)
	flushedState = e.flushState[xtime.UnixNano(alignedstartAtNanos[2])]
	consumedVal = flushedState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 589.0, consumedVal[0])

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.True(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	expectedOnFlushedRes[0].expiredTimes = nil
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestGaugeElemResendSumReset(t *testing.T) {
	alignedstartAtNanos := []int64{
		time.Unix(210, 0).UnixNano(),
		time.Unix(220, 0).UnixNano(),
	}
	gaugeVals := []float64{123.0, 246.0}
	aggregationTypes := maggregation.Types{maggregation.Sum}
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	pastBuffer := time.Second * 30
	opts := newTestOptions().
		SetBufferForPastTimedMetricFn(func(resolution time.Duration) time.Duration {
			return resolution + pastBuffer
		})
	pipe := applied.NewPipeline([]applied.OpUnion{
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.Reset},
		},
	})
	data := testGaugeData
	data.Pipeline = pipe
	data.AggTypes = aggregationTypes
	e := testGaugeElemWithData(t, alignedstartAtNanos, gaugeVals, data, opts, true)
	e.numForwardedTimes = 1

	// Consume one value.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(alignedstartAtNanos[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	exp := expectedLocalMetricsForGaugeWithVal(
		alignedstartAtNanos[1], 123.0, testStoragePolicy, maggregation.DefaultTypes)
	exp = append(exp, expectedLocalMetricsForGaugeWithVal(
		time.Unix(225, 0).UnixNano(), 0.0, testStoragePolicy, maggregation.DefaultTypes)...)
	require.Equal(t, exp, *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))
	a, err := e.find(xtime.UnixNano(testAlignedStarts[0]))
	require.NoError(t, err)
	v := a.lockedAgg
	a1, err := e.find(xtime.UnixNano(testAlignedStarts[1]))
	v1 := a1.lockedAgg
	require.NoError(t, err)
	require.False(t, v.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[0])].flushed)
	require.True(t, v1.dirty)
	require.False(t, e.flushState[xtime.UnixNano(testAlignedStarts[1])].flushed)

	// Update the first value after flushing
	require.NoError(t, e.AddUnique(time.Unix(0, testAlignedStarts[0]), aggregated.ForwardedMetric{
		Values:     []float64{122},
		PrevValues: []float64{123},
		Version:    1,
	}, metadata.ForwardMetadata{
		ResendEnabled: true,
		SourceID:      1,
	}))
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(testAlignedStarts[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	exp = expectedLocalMetricsForGaugeWithVal(
		alignedstartAtNanos[1], 122.0, testStoragePolicy, maggregation.DefaultTypes)
	exp = append(exp, expectedLocalMetricsForGaugeWithVal(
		time.Unix(225, 0).UnixNano(), 0.0, testStoragePolicy, maggregation.DefaultTypes)...)
	require.Equal(t, exp, *localRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 2, len(e.values))
	require.False(t, v.dirty)
	require.True(t, e.flushState[xtime.UnixNano(testAlignedStarts[0])].flushed)
	require.True(t, v1.dirty)
	require.False(t, e.flushState[xtime.UnixNano(testAlignedStarts[1])].flushed)
}

func TestGaugeElemResendBufferForwarding(t *testing.T) {
	alignedstartAtNanos := []int64{
		time.Unix(210, 0).UnixNano(),
		time.Unix(220, 0).UnixNano(),
		time.Unix(230, 0).UnixNano(),
		time.Unix(240, 0).UnixNano(),
	}
	gaugeVals := []float64{123.0, 456.0, 589.0}
	aggregationTypes := maggregation.Types{maggregation.Last}
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	bufferPast := time.Second * 30
	opts := newTestOptions().
		SetBufferForPastTimedMetricFn(func(resolution time.Duration) time.Duration {
			return resolution + bufferPast
		})
	sumRollup := applied.OpUnion{
		Type: pipeline.RollupOpType,
		Rollup: applied.RollupOp{
			ID:            []byte("foo.bar"),
			AggregationID: maggregation.MustCompressTypes(maggregation.Sum),
		},
	}
	pipe := applied.NewPipeline([]applied.OpUnion{
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.Increase},
		},
		sumRollup,
	})
	data := testGaugeData
	data.Pipeline = pipe
	data.AggTypes = aggregationTypes
	e := testGaugeElemWithData(t, alignedstartAtNanos[:3], gaugeVals, data, opts, true)

	aggKey := aggregationKey{
		aggregationID:     maggregation.MustCompressTypes(maggregation.Sum),
		storagePolicy:     testStoragePolicy,
		numForwardedTimes: 1,
	}
	expectedOnFlushedRes := []testOnForwardedFlushedData{
		{
			aggregationKey: aggKey,
		},
	}

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 3, len(e.values))

	// Consume one value.
	expectedForwardedRes := []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(220, 0).UnixNano(),
			value:          123.0,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(alignedstartAtNanos[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 3, len(e.values))
	require.Len(t, e.flushState, 1)
	fState := e.flushState[xtime.UnixNano(alignedstartAtNanos[0])]
	consumedVal := fState.consumedValues
	require.Len(t, consumedVal, 1)
	require.Equal(t, 123.0, consumedVal[0])

	// Consume all values.
	expectedForwardedRes = []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(230, 0).UnixNano(),
			value:          333,
		},
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(240, 0).UnixNano(),
			value:          133,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 3, len(e.values))
	require.Len(t, e.flushState, 3)
	require.Equal(t, 123.0, e.flushState[xtime.ToUnixNano(time.Unix(210, 0))].consumedValues[0])
	require.Equal(t, 456.0, e.flushState[xtime.ToUnixNano(time.Unix(220, 0))].consumedValues[0])
	require.Equal(t, 589.0, e.flushState[xtime.ToUnixNano(time.Unix(230, 0))].consumedValues[0])

	// Update a previous value
	mu := unaggregated.MetricUnion{
		GaugeVal: 124.0,
	}
	require.NoError(t, e.AddUnion(time.Unix(210, 0), mu, true))
	expectedForwardedRes = []testForwardedMetricWithMetadata{
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(220, 0).UnixNano(),
			value:          124,
		},
		{
			aggregationKey: aggKey,
			timeNanos:      time.Unix(230, 0).UnixNano(),
			value:          332,
		},
	}
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t,
		e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
			localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	verifyForwardedMetrics(t, expectedForwardedRes, *forwardRes)
	verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 3, len(e.values))
	require.Len(t, e.flushState, 3)
	require.Equal(t, 124.0, e.flushState[xtime.ToUnixNano(time.Unix(210, 0))].consumedValues[0])
	require.Equal(t, 456.0, e.flushState[xtime.ToUnixNano(time.Unix(220, 0))].consumedValues[0])
	require.Equal(t, 589.0, e.flushState[xtime.ToUnixNano(time.Unix(230, 0))].consumedValues[0])
}

func TestGaugeElemReset(t *testing.T) {
	alignedstartAtNanos := []int64{
		time.Unix(210, 0).UnixNano(),
		time.Unix(220, 0).UnixNano(),
		time.Unix(230, 0).UnixNano(),
		time.Unix(240, 0).UnixNano(),
	}
	gaugeVals := []float64{123.0, 456.0, 589.0}
	aggregationTypes := maggregation.Types{maggregation.Sum}
	isEarlierThanFn := isStandardMetricEarlierThan
	timestampNanosFn := standardMetricTimestampNanos
	opts := newTestOptions().SetDiscardNaNAggregatedValues(false)

	testPipeline := applied.NewPipeline([]applied.OpUnion{
		{
			Type:           pipeline.TransformationOpType,
			Transformation: pipeline.TransformationOp{Type: transformation.Reset},
		},
	})

	e := testGaugeElem(alignedstartAtNanos[:3], gaugeVals, aggregationTypes, testPipeline, opts)

	// Consume values before an early-enough time.
	localFn, localRes := testFlushLocalMetricFn()
	forwardFn, forwardRes := testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes := testOnForwardedFlushedFn()
	require.False(t, e.Consume(0, isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 3, len(e.values))

	// Consume one value.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[1], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, (*localRes)[0].timeNanos, alignedstartAtNanos[1])
	require.Equal(t, (*localRes)[0].value, 123.0)
	require.Equal(t, (*localRes)[1].timeNanos, alignedstartAtNanos[1]+int64(5*time.Second))
	require.Equal(t, (*localRes)[1].value, 0.0)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 3, len(e.values))
	for _, f := range e.flushState {
		require.Equal(t, 0, len(f.consumedValues))
	}

	// Consume all values.
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, (*localRes)[0].timeNanos, alignedstartAtNanos[2])
	require.Equal(t, (*localRes)[0].value, 456.0)
	require.Equal(t, (*localRes)[1].timeNanos, alignedstartAtNanos[2]+int64(5*time.Second))
	require.Equal(t, (*localRes)[1].value, 0.0)
	require.Equal(t, (*localRes)[2].timeNanos, alignedstartAtNanos[3])
	require.Equal(t, (*localRes)[2].value, 589.0)
	require.Equal(t, (*localRes)[3].timeNanos, alignedstartAtNanos[3]+int64(5*time.Second))
	require.Equal(t, (*localRes)[3].value, 0.0)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
	for _, f := range e.flushState {
		require.Equal(t, 0, len(f.consumedValues))
	}

	// Tombstone the element and discard all values.
	e.tombstoned = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, _ = testOnForwardedFlushedFn()
	require.True(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	// verifyOnForwardedFlushResult(t, expectedOnFlushedRes, *onForwardedFlushedRes)
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 1, len(e.values))

	// Reading and discarding values from a closed element is no op.
	e.closed = true
	localFn, localRes = testFlushLocalMetricFn()
	forwardFn, forwardRes = testFlushForwardedMetricFn()
	onForwardedFlushedFn, onForwardedFlushedRes = testOnForwardedFlushedFn()
	require.False(t, e.Consume(alignedstartAtNanos[3], isEarlierThanFn, timestampNanosFn, standardMetricTargetNanos,
		localFn, forwardFn, onForwardedFlushedFn, 0, consumeType))
	require.Equal(t, 0, len(*localRes))
	require.Equal(t, 0, len(*forwardRes))
	require.Equal(t, 0, len(*onForwardedFlushedRes))
	require.Equal(t, 1, len(e.values))
}

func TestGaugeElemClose(t *testing.T) {
	e := testGaugeElem(testAlignedStarts[:len(testAlignedStarts)-1], testGaugeVals,
		maggregation.DefaultTypes, applied.DefaultPipeline, newTestOptions())
	require.False(t, e.closed)

	// Closing the element.
	e.Close()

	// Closing a second time should have no impact.
	e.Close()

	require.True(t, e.closed)
	require.Nil(t, e.id)
	require.Equal(t, parsedPipeline{}, e.parsedPipeline)
	require.Nil(t, e.writeForwardedMetricFn)
	require.Nil(t, e.onForwardedAggregationWrittenFn)
	require.Nil(t, e.cachedSourceSets)
	require.Equal(t, 0, len(e.values))
	require.Equal(t, 0, len(e.toConsume))
	for _, f := range e.flushState {
		require.Equal(t, 0, len(f.consumedValues))
	}
	require.NotNil(t, e.values)
}

func TestGaugeFindOrCreateNoSourceSet(t *testing.T) {
	e, err := NewGaugeElem(testGaugeElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	inputs := []xtime.UnixNano{10, 10, 20, 10, 15}
	expected := make(map[xtime.UnixNano]bool)
	for _, inputNanos := range inputs {
		expected[inputNanos] = true
		res, err := e.findOrCreate(int64(inputNanos), createAggregationOptions{initSourceSet: false})
		require.NoError(t, err)
		require.Equal(t, e.values[inputNanos].lockedAgg, res)
		require.Equal(t, len(expected), len(e.values))
		for k := range e.values {
			_, ok := expected[k]
			require.True(t, ok)
		}
	}
}

func TestDirtyConsumption(t *testing.T) {
	e, err := NewCounterElem(testCounterElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)

	resolution := e.sp.Resolution().Window

	// Test 100 startAligned timestamps (and targetNanos is far enough to consume any of them)
	n := int64(100)
	targetNanos := resolution.Nanoseconds() * (n + 1)

	var (
		wg sync.WaitGroup
		// time by consume count
		consumed = make(map[xtime.UnixNano]int)
	)
	stop := atomic.NewBool(false)
	for i := int64(0); i < n; i++ {
		i := i
		wg.Add(1)

		// For each timestamp, repeatedly add an aggregation (will findOrCreate and make dirty)
		// and then consume the presently dirty points. This validates that new times are consumed
		// once added, and existing times are made dirty to be re-consumed.
		go func() {
			for !stop.Load() {
				require.NoError(t, e.AddUnion(time.Unix(0, i*resolution.Nanoseconds()), testCounter, false))

				e.Lock()
				e.dirtyToConsumeWithLock(targetNanos, resolution, isStandardMetricEarlierThan)

				for _, c := range e.toConsume {
					consumed[c.startAt]++
				}
				e.Unlock()
			}
			wg.Done()
		}()
	}

	time.Sleep(time.Second * 2)
	stop.Store(true)
	wg.Wait()

	require.Equal(t, 0, len(e.dirty), "all dirty should have been consumed")
	require.Equal(t, int(n), len(e.values), "values should have all N consumed times")
	require.Equal(t, int(n), len(consumed), "consumed should have all N consumed times")

	for _, count := range consumed {
		// Ensure times were consumed at least (a) upon first creation and (b) upon subsequent dirty
		require.True(t, count > 1)
	}
}

// confirms that when panics are enabled, we do panic, and when disabled,
// the code safely proceeds (and does not encounter some unexpected runtime panic)
func TestPanics(t *testing.T) {
	opts := newTestOptions().SetInstrumentOptions(instrument.NewOptions())
	elemData := testCounterElemData

	tests := []struct {
		name   string
		testFn func()
	}{
		{
			name: "panic if non-existent flush states are expired",
			testFn: func() {
				elem, err := NewCounterElem(elemData, NewElemOptions(opts))
				require.NoError(t, err)
				elem.flushStateToExpire = []xtime.UnixNano{0, 1, 2}
				elem.expireFlushState()
			},
		},
		{
			name: "panic if minStartTime moved beyond remaining values on Close",
			testFn: func() {
				elem, err := NewCounterElem(elemData, NewElemOptions(opts))
				require.NoError(t, err)
				elem.flushStateToExpire = []xtime.UnixNano{0, 1, 2}

				// Panic case where minStartTime > min(elem.values)
				elem.values[10] = timedCounter{
					lockedAgg: &lockedCounterAggregation{
						aggregation: newCounterAggregation(raggregation.NewCounter(elem.aggOpts)),
					},
				}
				elem.minStartTime = 11

				elem.Close()
			},
		},
		{
			name: "panic from dangling flushState on Close",
			testFn: func() {
				elem, err := NewCounterElem(elemData, NewElemOptions(opts))
				require.NoError(t, err)
				elem.flushStateToExpire = []xtime.UnixNano{0, 1, 2}

				// i.e. flushState to cleanup but no associated elem.values
				elem.flushState[10] = flushState{}

				elem.Close()
			},
		},
		{
			name: "panic if invalid flush state",
			testFn: func() {
				elem, err := NewCounterElem(elemData, NewElemOptions(opts))
				require.NoError(t, err)

				counter := timedCounter{
					lockedAgg: &lockedCounterAggregation{
						aggregation: newCounterAggregation(raggregation.NewCounter(elem.aggOpts)),
					},
				}

				elem.toConsume, _ = elem.appendConsumeStateWithLock(counter, elem.toConsume, nil)
				toConsume := elem.toConsume[0]
				toConsume.startAt = xtime.UnixNano(10)

				// Panic case is where dirty + flushed + !resendEnabled since only resendEnabled
				// allows for updating dirty/flushed aggs.
				toConsume.dirty = true
				toConsume.resendEnabled = false
				elem.flushState[toConsume.startAt] = flushState{
					flushed: true,
				}

				localFn, _ := testFlushLocalMetricFn()
				forwardFn, _ := testFlushForwardedMetricFn()
				elem.processValue(toConsume,
					standardMetricTimestampNanos,
					localFn,
					forwardFn,
					0, 0, 0,
					newFlushMetrics(tally.NewTestScope("", nil)))
			},
		},
		{
			name: "panic if no consumedValues",
			testFn: func() {
				elemData := testCounterElemData
				elemData.Pipeline = testPipeline
				elem, err := NewCounterElem(elemData, NewElemOptions(opts))
				require.NoError(t, err)

				counter := timedCounter{
					lockedAgg: &lockedCounterAggregation{
						aggregation: newCounterAggregation(raggregation.NewCounter(elem.aggOpts)),
					},
				}
				elem.toConsume, _ = elem.appendConsumeStateWithLock(counter, elem.toConsume, nil)
				toConsume := elem.toConsume[0]

				// Panic case is where startAt has a flushState entry but prevStartTime does not.
				toConsume.startAt = xtime.UnixNano(10)
				toConsume.prevStartTime = xtime.UnixNano(9)
				elem.flushState[toConsume.startAt] = flushState{
					flushed: true,
				}
				elem.flushState[toConsume.startAt] = flushState{}

				localFn, _ := testFlushLocalMetricFn()
				forwardFn, _ := testFlushForwardedMetricFn()
				elem.processValue(toConsume,
					standardMetricTimestampNanos,
					localFn,
					forwardFn,
					0, 0, 0,
					newFlushMetrics(tally.NewTestScope("", nil)))
			},
		},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			// Test panics enabled/disabled.
			instrument.SetShouldPanicEnvironmentVariable(true)
			assert.Panics(t, tt.testFn, "should panic but did not")
			instrument.SetShouldPanicEnvironmentVariable(false)
			tt.testFn()
		})
	}
}

// nolint: dupl
func TestGaugeFindOrCreateWithSourceSet(t *testing.T) {
	e, err := NewGaugeElem(testGaugeElemData, NewElemOptions(newTestOptions()))
	require.NoError(t, err)
	e.cachedSourceSets = make([]map[uint32]*bitset.BitSet, 0)

	inputs := []xtime.UnixNano{10, 20}
	expected := make(map[xtime.UnixNano]bool)
	for _, inputNanos := range inputs {
		expected[inputNanos] = true
		res, err := e.findOrCreate(int64(inputNanos), createAggregationOptions{initSourceSet: true})
		require.NoError(t, err)
		require.Equal(t, e.values[inputNanos].lockedAgg, res)
		require.NotNil(t, e.values[inputNanos].lockedAgg.sourcesSeen)
		require.Equal(t, len(expected), len(e.values))
		for k := range e.values {
			_, ok := expected[k]
			require.True(t, ok)
		}
	}
	require.Equal(t, 0, len(e.cachedSourceSets))
}

func TestExpireValues(t *testing.T) {
	opts := newTestOptions()
	resolutionDuration := opts.DefaultStoragePolicies()[0].Resolution().Window
	resolution := xtime.UnixNano(resolutionDuration)
	bufferPastDuration := opts.BufferForPastTimedMetricFn()(resolutionDuration)

	valsNoGaps := []xtime.UnixNano{
		resolution * 1,
		resolution * 2,
		resolution * 3,
	}
	valsGaps := []xtime.UnixNano{
		resolution * 1,
		resolution * 5,
		resolution * 6,
		resolution * 10,
	}

	for _, test := range []struct {
		name             string
		targetNanos      xtime.UnixNano
		resendEnabled    bool
		values           []xtime.UnixNano
		expectedToExpire []xtime.UnixNano
		expectedValues   []xtime.UnixNano
	}{
		// no gaps - resend disabled
		{
			name:             "no gaps - resend disabled - zero target",
			targetNanos:      resolution * 0,
			resendEnabled:    false,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend disabled - target val 0",
			targetNanos:      valsNoGaps[0],
			resendEnabled:    false,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend disabled - target val 1",
			targetNanos:      valsNoGaps[1],
			resendEnabled:    false,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend disabled - target val 2",
			targetNanos:      valsNoGaps[2],
			resendEnabled:    false,
			values:           valsNoGaps,
			expectedToExpire: valsNoGaps[0:1], expectedValues: valsNoGaps[1:],
		},
		{
			name:             "no gaps - resend disabled - target val 2 + 1 resolution",
			targetNanos:      valsNoGaps[2].Add(resolutionDuration),
			resendEnabled:    false,
			values:           valsNoGaps,
			expectedToExpire: valsNoGaps[0:2], expectedValues: valsNoGaps[2:],
		},
		{
			name:             "no gaps - resend disabled - target val 2 + 2 resolution",
			targetNanos:      valsNoGaps[2].Add(2 * resolutionDuration),
			resendEnabled:    false,
			values:           valsNoGaps,
			expectedToExpire: valsNoGaps[0:2], expectedValues: valsNoGaps[2:],
		},
		// no gaps - resend enabled
		{
			name:             "no gaps - resend enabled - target zero",
			targetNanos:      resolution * 0,
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 0",
			targetNanos:      valsNoGaps[0],
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:          "no gaps - resend enabled - target val 1",
			targetNanos:   valsNoGaps[1],
			resendEnabled: true, values: valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 2",
			targetNanos:      valsNoGaps[2],
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 2 + 1 resolution",
			targetNanos:      valsNoGaps[2].Add(resolutionDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 2 + 2 resolution",
			targetNanos:      valsNoGaps[2].Add(2 * resolutionDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 0 + buffer past",
			targetNanos:      valsNoGaps[0].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 1 + buffer past",
			targetNanos:      valsNoGaps[1].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsNoGaps,
		},
		{
			name:             "no gaps - resend enabled - target val 2 + buffer past",
			targetNanos:      valsNoGaps[2].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: valsNoGaps[0:1], expectedValues: valsNoGaps[1:],
		},
		{
			name:             "no gaps - resend enabled - target val 2 + 1 resolution + buffer past",
			targetNanos:      valsNoGaps[2].Add(resolutionDuration).Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: valsNoGaps[0:2], expectedValues: valsNoGaps[2:],
		},
		{
			name:             "no gaps - resend enabled - target val 2 + 2 resolution + buffer past",
			targetNanos:      valsNoGaps[2].Add(2 * resolutionDuration).Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsNoGaps,
			expectedToExpire: valsNoGaps[0:2], expectedValues: valsNoGaps[2:],
		},
		// gaps - resend disabled
		{
			name:             "gaps - resend disabled - zero target",
			targetNanos:      resolution * 0,
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend disabled - target val 0",
			targetNanos:      valsGaps[0],
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend disabled - target val 1",
			targetNanos:      valsGaps[1],
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend disabled - target val 2",
			targetNanos:      valsGaps[2],
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:1], expectedValues: valsGaps[1:],
		},
		{
			name:             "gaps - resend disabled - target between val 2 and val 3 (1 resolution)",
			targetNanos:      valsGaps[2].Add(resolutionDuration),
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:2], expectedValues: valsGaps[2:],
		},
		{
			name:             "gaps - resend disabled - target between val 2 and val 3 (2 resolution)",
			targetNanos:      valsGaps[2].Add(2 * resolutionDuration),
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:2], expectedValues: valsGaps[2:],
		},
		{
			name:             "gaps - resend disabled - target val 3",
			targetNanos:      valsGaps[3],
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:2], expectedValues: valsGaps[2:],
		},
		{
			name:             "gaps - resend disabled - target val 3 + 1 resolution",
			targetNanos:      valsGaps[3].Add(resolutionDuration),
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:3], expectedValues: valsGaps[3:],
		},
		{
			name:             "gaps - resend disabled - target val 3 + 2 resolution",
			targetNanos:      valsGaps[3].Add(2 * resolutionDuration),
			resendEnabled:    false,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:3], expectedValues: valsGaps[3:],
		},
		// gaps - resend enabled
		{
			name:          "gaps - resend enabled - target zero",
			targetNanos:   resolution * 0,
			resendEnabled: true,
			values:        valsGaps, expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 0",
			targetNanos:      valsGaps[0],
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 1",
			targetNanos:      valsGaps[1],
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 2",
			targetNanos:      valsGaps[2],
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 3",
			targetNanos:      valsGaps[3],
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 3 + 1 resolution",
			targetNanos:      valsGaps[3].Add(resolutionDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 3 + 2 resolution",
			targetNanos:      valsGaps[3].Add(2 * resolutionDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 0 + buffer past",
			targetNanos:      valsGaps[0].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 1 + buffer past",
			targetNanos:      valsGaps[1].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: []xtime.UnixNano{}, expectedValues: valsGaps,
		},
		{
			name:             "gaps - resend enabled - target val 2 + buffer past",
			targetNanos:      valsGaps[2].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:1], expectedValues: valsGaps[1:],
		},
		{
			name:             "gaps - resend enabled - target between val 2 and 3 (1 resolution) + buffer past",
			targetNanos:      valsGaps[2].Add(resolutionDuration).Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:2], expectedValues: valsGaps[2:],
		},
		{
			name:             "gaps - resend enabled - target between val 2 and 3 (2 resolution) + buffer past",
			targetNanos:      valsGaps[2].Add(2 * resolutionDuration).Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:2], expectedValues: valsGaps[2:],
		},
		{
			name:             "gaps - resend enabled - target val 3 + buffer past",
			targetNanos:      valsGaps[3].Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:2], expectedValues: valsGaps[2:],
		},
		{
			name:             "gaps - resend enabled - target val 3 + 1 resolution + buffer past",
			targetNanos:      valsGaps[3].Add(resolutionDuration).Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:3], expectedValues: valsGaps[3:],
		},
		{
			name:             "gaps - resend enabled - target val 3 + 2 resolution + buffer past",
			targetNanos:      valsGaps[3].Add(2 * resolutionDuration).Add(bufferPastDuration),
			resendEnabled:    true,
			values:           valsGaps,
			expectedToExpire: valsGaps[0:3], expectedValues: valsGaps[3:],
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			e, err := NewCounterElem(testCounterElemData, NewElemOptions(opts))
			require.NoError(t, err)
			require.Equal(t, 0, len(e.flushStateToExpire))

			// Add initial toExpire since expireValues should reset this.
			e.flushStateToExpire = make([]xtime.UnixNano, 10)

			// Add test values.
			for _, v := range test.values {
				a, err := e.findOrCreate(int64(v), createAggregationOptions{})
				a.resendEnabled = test.resendEnabled
				// need to manually seed the flush state since we don't call Consume(), which takes care of setting
				// the flush state for expireValuesWithLock to use.
				e.flushState[v] = flushState{latestResendEnabled: test.resendEnabled}
				require.NoError(t, err)
			}

			// Expire up to target.
			e.expireValuesWithLock(int64(test.targetNanos), isStandardMetricEarlierThan,
				newFlushMetrics(tally.NewTestScope("", nil)))

			// Validate toExpire and remaining values.
			require.Equal(t, len(test.expectedToExpire), len(e.flushStateToExpire))
			for i, toExpire := range test.expectedToExpire {
				require.Equal(t, toExpire, e.flushStateToExpire[i], "missing expire")
			}
			require.Equal(t, len(test.expectedValues), len(e.values))
			for _, value := range test.expectedValues {
				_, ok := e.values[value]
				require.True(t, ok, "missing value")
			}
		})
	}
}

type testIndexData struct {
	index int
	data  []int64
}

type testSuffixAndValue struct {
	aggType maggregation.Type
	value   float64
}

type testLocalMetricWithMetadata struct {
	idPrefix  []byte
	id        id.RawID
	idSuffix  []byte
	timeNanos int64
	value     float64
	sp        policy.StoragePolicy
}

type testForwardedMetricWithMetadata struct {
	aggregationKey aggregationKey
	timeNanos      int64
	value          float64
}

type testOnForwardedFlushedData struct {
	aggregationKey aggregationKey
	expiredTimes   []xtime.UnixNano
}

func testFlushLocalMetricFn() (
	flushLocalMetricFn,
	*[]testLocalMetricWithMetadata,
) {
	var result []testLocalMetricWithMetadata
	return func(
		idPrefix []byte,
		id id.RawID,
		idSuffix []byte,
		timeNanos int64,
		value float64,
		annotation []byte,
		sp policy.StoragePolicy,
	) {
		result = append(result, testLocalMetricWithMetadata{
			idPrefix:  idPrefix,
			id:        id,
			idSuffix:  idSuffix,
			timeNanos: timeNanos,
			value:     value,
			sp:        sp,
		})
	}, &result
}

func testFlushForwardedMetricFn() (
	flushForwardedMetricFn,
	*[]testForwardedMetricWithMetadata,
) {
	var result []testForwardedMetricWithMetadata
	return func(
		writeFn writeForwardedMetricFn,
		aggregationKey aggregationKey,
		timeNanos int64,
		value float64,
		prevValue float64,
		annotation []byte,
		resendEnabled bool,
	) {
		result = append(result, testForwardedMetricWithMetadata{
			aggregationKey: aggregationKey,
			timeNanos:      timeNanos,
			value:          value,
		})
	}, &result
}

func testOnForwardedFlushedFn() (
	onForwardingElemFlushedFn,
	*[]testOnForwardedFlushedData,
) {
	var result []testOnForwardedFlushedData
	return func(
		onDoneFn onForwardedAggregationDoneFn,
		aggregationKey aggregationKey,
		expiredTimes []xtime.UnixNano,
	) {
		result = append(result, testOnForwardedFlushedData{
			aggregationKey: aggregationKey,
			expiredTimes:   expiredTimes,
		})
	}, &result
}

func testCounterElem(
	alignedstartAtNanos []int64,
	counterVals []int64,
	aggTypes maggregation.Types,
	pipeline applied.Pipeline,
	opts Options,
) *CounterElem {
	elemData := testCounterElemData
	elemData.AggTypes = aggTypes
	elemData.Pipeline = pipeline
	e := MustNewCounterElem(elemData, NewElemOptions(opts))
	for i, aligned := range alignedstartAtNanos {
		counter := &lockedCounterAggregation{
			aggregation: newCounterAggregation(raggregation.NewCounter(e.aggOpts)),
		}
		counter.dirty = true
		counter.aggregation.Update(time.Unix(0, aligned), counterVals[i], nil)
		startAligned := xtime.UnixNano(aligned)
		e.values[startAligned] = timedCounter{
			startAt:   startAligned,
			lockedAgg: counter,
		}
		e.dirty = append(e.dirty, startAligned)
	}
	e.minStartTime = xtime.UnixNano(alignedstartAtNanos[0])
	e.maxStartTime = xtime.UnixNano(alignedstartAtNanos[len(alignedstartAtNanos)-1])
	return e
}

func testTimerElem(
	alignedstartAtNanos []int64,
	timerBatches [][]float64,
	aggTypes maggregation.Types,
	pipeline applied.Pipeline,
	opts Options,
) *TimerElem {
	elemData := testTimerElemData
	elemData.Pipeline = pipeline
	elemData.AggTypes = aggTypes
	e := MustNewTimerElem(elemData, NewElemOptions(opts))
	for i, aligned := range alignedstartAtNanos {
		newTimer := raggregation.NewTimer(opts.AggregationTypesOptions().Quantiles(), opts.StreamOptions(), e.aggOpts)
		timer := &lockedTimerAggregation{
			aggregation: newTimerAggregation(newTimer),
		}
		timer.dirty = true
		timer.aggregation.AddBatch(time.Now(), timerBatches[i], nil)
		startAligned := xtime.UnixNano(aligned)
		e.values[startAligned] = timedTimer{
			startAt:   startAligned,
			lockedAgg: timer,
		}
		e.dirty = append(e.dirty, startAligned)
	}
	e.minStartTime = xtime.UnixNano(alignedstartAtNanos[0])
	e.maxStartTime = xtime.UnixNano(alignedstartAtNanos[len(alignedstartAtNanos)-1])
	return e
}

var testGaugeData = ElemData{
	ID:                testGaugeID,
	StoragePolicy:     testStoragePolicy,
	Pipeline:          testPipeline,
	NumForwardedTimes: testNumForwardedTimes,
}

func testGaugeElem(
	alignedstartAtNanos []int64,
	gaugeVals []float64,
	aggTypes maggregation.Types,
	pipeline applied.Pipeline,
	opts Options,
) *GaugeElem {
	data := testGaugeData
	data.Pipeline = pipeline
	data.AggTypes = aggTypes
	return testGaugeElemWithData(&testing.T{}, alignedstartAtNanos, gaugeVals, data, opts, false)
}

func testGaugeElemWithData(
	t *testing.T,
	alignedstartAtNanos []int64,
	gaugeVals []float64,
	data ElemData,
	opts Options,
	resendEnabled bool,
) *GaugeElem {
	e := MustNewGaugeElem(data, NewElemOptions(opts))
	require.NoError(t, e.ResetSetData(data))
	for i, aligned := range alignedstartAtNanos {
		gauge := &lockedGaugeAggregation{
			aggregation:   newGaugeAggregation(raggregation.NewGauge(e.aggOpts)),
			sourcesSeen:   make(map[uint32]*bitset.BitSet),
			resendEnabled: resendEnabled,
		}
		gauge.dirty = true
		// offset the timestamp by 1 so the gauge value can be updated using the aligned timestamp later.
		gauge.aggregation.Update(time.Unix(0, aligned-1), gaugeVals[i], nil)
		startAligned := xtime.UnixNano(aligned)
		e.values[startAligned] = timedGauge{
			startAt:   startAligned,
			lockedAgg: gauge,
		}
		e.dirty = append(e.dirty, startAligned)
	}
	e.minStartTime = xtime.UnixNano(alignedstartAtNanos[0])
	e.maxStartTime = xtime.UnixNano(alignedstartAtNanos[len(alignedstartAtNanos)-1])
	return e
}

func expectCounterSuffix(aggType maggregation.Type) []byte {
	return testOpts.AggregationTypesOptions().TypeStringForCounter(aggType)
}

func expectTimerSuffix(aggType maggregation.Type) []byte {
	return testOpts.AggregationTypesOptions().TypeStringForTimer(aggType)
}

func expectGaugeSuffix(aggType maggregation.Type) []byte {
	return testOpts.AggregationTypesOptions().TypeStringForGauge(aggType)
}

func expectedLocalMetricsForCounter(
	timeNanos int64,
	sp policy.StoragePolicy,
	aggTypes maggregation.Types,
) []testLocalMetricWithMetadata {
	if !aggTypes.IsDefault() {
		var res []testLocalMetricWithMetadata
		for _, aggType := range aggTypes {
			res = append(res, testLocalMetricWithMetadata{
				idPrefix:  []byte("stats.counts."),
				id:        testCounterID,
				idSuffix:  expectCounterSuffix(aggType),
				timeNanos: timeNanos,
				value:     float64(testCounter.CounterVal),
				sp:        sp,
			})
		}
		return res
	}
	return []testLocalMetricWithMetadata{
		{
			idPrefix:  []byte("stats.counts."),
			id:        testCounterID,
			idSuffix:  nil,
			timeNanos: timeNanos,
			value:     float64(testCounter.CounterVal),
			sp:        sp,
		},
	}
}

func expectedLocalMetricsForTimer(
	timeNanos int64,
	sp policy.StoragePolicy,
	aggTypes maggregation.Types,
) []testLocalMetricWithMetadata {
	// This needs to be a list as the order of the result matters in some test.
	data := []testSuffixAndValue{
		{maggregation.Sum, 18.0},
		{maggregation.SumSq, 83.38},
		{maggregation.Mean, 3.6},
		{maggregation.Min, 1.0},
		{maggregation.Max, 6.5},
		{maggregation.Count, 5.0},
		{maggregation.Stdev, 2.15522620622523},
		{maggregation.Median, 3.5},
		{maggregation.P50, 3.5},
		{maggregation.P95, 6.5},
		{maggregation.P99, 6.5},
	}
	var expected []testLocalMetricWithMetadata
	if !aggTypes.IsDefault() {
		for _, aggType := range aggTypes {
			for _, d := range data {
				if d.aggType == aggType {
					expected = append(expected, testLocalMetricWithMetadata{
						idPrefix:  []byte("stats.timers."),
						id:        testBatchTimerID,
						idSuffix:  expectTimerSuffix(aggType),
						timeNanos: timeNanos,
						value:     d.value,
						sp:        sp,
					})
				}
			}
		}
		return expected
	}

	for _, d := range data {
		expected = append(expected, testLocalMetricWithMetadata{
			idPrefix:  []byte("stats.timers."),
			id:        testBatchTimerID,
			idSuffix:  expectTimerSuffix(d.aggType),
			timeNanos: timeNanos,
			value:     d.value,
			sp:        sp,
		})
	}
	return expected
}

func expectedLocalMetricsForGauge(
	timeNanos int64,
	sp policy.StoragePolicy,
	aggTypes maggregation.Types,
) []testLocalMetricWithMetadata {
	return expectedLocalMetricsForGaugeWithVal(timeNanos, testGauge.GaugeVal, sp, aggTypes)
}

func expectedLocalMetricsForGaugeWithVal(
	timeNanos int64,
	val float64,
	sp policy.StoragePolicy,
	aggTypes maggregation.Types,
) []testLocalMetricWithMetadata {
	if !aggTypes.IsDefault() {
		var res []testLocalMetricWithMetadata
		for _, aggType := range aggTypes {
			res = append(res, testLocalMetricWithMetadata{
				idPrefix:  []byte("stats.gauges."),
				id:        testGaugeID,
				idSuffix:  expectGaugeSuffix(aggType),
				timeNanos: timeNanos,
				value:     val,
				sp:        sp,
			})
		}
		return res
	}
	return []testLocalMetricWithMetadata{
		{
			idPrefix:  []byte("stats.gauges."),
			id:        testGaugeID,
			idSuffix:  nil,
			timeNanos: timeNanos,
			value:     val,
			sp:        sp,
		},
	}
}

func verifyForwardedMetrics(t *testing.T, expected, actual []testForwardedMetricWithMetadata) {
	require.Equal(t, len(expected), len(actual))
	for i := 0; i < len(expected); i++ {
		require.True(t, expected[i].aggregationKey.Equal(actual[i].aggregationKey))
		require.Equal(t, expected[i].timeNanos, actual[i].timeNanos)
		if math.IsNaN(expected[i].value) {
			require.True(t, math.IsNaN(actual[i].value))
		} else {
			require.Equal(t, expected[i].value, actual[i].value)
		}
	}
}

func verifyOnForwardedFlushResult(t *testing.T, expected, actual []testOnForwardedFlushedData) {
	require.Equal(t, len(expected), len(actual))
	for i := 0; i < len(expected); i++ {
		require.True(t, expected[i].aggregationKey.Equal(actual[i].aggregationKey))
		require.Equal(t, len(expected[i].expiredTimes), len(actual[i].expiredTimes))
		for j := 0; j < len(expected[i].expiredTimes); j++ {
			require.Equal(t, expected[i].expiredTimes[j], actual[i].expiredTimes[j])
		}
	}
}


<|start_filename|>src/query/util/json/writer.go<|end_filename|>
// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// Package json contains logic for writing JSON.
package json

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"math"
	"strconv"
)

var (
	errContainerMismatch  = errors.New("container mismatch")
	errNotInContainer     = errors.New("not in container")
	errFieldNotAllowed    = errors.New("field not allowed")
	errValueNotAllowed    = errors.New("value not allowed")
	errContainerStillOpen = errors.New("container still open")
)

type (
	writeState    int
	containerType int
)

const (
	object containerType = iota
	array
)

const (
	writeStart writeState = iota
	writeBeforeFirstField
	writeBeforeNthField
	writeBeforeFieldValue
	writeBeforeFirstArrayElement
	writeBeforeNthArrayElement
	writeEnd
)

var writeValueAllowed = map[writeState]struct{}{
	writeStart:                   {},
	writeBeforeFieldValue:        {},
	writeBeforeFirstArrayElement: {},
	writeBeforeNthArrayElement:   {},
}

func (s writeState) isValueAllowed() bool {
	_, allowed := writeValueAllowed[s]
	return allowed
}

// A Writer can be used to directly stream JSON results without going through
// an intermediate object layer.
type Writer interface {
	BeginObject()
	BeginObjectField(name string)
	BeginObjectBytesField(name []byte)
	EndObject()
	BeginArray()
	EndArray()
	WriteBool(b bool)
	WriteNull()
	WriteFloat64(n float64)
	WriteInt(n int)
	WriteString(s string)
	WriteBytesString(s []byte)
	Flush() error
	Close() error
}

type writer struct {
	w          *bufio.Writer
	buff       *bytes.Reader
	state      writeState
	containers []containerType
	err        error
}

// NewWriter creates a new JSON token writer
func NewWriter(w io.Writer) Writer {
	return &writer{
		w:          bufio.NewWriter(w),
		buff:       bytes.NewReader(nil),
		containers: make([]containerType, 0, 5),
	}
}

// BeginObject begins a new object
func (w *writer) BeginObject() {
	if !w.beginValue() {
		return
	}

	w.containers = append(w.containers, object)
	w.err = w.w.WriteByte('{')
	w.state = writeBeforeFirstField
}

// BeginObjectField begins a new object field with the given name
func (w *writer) BeginObjectField(name string) {
	w.err = w.beginObjectFieldStart()
	if w.err != nil {
		return
	}

	w.writeString(name)
	if w.err != nil {
		return
	}

	w.err = w.beginObjectFieldEnd()
}

func (w *writer) beginObjectFieldStart() error {
	if w.err != nil {
		return w.err
	}

	if w.state != writeBeforeFirstField && w.state != writeBeforeNthField {
		return errFieldNotAllowed
	}

	if w.state == writeBeforeNthField {
		if err := w.w.WriteByte(','); err != nil {
			return err
		}
	}

	return nil
}

func (w *writer) beginObjectFieldEnd() error {
	if err := w.w.WriteByte(':'); err != nil {
		return err
	}
	w.state = writeBeforeFieldValue
	return nil
}

func (w *writer) BeginObjectBytesField(name []byte) {
	w.err = w.beginObjectFieldStart()
	if w.err != nil {
		return
	}

	w.writeBytesString(name)
	if w.err != nil {
		return
	}

	w.err = w.beginObjectFieldEnd()
}

// EndObject finishes an open object
func (w *writer) EndObject() {
	if !w.endContainer(object) {
		return
	}

	w.err = w.w.WriteByte('}')
}

// BeginArray begins a new array value
func (w *writer) BeginArray() {
	if !w.beginValue() {
		return
	}

	w.containers = append(w.containers, array)
	w.err = w.w.WriteByte('[')
	w.state = writeBeforeFirstArrayElement
}

// EndArray finishes an array value
func (w *writer) EndArray() {
	if !w.endContainer(array) {
		return
	}

	w.err = w.w.WriteByte(']')
}

// endContainer finishes a container of the given type
func (w *writer) endContainer(expected containerType) bool {
	if w.err != nil {
		return false
	}

	if len(w.containers) == 0 {
		w.err = errNotInContainer
		return false
	}

	container := w.containers[len(w.containers)-1]
	w.containers = w.containers[:len(w.containers)-1]
	if container != expected {
		w.err = errContainerMismatch
		return false
	}

	w.endValue()
	return true
}

// WriteBool writes a boolean value
func (w *writer) WriteBool(b bool) {
	if !w.beginValue() {
		return
	}

	if b {
		_, w.err = w.w.WriteString("true")
	} else {
		_, w.err = w.w.WriteString("false")
	}
	w.endValue()
}

func (w *writer) writeNull() {
	_, w.err = w.w.WriteString("null")
}

// WriteNull writes a null value
func (w *writer) WriteNull() {
	if !w.beginValue() {
		return
	}

	w.writeNull()
	w.endValue()
}

// WriteFloat64 writes a float value
func (w *writer) WriteFloat64(n float64) {
	if !w.beginValue() {
		return
	}

	// JSON does not support NaNs or infinity
	if math.IsNaN(n) || math.IsInf(n, 0) {
		w.writeNull()
	} else {
		_, w.err = fmt.Fprintf(w.w, "%f", n)
	}

	w.endValue()
}

// WriteInt writes an int value
func (w *writer) WriteInt(n int) {
	if !w.beginValue() {
		return
	}

	_, w.err = w.w.WriteString(strconv.Itoa(n))
	w.endValue()
}

// WriteString writes a string value
func (w *writer) WriteString(s string) {
	if !w.beginValue() {
		return
	}

	w.writeString(s)

	w.endValue()
}

func (w *writer) writeString(s string) {
	if w.err = w.w.WriteByte('"'); w.err != nil {
		return
	}

	for _, c := range s {
		w.writeRune(c)
		if w.err != nil {
			return
		}
	}

	w.err = w.w.WriteByte('"')
}

func (w *writer) WriteBytesString(s []byte) {
	if !w.beginValue() {
		return
	}

	w.writeBytesString(s)

	w.endValue()
}

func (w *writer) writeBytesString(s []byte) {
	if w.err = w.w.WriteByte('"'); w.err != nil {
		return
	}

	w.buff.Reset(s)
	defer w.buff.Reset(nil) // Free holding onto byte slice.

	for {
		c, _, err := w.buff.ReadRune()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			w.err = err
			return
		}
		w.writeRune(c)
	}

	w.err = w.w.WriteByte('"')
}

func (w *writer) writeRune(r rune) {
	if r <= 31 || r == '"' || r == '\\' {
		if w.err = w.w.WriteByte('\\'); w.err != nil {
			return
		}

		switch r {
		case '"':
			if w.err = w.w.WriteByte('"'); w.err != nil {
				return
			}
		case '\\':
			if w.err = w.w.WriteByte('\\'); w.err != nil {
				return
			}
		case '\n':
			if w.err = w.w.WriteByte('n'); w.err != nil {
				return
			}
		case '\r':
			if w.err = w.w.WriteByte('r'); w.err != nil {
				return
			}
		case '\t':
			if w.err = w.w.WriteByte('t'); w.err != nil {
				return
			}
		default:
			codePoint := fmt.Sprintf("%U", r)
			if w.err = w.w.WriteByte('u'); w.err != nil {
				return
			}
			if _, w.err = w.w.WriteString(codePoint[2:]); w.err != nil {
				return
			}
		}

		return
	}

	if _, w.err = w.w.WriteRune(r); w.err != nil {
		return
	}
}

// beginValue begins a new value, confirming that the current position of the
// writer allows a value
func (w *writer) beginValue() bool {
	if w.err != nil {
		return false
	}

	if !w.state.isValueAllowed() {
		w.err = errValueNotAllowed
		return false
	}

	if w.state == writeBeforeNthArrayElement {
		if w.err = w.w.WriteByte(','); w.err != nil {
			return false
		}
	}

	return true
}

// endValue marks as value as being complete
func (w *writer) endValue() {
	if len(w.containers) == 0 {
		// End of top level object
		w.state = writeEnd
		return
	}

	c := w.containers[len(w.containers)-1]
	switch c {
	case object:
		w.state = writeBeforeNthField
	case array:
		w.state = writeBeforeNthArrayElement
	default:
		panic(fmt.Sprintf("unknown container type %d", c))
	}
}

// Flush flushes the writer
func (w *writer) Flush() error {
	if w.err != nil {
		return w.err
	}

	w.err = w.w.Flush()
	return w.err
}

// Close closes the writer, returning any write errors that have occurred
func (w *writer) Close() error {
	if w.err != nil {
		return w.err
	}

	if len(w.containers) > 0 {
		w.err = errContainerStillOpen
		return w.err
	}

	w.err = w.w.Flush()
	return w.err
}


<|start_filename|>src/query/api/v1/middleware/request_classification.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package middleware

import (
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/m3db/m3/src/cmd/services/m3query/config"
	"github.com/m3db/m3/src/query/api/v1/handler/prometheus"
	"github.com/m3db/m3/src/query/api/v1/route"
	"github.com/m3db/m3/src/x/headers"

	"github.com/prometheus/prometheus/promql/parser"
	"github.com/uber-go/tally"
	"go.uber.org/zap"
)

const (
	// promDefaultLookback is the default lookback in Prometheus, that affects
	// the actual range it fetches datapoints from.
	promDefaultLookback = time.Minute * 5
)

type classificationMetrics struct {
	classifiedResult   tally.Counter
	classifiedDuration tally.Counter
	notClassifiable    tally.Counter
	badParams          tally.Counter
	error              tally.Counter
}

func newClassificationMetrics(scope tally.Scope) *classificationMetrics {
	buildCounter := func(status string) tally.Counter {
		return scope.Tagged(map[string]string{"status": status}).Counter("count")
	}

	return &classificationMetrics{
		classifiedResult:   buildCounter("classified_result"),
		classifiedDuration: buildCounter("classified_duration"),
		notClassifiable:    buildCounter("not_classifiable"),
		badParams:          buildCounter("bad_params"),
		error:              buildCounter("error"),
	}
}

func retrieveQueryRange(expr parser.Node, rangeSoFar time.Duration) time.Duration {
	var (
		queryRange  = rangeSoFar
		useLookback = rangeSoFar == 0
	)

	parser.Inspect(expr, func(node parser.Node, path []parser.Node) error {
		switch n := node.(type) {
		case *parser.SubqueryExpr:
			useLookback = false
			d := retrieveQueryRange(n.Expr, n.Range+n.OriginalOffset+rangeSoFar)
			if d > queryRange {
				queryRange = d
			}
		case *parser.MatrixSelector:
			useLookback = false
			// nolint
			v := n.VectorSelector.(*parser.VectorSelector)
			if d := v.OriginalOffset + n.Range + rangeSoFar; d > queryRange {
				queryRange = d
			}
		case *parser.VectorSelector:
			if d := n.OriginalOffset + rangeSoFar; d > queryRange {
				queryRange = d
			}
		}

		return nil
	})

	if useLookback {
		// NB: no range provided in any query selectors; use lookback time.
		queryRange += promDefaultLookback
	}

	return queryRange
}

type classificationTags map[string]string

const (
	resultsClassification  = "results_bucket"
	durationClassification = "duration_bucket"
	unclassified           = "unclassified"
)

func newClassificationTags() classificationTags {
	return map[string]string{
		resultsClassification:  unclassified,
		durationClassification: unclassified,
	}
}

// classifyRequest determines the bucket(s) that important request metrics
// will fall into. Supported request metrics at this time are the number of results
// fetched and the time range of the request. The number of buckets and their values
// can be configured by the client. Currently, only query (i.e. query and query_range)
// and label (i.e. label names and label values) endpoints are supported.
func classifyRequest(
	w http.ResponseWriter,
	r *http.Request,
	metrics *classificationMetrics,
	opts Options,
	requestStart time.Time,
	path string,
) classificationTags {
	var (
		cfg  = opts.Metrics.Config
		tags = newClassificationTags()
	)

	// NB(nate): have to check for /label/*/values this way since the URL is templated
	labelValues := strings.Contains(path, "/label/") && strings.Contains(path, "/values")
	if path == route.QueryRangeURL || path == route.QueryURL {
		if opts.Metrics.ParseQueryParams == nil {
			metrics.notClassifiable.Inc(1)
			return tags
		}

		params, err := opts.Metrics.ParseQueryParams(r, requestStart)
		if err != nil || params.Query == "" {
			metrics.badParams.Inc(1)
			return tags
		}

		if tags, err = classifyForQueryEndpoints(
			cfg.QueryEndpointsClassification, params, w, metrics,
		); err != nil {
			opts.InstrumentOpts.Logger().Error(
				"failed to classify query endpoint request", zap.Error(err),
			)
			metrics.error.Inc(1)
			return tags
		}
	} else if path == route.LabelNamesURL || labelValues {
		start, end, err := prometheus.ParseStartAndEnd(r, opts.Metrics.ParseOptions)
		if err != nil {
			metrics.badParams.Inc(1)
			return tags
		}
		if tags, err = classifyForLabelEndpoints(
			cfg.LabelEndpointsClassification, start, end, w, metrics,
		); err != nil {
			opts.InstrumentOpts.Logger().Error(
				"failed to classify label endpoint request", zap.Error(err),
			)
			metrics.error.Inc(1)
			return tags
		}
	}

	return tags
}

func classifyForQueryEndpoints(
	cfg config.QueryClassificationConfig,
	params QueryParams,
	w http.ResponseWriter,
	metrics *classificationMetrics,
) (classificationTags, error) {
	resultsBuckets := cfg.ResultsBuckets
	tags := newClassificationTags()
	if len(resultsBuckets) > 0 {
		fetchedCount := w.Header().Get(headers.FetchedSeriesCount)
		if fetchedCount == "" {
			fetchedCount = "0"
		}
		fetched, err := strconv.Atoi(fetchedCount)
		if err != nil {
			return newClassificationTags(), err
		}

		tags[resultsClassification] = strconv.Itoa(resultsBuckets[0])
		for _, bucket := range resultsBuckets {
			if fetched >= bucket {
				tags[resultsClassification] = strconv.Itoa(bucket)
			}
		}
		metrics.classifiedResult.Inc(1)
	}

	durationBuckets := cfg.DurationBuckets
	if len(durationBuckets) > 0 {
		expr, err := parser.ParseExpr(params.Query)
		if err != nil {
			return newClassificationTags(), err
		}

		queryRange := retrieveQueryRange(expr, 0)
		duration := params.Range()
		totalDuration := duration + queryRange

		tags[durationClassification] = durationBuckets[0].String()
		for _, bucket := range durationBuckets {
			if totalDuration >= bucket {
				tags[durationClassification] = bucket.String()
			}
		}
		metrics.classifiedDuration.Inc(1)
	}

	return tags, nil
}

func classifyForLabelEndpoints(
	cfg config.QueryClassificationConfig,
	start time.Time,
	end time.Time,
	w http.ResponseWriter,
	metrics *classificationMetrics,
) (classificationTags, error) {
	resultsBuckets := cfg.ResultsBuckets
	tags := newClassificationTags()
	if len(resultsBuckets) > 0 {
		fetchedCount := w.Header().Get(headers.FetchedMetadataCount)
		if fetchedCount == "" {
			fetchedCount = "0"
		}
		fetched, err := strconv.Atoi(fetchedCount)
		if err != nil {
			return newClassificationTags(), err
		}

		tags[resultsClassification] = strconv.Itoa(resultsBuckets[0])
		for _, bucket := range resultsBuckets {
			if fetched >= bucket {
				tags[resultsClassification] = strconv.Itoa(bucket)
			}
		}
		metrics.classifiedResult.Inc(1)
	}

	durationBuckets := cfg.DurationBuckets
	if len(durationBuckets) > 0 {
		duration := end.Sub(start)
		tags[durationClassification] = durationBuckets[0].String()
		for _, bucket := range durationBuckets {
			if duration >= bucket {
				tags[durationClassification] = bucket.String()
			}
		}
		metrics.classifiedDuration.Inc(1)
	}

	return tags, nil
}


<|start_filename|>src/integration/resources/inprocess/cluster.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package inprocess

import (
	"errors"
	"fmt"
	"net"
	"strconv"

	aggcfg "github.com/m3db/m3/src/cmd/services/m3aggregator/config"
	dbcfg "github.com/m3db/m3/src/cmd/services/m3dbnode/config"
	coordinatorcfg "github.com/m3db/m3/src/cmd/services/m3query/config"
	"github.com/m3db/m3/src/dbnode/client"
	"github.com/m3db/m3/src/dbnode/discovery"
	"github.com/m3db/m3/src/dbnode/environment"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/integration/resources"
	nettest "github.com/m3db/m3/src/integration/resources/net"
	"github.com/m3db/m3/src/query/storage/m3"
	xconfig "github.com/m3db/m3/src/x/config"
	"github.com/m3db/m3/src/x/config/hostid"
	xerrors "github.com/m3db/m3/src/x/errors"

	"github.com/google/uuid"
	"go.uber.org/zap"
	"gopkg.in/yaml.v2"
)

// ClusterConfigs contain the input config to use for components within
// the cluster. There is one default configuration for each type of component.
// Given a set of ClusterConfigs, the function NewCluster can spin up an m3 cluster.
// Or one can use GenerateClusterSpecification to get the per-instance configuration
// and options based on the given ClusterConfigs.
type ClusterConfigs struct {
	// DBNode is the configuration for db nodes.
	DBNode dbcfg.Configuration
	// Coordinator is the configuration for the coordinator.
	Coordinator coordinatorcfg.Configuration
	// Aggregator is the configuration for aggregators.
	// If Aggregator is nil, the cluster contains only m3coordinator and dbnodes.
	Aggregator *aggcfg.Configuration
}

// ClusterSpecification contain the per-instance configuration and options to use
// for starting each components within the cluster.
// The function NewClusterFromSpecification will spin up an m3 cluster
// with the given ClusterSpecification.
type ClusterSpecification struct {
	// Configs contains the per-instance configuration for all components in the cluster.
	Configs PerInstanceConfigs
	// Options contains the per-insatance options for setting up the cluster.
	Options PerInstanceOptions
}

// PerInstanceConfigs contain the per-instance configuration for all components.
type PerInstanceConfigs struct {
	// DBNodes contains the per-instance configuration for db nodes.
	DBNodes []dbcfg.Configuration
	// Coordinator is the configuration for the coordinator.
	Coordinator coordinatorcfg.Configuration
	// Aggregators is the configuration for aggregators.
	// If Aggregators is nil, the cluster contains only m3coordinator and dbnodes.
	Aggregators []aggcfg.Configuration
}

// PerInstanceOptions contain the per-instance options for setting up the cluster.
type PerInstanceOptions struct {
	// DBNodes contains the per-instance options for db nodes in the cluster.
	DBNode []DBNodeOptions
}

// NewClusterConfigsFromConfigFile creates a new ClusterConfigs object from the
// provided filepaths for dbnode and coordinator configuration.
func NewClusterConfigsFromConfigFile(
	pathToDBNodeCfg string,
	pathToCoordCfg string,
	pathToAggCfg string,
) (ClusterConfigs, error) {
	var dCfg dbcfg.Configuration
	if err := xconfig.LoadFile(&dCfg, pathToDBNodeCfg, xconfig.Options{}); err != nil {
		return ClusterConfigs{}, err
	}

	var cCfg coordinatorcfg.Configuration
	if err := xconfig.LoadFile(&cCfg, pathToCoordCfg, xconfig.Options{}); err != nil {
		return ClusterConfigs{}, err
	}

	var aCfg aggcfg.Configuration
	if pathToAggCfg != "" {
		if err := xconfig.LoadFile(&aCfg, pathToAggCfg, xconfig.Options{}); err != nil {
			return ClusterConfigs{}, err
		}
	}

	return ClusterConfigs{
		DBNode:      dCfg,
		Coordinator: cCfg,
		Aggregator:  &aCfg,
	}, nil
}

// NewClusterConfigsFromYAML creates a new ClusterConfigs object from YAML strings
// representing component configs.
func NewClusterConfigsFromYAML(dbnodeYaml string, coordYaml string, aggYaml string) (ClusterConfigs, error) {
	var dbCfg dbcfg.Configuration
	if err := yaml.Unmarshal([]byte(dbnodeYaml), &dbCfg); err != nil {
		return ClusterConfigs{}, err
	}

	var coordCfg coordinatorcfg.Configuration
	if err := yaml.Unmarshal([]byte(coordYaml), &coordCfg); err != nil {
		return ClusterConfigs{}, err
	}

	var aggCfg aggcfg.Configuration
	if aggYaml != "" {
		if err := yaml.Unmarshal([]byte(aggYaml), &aggCfg); err != nil {
			return ClusterConfigs{}, err
		}
	}

	return ClusterConfigs{
		Coordinator: coordCfg,
		DBNode:      dbCfg,
		Aggregator:  &aggCfg,
	}, nil
}

// NewCluster creates a new M3 cluster based on the ClusterOptions provided.
// Expects at least a coordinator, a dbnode and an aggregator config.
func NewCluster(
	configs ClusterConfigs,
	opts resources.ClusterOptions,
) (resources.M3Resources, error) {
	fullConfigs, err := GenerateClusterSpecification(configs, opts)
	if err != nil {
		return nil, err
	}

	return NewClusterFromSpecification(fullConfigs, opts)
}

// NewClusterFromSpecification creates a new M3 cluster with the given ClusterSpecification.
func NewClusterFromSpecification(
	specs ClusterSpecification,
	opts resources.ClusterOptions,
) (resources.M3Resources, error) {
	if err := opts.Validate(); err != nil {
		return nil, err
	}

	logger, err := resources.NewLogger()
	if err != nil {
		return nil, err
	}

	var (
		coord resources.Coordinator
		nodes = make(resources.Nodes, 0, len(specs.Configs.DBNodes))
		aggs  = make(resources.Aggregators, 0, len(specs.Configs.Aggregators))
	)

	fs.DisableIndexClaimsManagersCheckUnsafe()

	// Ensure that once we start creating resources, they all get cleaned up even if the function
	// fails half way.
	defer func() {
		if err != nil {
			cleanup(logger, nodes, coord, aggs)
		}
	}()

	for i := 0; i < len(specs.Configs.DBNodes); i++ {
		var node resources.Node
		node, err = NewDBNode(specs.Configs.DBNodes[i], specs.Options.DBNode[i])
		if err != nil {
			return nil, err
		}
		nodes = append(nodes, node)
	}

	for _, aggCfg := range specs.Configs.Aggregators {
		var agg resources.Aggregator
		agg, err = NewAggregator(aggCfg, AggregatorOptions{
			GeneratePorts:  true,
			GenerateHostID: false,
		})
		if err != nil {
			return nil, err
		}
		aggs = append(aggs, agg)
	}

	coord, err = NewCoordinator(
		specs.Configs.Coordinator,
		CoordinatorOptions{GeneratePorts: opts.Coordinator.GeneratePorts},
	)
	if err != nil {
		return nil, err
	}

	if err = ConfigurePlacementsForAggregation(nodes, coord, aggs, specs, opts); err != nil {
		return nil, err
	}

	// Start all the configured resources.
	m3 := NewM3Resources(ResourceOptions{
		Coordinator: coord,
		DBNodes:     nodes,
		Aggregators: aggs,
	})
	m3.Start()

	if err = resources.SetupCluster(m3, opts); err != nil {
		return nil, err
	}

	return m3, nil
}

// ConfigurePlacementsForAggregation sets up the correct placement information for
// coordinators and aggregators when aggregation is enabled.
func ConfigurePlacementsForAggregation(
	nodes resources.Nodes,
	coord resources.Coordinator,
	aggs resources.Aggregators,
	specs ClusterSpecification,
	opts resources.ClusterOptions,
) error {
	if len(aggs) == 0 {
		return nil
	}

	coordAPI := coord
	hostDetails, err := coord.HostDetails()
	if err != nil {
		return err
	}

	// With remote aggregation enabled, aggregation is not handled within the coordinator.
	// When this is true, the coordinator will fail to start until placement is updated with
	// aggregation related information. As such, use the coordinator embedded within the dbnode
	// to configure the placement and topics.
	if specs.Configs.Coordinator.Downsample.RemoteAggregator != nil {
		if len(specs.Configs.DBNodes) == 0 ||
			specs.Configs.DBNodes[0].Coordinator == nil {
			return errors.New("remote aggregation requires at least one DB node" +
				" running an embedded coordinator for placement and topic configuration")
		}

		embedded, err := NewEmbeddedCoordinator(nodes[0].(*DBNode))
		if err != nil {
			return nil
		}

		coordAPI = embedded
	} else {
		// TODO(nate): Remove this in a follow up. If we're not doing remote aggregation
		// we should not be starting aggs which is what requires the coordinator to get started.
		// Once we've refactored existing tests that have aggs w/o remote aggregation enabled,
		// this should be killable.
		coord.Start()
	}

	if err = coordAPI.WaitForNamespace(""); err != nil {
		return err
	}

	if err = resources.SetupPlacement(coordAPI, *hostDetails, aggs, *opts.Aggregator); err != nil {
		return err
	}

	aggInstanceInfo, err := aggs[0].HostDetails()
	if err != nil {
		return err
	}

	return resources.SetupM3MsgTopics(coordAPI, *aggInstanceInfo, opts)
}

// GenerateClusterSpecification generates the per-instance configuration and options
// for the cluster set up based on the given input configuation and options.
func GenerateClusterSpecification(
	configs ClusterConfigs,
	opts resources.ClusterOptions,
) (ClusterSpecification, error) {
	if err := opts.Validate(); err != nil {
		return ClusterSpecification{}, err
	}

	nodeCfgs, nodeOpts, envConfig, err := GenerateDBNodeConfigsForCluster(configs, opts.DBNode)
	if err != nil {
		return ClusterSpecification{}, err
	}

	coordConfig := configs.Coordinator
	// TODO(nate): refactor to support having envconfig if no DB.
	if len(coordConfig.Clusters) > 0 {
		coordConfig.Clusters[0].Client.EnvironmentConfig = &envConfig
	} else {
		coordConfig.Clusters = m3.ClustersStaticConfiguration{
			{
				Client: client.Configuration{
					EnvironmentConfig: &envConfig,
				},
			},
		}
	}

	var aggCfgs []aggcfg.Configuration
	if opts.Aggregator != nil {
		aggCfgs, err = GenerateAggregatorConfigsForCluster(configs, opts.Aggregator)
		if err != nil {
			return ClusterSpecification{}, err
		}
	}

	return ClusterSpecification{
		Configs: PerInstanceConfigs{
			DBNodes:     nodeCfgs,
			Coordinator: coordConfig,
			Aggregators: aggCfgs,
		},
		Options: PerInstanceOptions{
			DBNode: nodeOpts,
		},
	}, nil
}

// GenerateDBNodeConfigsForCluster generates the unique configs and options
// for each DB node that will be instantiated. Additionally, provides
// default environment config that can be used to connect to embedded KV
// within the DB nodes.
func GenerateDBNodeConfigsForCluster(
	configs ClusterConfigs,
	opts *resources.DBNodeClusterOptions,
) ([]dbcfg.Configuration, []DBNodeOptions, environment.Configuration, error) {
	if opts == nil {
		return nil, nil, environment.Configuration{}, errors.New("dbnode cluster options is nil")
	}

	var (
		numNodes            = opts.RF * opts.NumInstances
		generatePortsAndIDs = numNodes > 1
	)

	// TODO(nate): eventually support clients specifying their own discovery stanza.
	// Practically, this should cover 99% of cases.
	//
	// Generate a discovery config with the dbnode using the generated hostID marked as
	// the etcd server (i.e. seed node).
	hostID := uuid.NewString()
	defaultDBNodesCfg := configs.DBNode
	discoveryCfg, envConfig, err := generateDefaultDiscoveryConfig(
		defaultDBNodesCfg,
		hostID,
		generatePortsAndIDs)
	if err != nil {
		return nil, nil, environment.Configuration{}, err
	}

	var (
		defaultDBNodeOpts = DBNodeOptions{
			GenerateHostID: generatePortsAndIDs,
			GeneratePorts:  generatePortsAndIDs,
			Start:          true,
		}
		cfgs     = make([]dbcfg.Configuration, 0, numNodes)
		nodeOpts = make([]DBNodeOptions, 0, numNodes)
	)
	for i := 0; i < int(numNodes); i++ {
		var cfg dbcfg.Configuration
		cfg, err = defaultDBNodesCfg.DeepCopy()
		if err != nil {
			return nil, nil, environment.Configuration{}, err
		}
		dbnodeOpts := defaultDBNodeOpts

		if i == 0 {
			// Mark the initial node as the etcd seed node.
			dbnodeOpts.GenerateHostID = false
			cfg.DB.HostID = &hostid.Configuration{
				Resolver: hostid.ConfigResolver,
				Value:    &hostID,
			}
		}
		cfg.DB.Discovery = &discoveryCfg

		cfgs = append(cfgs, cfg)
		nodeOpts = append(nodeOpts, dbnodeOpts)
	}

	return cfgs, nodeOpts, envConfig, nil
}

// generateDefaultDiscoveryConfig handles creating the correct config
// for having an embedded ETCD server with the correct server and
// client configuration.
func generateDefaultDiscoveryConfig(
	cfg dbcfg.Configuration,
	hostID string,
	generateETCDPorts bool,
) (discovery.Configuration, environment.Configuration, error) {
	discoveryConfig := cfg.DB.DiscoveryOrDefault()
	envConfig, err := discoveryConfig.EnvironmentConfig(hostID)
	if err != nil {
		return discovery.Configuration{}, environment.Configuration{}, err
	}

	var (
		etcdClientPort = dbcfg.DefaultEtcdClientPort
		etcdServerPort = dbcfg.DefaultEtcdServerPort
	)
	if generateETCDPorts {
		etcdClientPort, err = nettest.GetAvailablePort()
		if err != nil {
			return discovery.Configuration{}, environment.Configuration{}, err
		}

		etcdServerPort, err = nettest.GetAvailablePort()
		if err != nil {
			return discovery.Configuration{}, environment.Configuration{}, err
		}
	}

	etcdServerURL := fmt.Sprintf("http://0.0.0.0:%d", etcdServerPort)
	etcdClientAddr := net.JoinHostPort("0.0.0.0", strconv.Itoa(etcdClientPort))
	etcdClientURL := fmt.Sprintf("http://0.0.0.0:%d", etcdClientPort)

	envConfig.SeedNodes.InitialCluster[0].Endpoint = etcdServerURL
	envConfig.SeedNodes.InitialCluster[0].HostID = hostID
	envConfig.Services[0].Service.ETCDClusters[0].Endpoints = []string{etcdClientAddr}
	if generateETCDPorts {
		envConfig.SeedNodes.ListenPeerUrls = []string{etcdServerURL}
		envConfig.SeedNodes.ListenClientUrls = []string{etcdClientURL}
		envConfig.SeedNodes.InitialAdvertisePeerUrls = []string{etcdServerURL}
		envConfig.SeedNodes.AdvertiseClientUrls = []string{etcdClientURL}
	}

	configType := discovery.ConfigType
	return discovery.Configuration{
		Type:   &configType,
		Config: &envConfig,
	}, envConfig, nil
}

func cleanup(logger *zap.Logger, nodes resources.Nodes, coord resources.Coordinator, aggs resources.Aggregators) {
	var multiErr xerrors.MultiError
	for _, n := range nodes {
		multiErr = multiErr.Add(n.Close())
	}

	if coord != nil {
		multiErr = multiErr.Add(coord.Close())
	}

	for _, a := range aggs {
		multiErr = multiErr.Add(a.Close())
	}

	if !multiErr.Empty() {
		logger.Warn("failed closing resources", zap.Error(multiErr.FinalError()))
	}
}

// GenerateAggregatorConfigsForCluster generates the unique configs for each aggregator instance.
func GenerateAggregatorConfigsForCluster(
	configs ClusterConfigs,
	opts *resources.AggregatorClusterOptions,
) ([]aggcfg.Configuration, error) {
	if configs.Aggregator == nil {
		return nil, nil
	}

	cfgs := make([]aggcfg.Configuration, 0, int(opts.NumInstances))
	for i := 0; i < int(opts.NumInstances); i++ {
		cfg, err := configs.Aggregator.DeepCopy()
		if err != nil {
			return nil, err
		}

		hostID := fmt.Sprintf("m3aggregator%02d", i)
		aggCfg := cfg.AggregatorOrDefault()
		aggCfg.HostID = &hostid.Configuration{
			Resolver: hostid.ConfigResolver,
			Value:    &hostID,
		}
		cfg.Aggregator = &aggCfg

		cfgs = append(cfgs, cfg)
	}

	return cfgs, nil
}


<|start_filename|>src/cluster/placement/compress.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package placement

import (
	"github.com/m3db/m3/src/cluster/generated/proto/placementpb"

	"github.com/gogo/protobuf/proto"
	"github.com/klauspost/compress/zstd"
)

const _decoderInitialBufferSize = 131072

func compressPlacementProto(p *placementpb.Placement) ([]byte, error) {
	if p == nil {
		return nil, errNilPlacementSnapshotsProto
	}

	w, err := zstd.NewWriter(
		nil,
		zstd.WithEncoderCRC(true),
		zstd.WithEncoderConcurrency(1),
		zstd.WithEncoderLevel(zstd.SpeedBestCompression))
	if err != nil {
		return nil, err
	}
	defer w.Close() // nolint:errcheck

	uncompressed, _ := p.Marshal()

	compressed := make([]byte, 0, len(uncompressed)/2) // prealloc, assuming at least 50% space savings
	compressed = w.EncodeAll(uncompressed, compressed)

	return compressed, nil
}

func decompressPlacementProto(compressed []byte) (*placementpb.Placement, error) {
	if compressed == nil {
		return nil, errNilValue
	}

	decompressed := make([]byte, 0, _decoderInitialBufferSize)
	r, err := zstd.NewReader(
		nil,
		zstd.WithDecoderConcurrency(1),
		zstd.WithDecoderLowmem(false),
	)
	if err != nil {
		return nil, err
	}
	defer r.Close() // nolint:errcheck

	decompressed, err = r.DecodeAll(compressed, decompressed)
	if err != nil {
		return nil, err
	}

	result := &placementpb.Placement{}
	if err := proto.Unmarshal(decompressed, result); err != nil {
		return nil, err
	}

	return result, nil
}


<|start_filename|>src/query/storage/promremote/query_converter_test.go<|end_filename|>
// Copyright (c) 2021  Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package promremote

import (
	"testing"
	"time"

	"github.com/m3db/m3/src/query/models"
	"github.com/m3db/m3/src/query/storage"
	"github.com/m3db/m3/src/query/ts"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/prometheus/prometheus/prompb"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestWriteQueryConverter(t *testing.T) {
	now := xtime.Now()
	dp := ts.Datapoint{
		Timestamp: now,
		Value:     42,
	}
	tag := models.Tag{
		Name:  []byte("test_tag_name"),
		Value: []byte("test_tag_value"),
	}
	convertedToLabel := prompb.Label{
		Name:  "test_tag_name",
		Value: "test_tag_value",
	}
	covertedToSample := prompb.Sample{
		Timestamp: now.ToNormalizedTime(time.Millisecond),
		Value:     42,
	}

	tcs := []struct {
		name     string
		input    storage.WriteQueryOptions
		expected *prompb.WriteRequest
	}{
		{
			name: "single datapoint",
			input: storage.WriteQueryOptions{
				Tags: models.Tags{
					Opts: models.NewTagOptions(),
					Tags: []models.Tag{tag},
				},
				Datapoints: ts.Datapoints{dp},
				Unit:       xtime.Millisecond,
			},
			expected: promWriteRequest(prompb.TimeSeries{
				Labels:  []prompb.Label{convertedToLabel},
				Samples: []prompb.Sample{covertedToSample},
			}),
		},
		{
			name: "duplicate tags and samples",
			input: storage.WriteQueryOptions{
				Tags: models.Tags{
					Opts: models.NewTagOptions().SetAllowTagNameDuplicates(true),
					Tags: []models.Tag{tag, tag},
				},
				Datapoints: ts.Datapoints{dp, dp},
				Unit:       xtime.Millisecond,
			},
			expected: promWriteRequest(prompb.TimeSeries{
				Labels:  []prompb.Label{convertedToLabel, convertedToLabel},
				Samples: []prompb.Sample{covertedToSample, covertedToSample},
			}),
		},
		{
			name: "overrides metric name tag",
			input: storage.WriteQueryOptions{
				Tags: models.Tags{
					Opts: models.NewTagOptions().SetMetricName(tag.Name),
					Tags: []models.Tag{tag},
				},
				Datapoints: ts.Datapoints{dp},
				Unit:       xtime.Millisecond,
			},
			expected: promWriteRequest(prompb.TimeSeries{
				Labels: []prompb.Label{{
					Name:  "__name__",
					Value: convertedToLabel.Value,
				}},
				Samples: []prompb.Sample{covertedToSample},
			}),
		},
		{
			name: "overrides bucket name name tag",
			input: storage.WriteQueryOptions{
				Tags: models.Tags{
					Opts: models.NewTagOptions().SetBucketName(tag.Name),
					Tags: []models.Tag{tag},
				},
				Datapoints: ts.Datapoints{dp},
				Unit:       xtime.Millisecond,
			},
			expected: promWriteRequest(prompb.TimeSeries{
				Labels: []prompb.Label{{
					Name:  "le",
					Value: convertedToLabel.Value,
				}},
				Samples: []prompb.Sample{covertedToSample},
			}),
		},
	}

	for _, tc := range tcs {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			q, err := storage.NewWriteQuery(tc.input)
			require.NoError(t, err)
			assert.Equal(t, tc.expected, convertWriteQuery(q))
		})
	}
}

func TestConvertQueryNil(t *testing.T) {
	assert.Nil(t, convertWriteQuery(nil))
}

func TestEncodeWriteQuery(t *testing.T) {
	data, err := convertAndEncodeWriteQuery(nil)
	require.Error(t, err)
	assert.Len(t, data, 0)
	assert.Contains(t, err.Error(), "received nil query")
}

func promWriteRequest(ts prompb.TimeSeries) *prompb.WriteRequest {
	return &prompb.WriteRequest{Timeseries: []prompb.TimeSeries{ts}}
}


<|start_filename|>src/dbnode/integration/repair_force_only_compare_test.go<|end_filename|>
//go:build integration
// +build integration

// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package integration

import (
	"testing"
	"time"

	"github.com/m3db/m3/src/dbnode/integration/generate"
	"github.com/m3db/m3/src/dbnode/namespace"
	"github.com/m3db/m3/src/dbnode/retention"
	"github.com/m3db/m3/src/dbnode/storage/repair"
	xtest "github.com/m3db/m3/src/x/test"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/stretchr/testify/require"
	"github.com/uber-go/tally"
)

func TestRepairForceAndOnlyCompare(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Test both disjoint and shared series repair.
	genRepairData := genRepairDatafn(func(now xtime.UnixNano, blockSize time.Duration) (
		node0Data generate.SeriesBlocksByStart,
		node1Data generate.SeriesBlocksByStart,
		node2Data generate.SeriesBlocksByStart,
		allData generate.SeriesBlocksByStart,
	) {
		currBlockStart := now.Truncate(blockSize)
		node0Data = generate.BlocksByStart([]generate.BlockConfig{
			{IDs: []string{"foo"}, NumPoints: 90, Start: currBlockStart.Add(-4 * blockSize)},
			{IDs: []string{"foo", "baz"}, NumPoints: 90, Start: currBlockStart.Add(-3 * blockSize)},
		})
		node1Data = generate.BlocksByStart([]generate.BlockConfig{
			{IDs: []string{"bar"}, NumPoints: 90, Start: currBlockStart.Add(-4 * blockSize)},
			{IDs: []string{"foo", "baz"}, NumPoints: 90, Start: currBlockStart.Add(-3 * blockSize)},
		})

		allData = make(map[xtime.UnixNano]generate.SeriesBlock)
		for start, data := range node0Data {
			allData[start] = append(allData[start], data...)
		}
		for start, data := range node1Data {
			allData[start] = append(allData[start], data...)
		}
		for start, data := range node2Data {
			allData[start] = append(allData[start], data...)
		}

		return node0Data, node1Data, node2Data, allData
	})

	// Test setups.
	log := xtest.NewLogger(t)
	retentionOpts := retention.NewOptions().
		SetRetentionPeriod(20 * time.Hour).
		SetBlockSize(2 * time.Hour).
		SetBufferPast(10 * time.Minute).
		SetBufferFuture(2 * time.Minute)
	nsOpts := namespace.NewOptions().
		// Test needing to force enable repairs.
		SetRepairEnabled(false).
		SetRetentionOptions(retentionOpts)
	namesp, err := namespace.NewMetadata(testNamespaces[0], nsOpts)
	require.NoError(t, err)
	opts := NewTestOptions(t).
		SetNamespaces([]namespace.Metadata{namesp}).
		// Use TChannel clients for writing / reading because we want to target individual nodes at a time
		// and not write/read all nodes in the cluster.
		SetUseTChannelClientForWriting(true).
		SetUseTChannelClientForReading(true)

	setupOpts := []BootstrappableTestSetupOptions{
		{
			DisablePeersBootstrapper: true,
			EnableRepairs:            true,
			// Test forcing repair of type compare only repair.
			ForceRepairs: true,
			RepairType:   repair.OnlyCompareRepair,
		},
		{
			DisablePeersBootstrapper: true,
			EnableRepairs:            true,
			// Test forcing repair of type compare only repair.
			ForceRepairs: true,
			RepairType:   repair.OnlyCompareRepair,
		},
		{
			DisablePeersBootstrapper: true,
			EnableRepairs:            true,
			// Test forcing repair of type compare only repair.
			ForceRepairs: true,
			RepairType:   repair.OnlyCompareRepair,
		},
	}

	// nolint: govet
	setups, closeFn := NewDefaultBootstrappableTestSetups(t, opts, setupOpts)
	defer closeFn()

	// Ensure that the current time is set such that the previous block is flushable.
	blockSize := retentionOpts.BlockSize()
	now := setups[0].NowFn()().Truncate(blockSize).Add(retentionOpts.BufferPast()).Add(time.Second)
	for _, setup := range setups {
		setup.SetNowFn(now)
	}

	node0Data, node1Data, node2Data, _ := genRepairData(now, blockSize)
	if node0Data != nil {
		require.NoError(t, writeTestDataToDisk(namesp, setups[0], node0Data, 0))
	}
	if node1Data != nil {
		require.NoError(t, writeTestDataToDisk(namesp, setups[1], node1Data, 0))
	}
	if node2Data != nil {
		require.NoError(t, writeTestDataToDisk(namesp, setups[2], node2Data, 0))
	}

	// Start the servers with filesystem bootstrappers.
	setups.parallel(func(s TestSetup) {
		if err := s.StartServer(); err != nil {
			panic(err)
		}
	})
	log.Debug("servers are now up")

	// Stop the servers.
	defer func() {
		setups.parallel(func(s TestSetup) {
			require.NoError(t, s.StopServer())
		})
		log.Debug("servers are now down")
	}()

	// Wait for repairs to occur at least once per node.
	log.Debug("waiting for repairs to run")
	var runSuccessPerNodeCounters []tally.CounterSnapshot
	require.True(t, waitUntil(func() bool {
		var successCounters []tally.CounterSnapshot
		for _, setup := range setups {
			scope := setup.Scope()
			for _, v := range scope.Snapshot().Counters() {
				if v.Name() != "repair.run" {
					continue
				}
				repairType, ok := v.Tags()["repair_type"]
				if !ok || repairType != "only_compare" {
					continue
				}
				if v.Value() > 0 {
					successCounters = append(successCounters, v)
					break
				}
			}
		}

		// Check if all counters are success.
		successAll := len(successCounters) == len(setups)
		if successAll {
			runSuccessPerNodeCounters = successCounters
			return true
		}
		return false
	}, 60*time.Second))

	// Verify that the repair runs only ran comparisons without repairing data.
	log.Debug("verifying repairs that ran")
	require.Equal(t, len(setups), len(runSuccessPerNodeCounters),
		"unexpected number of successful nodes ran repairs")
	for _, counter := range runSuccessPerNodeCounters {
		repairType, ok := counter.Tags()["repair_type"]
		require.True(t, ok)
		require.Equal(t, "only_compare", repairType)
		require.True(t, counter.Value() > 0)
	}

	// Verify data did not change (repair type is compare only).
	verifySeriesMaps(t, setups[0], namesp.ID(), node0Data)
	verifySeriesMaps(t, setups[1], namesp.ID(), node1Data)
	verifySeriesMaps(t, setups[2], namesp.ID(), node2Data)
}


<|start_filename|>src/cmd/tools/split_shards/main/main.go<|end_filename|>
// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	iofs "io/fs"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/m3db/m3/src/dbnode/persist"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/dbnode/sharding"
	xerrors "github.com/m3db/m3/src/x/errors"
	"github.com/m3db/m3/src/x/ident"
	xtime "github.com/m3db/m3/src/x/time"

	"github.com/pborman/getopt"
	"go.uber.org/zap"
)

var checkpointPattern = regexp.MustCompile(`/data/(\w+)/([0-9]+)/fileset-([0-9]+)-([0-9]+)-checkpoint.db$`)

func main() {
	var (
		optSrcPath       = getopt.StringLong("src-path", 's', "", "Source path [e.g. /temp/lib/m3db/data]")
		optDstPathPrefix = getopt.StringLong("dst-path", 'd', "", "Destination path prefix [e.g. /var/lib/m3db/data]")
		optBlockUntil    = getopt.Int64Long("block-until", 'b', 0, "Block Until Time, exclusive [in nsec]")
		optShards        = getopt.Uint32Long("src-shards", 'h', 0, "Original (source) number of shards")
		optFactor        = getopt.IntLong("factor", 'f', 0, "Integer factor to increase the number of shards by")
	)
	getopt.Parse()

	rawLogger, err := zap.NewDevelopment()
	if err != nil {
		log.Fatalf("unable to create logger: %+v", err)
	}
	logger := rawLogger.Sugar()

	if *optSrcPath == "" ||
		*optDstPathPrefix == "" ||
		*optSrcPath == *optDstPathPrefix ||
		*optBlockUntil <= 0 ||
		*optShards == 0 ||
		*optFactor <= 0 {
		getopt.Usage()
		os.Exit(1)
	}

	var (
		srcFilesetLocation = dropDataSuffix(*optSrcPath)
		dstFilesetLocation = dropDataSuffix(*optDstPathPrefix)

		srcFsOpts = fs.NewOptions().SetFilePathPrefix(srcFilesetLocation)
		dstFsOpts = fs.NewOptions().SetFilePathPrefix(dstFilesetLocation)
	)

	// Not using bytes pool with streaming reads/writes to avoid the fixed memory overhead.
	srcReader, err := fs.NewReader(nil, srcFsOpts)
	if err != nil {
		logger.Fatalf("could not create srcReader: %+v", err)
	}

	dstReaders := make([]fs.DataFileSetReader, *optFactor)
	dstWriters := make([]fs.StreamingWriter, *optFactor)
	for i := range dstWriters {
		dstReaders[i], err = fs.NewReader(nil, dstFsOpts)
		if err != nil {
			logger.Fatalf("could not create dstReader, %+v", err)
		}
		dstWriters[i], err = fs.NewStreamingWriter(dstFsOpts)
		if err != nil {
			logger.Fatalf("could not create writer, %+v", err)
		}
	}

	hashFn := sharding.DefaultHashFn(int(*optShards) * (*optFactor))

	start := time.Now()

	if err := filepath.WalkDir(*optSrcPath, func(path string, d iofs.DirEntry, err error) error {
		if err != nil || d.IsDir() || !strings.HasSuffix(d.Name(), "-checkpoint.db") {
			return err
		}
		fmt.Printf("%s - %s\n", time.Now().Local(), path) // nolint: forbidigo
		pathParts := checkpointPattern.FindStringSubmatch(path)
		if len(pathParts) != 5 {
			return fmt.Errorf("failed to parse path %s", path)
		}

		var (
			namespace        = pathParts[1]
			shard, err2      = strconv.Atoi(pathParts[2])
			blockStart, err3 = strconv.Atoi(pathParts[3])
			volume, err4     = strconv.Atoi(pathParts[4])
		)
		if err = xerrors.FirstError(err2, err3, err4); err != nil {
			return err
		}

		if blockStart >= int(*optBlockUntil) {
			fmt.Println(" - skip (too recent)") // nolint: forbidigo
			return nil
		}

		if err = splitFileSet(
			srcReader, dstWriters, hashFn, *optShards, *optFactor, namespace, uint32(shard),
			xtime.UnixNano(blockStart), volume); err != nil {
			if strings.Contains(err.Error(), "no such file or directory") {
				fmt.Println(" - skip (incomplete fileset)") // nolint: forbidigo
				return nil
			}
			return err
		}

		err = verifySplitShards(
			srcReader, dstReaders, hashFn, *optShards, namespace, uint32(shard),
			xtime.UnixNano(blockStart), volume)
		if err != nil && strings.Contains(err.Error(), "no such file or directory") {
			return nil
		}
		return err
	}); err != nil {
		logger.Fatalf("unable to walk the source dir: %+v", err)
	}

	runTime := time.Since(start)
	fmt.Printf("Running time: %s\n", runTime) // nolint: forbidigo
}

func splitFileSet(
	srcReader fs.DataFileSetReader,
	dstWriters []fs.StreamingWriter,
	hashFn sharding.HashFn,
	srcNumShards uint32,
	factor int,
	namespace string,
	srcShard uint32,
	blockStart xtime.UnixNano,
	volume int,
) error {
	if srcShard >= srcNumShards {
		return fmt.Errorf("unexpected source shard ID %d (must be under %d)", srcShard, srcNumShards)
	}

	readOpts := fs.DataReaderOpenOptions{
		Identifier: fs.FileSetFileIdentifier{
			Namespace:   ident.StringID(namespace),
			Shard:       srcShard,
			BlockStart:  blockStart,
			VolumeIndex: volume,
		},
		FileSetType:      persist.FileSetFlushType,
		StreamingEnabled: true,
	}

	err := srcReader.Open(readOpts)
	if err != nil {
		return fmt.Errorf("unable to open srcReader: %w", err)
	}

	plannedRecordsCount := uint(srcReader.Entries() / factor)
	if plannedRecordsCount == 0 {
		plannedRecordsCount = 1
	}

	for i := range dstWriters {
		writeOpts := fs.StreamingWriterOpenOptions{
			NamespaceID:         ident.StringID(namespace),
			ShardID:             mapToDstShard(srcNumShards, i, srcShard),
			BlockStart:          blockStart,
			BlockSize:           srcReader.Status().BlockSize,
			VolumeIndex:         volume + 1,
			PlannedRecordsCount: plannedRecordsCount,
		}
		if err := dstWriters[i].Open(writeOpts); err != nil {
			return fmt.Errorf("unable to open dstWriters[%d]: %w", i, err)
		}
	}

	dataHolder := make([][]byte, 1)
	for {
		entry, err := srcReader.StreamingRead()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return fmt.Errorf("read error: %w", err)
		}

		newShardID := hashFn(entry.ID)
		if newShardID%srcNumShards != srcShard {
			return fmt.Errorf("mismatched shards, %d to %d", srcShard, newShardID)
		}
		writer := dstWriters[newShardID/srcNumShards]

		dataHolder[0] = entry.Data
		if err := writer.WriteAll(entry.ID, entry.EncodedTags, dataHolder, entry.DataChecksum); err != nil {
			return err
		}
	}

	for i := range dstWriters {
		if err := dstWriters[i].Close(); err != nil {
			return err
		}
	}

	return srcReader.Close()
}

func verifySplitShards(
	srcReader fs.DataFileSetReader,
	dstReaders []fs.DataFileSetReader,
	hashFn sharding.HashFn,
	srcNumShards uint32,
	namespace string,
	srcShard uint32,
	blockStart xtime.UnixNano,
	volume int,
) error {
	if srcShard >= srcNumShards {
		return fmt.Errorf("unexpected source shard ID %d (must be under %d)", srcShard, srcNumShards)
	}

	srcReadOpts := fs.DataReaderOpenOptions{
		Identifier: fs.FileSetFileIdentifier{
			Namespace:   ident.StringID(namespace),
			Shard:       srcShard,
			BlockStart:  blockStart,
			VolumeIndex: volume,
		},
		FileSetType:      persist.FileSetFlushType,
		StreamingEnabled: true,
	}

	err := srcReader.Open(srcReadOpts)
	if err != nil {
		return fmt.Errorf("unable to open srcReader: %w", err)
	}

	dstEntries := 0
	for i := range dstReaders {
		dstReadOpts := fs.DataReaderOpenOptions{
			Identifier: fs.FileSetFileIdentifier{
				Namespace:   ident.StringID(namespace),
				Shard:       mapToDstShard(srcNumShards, i, srcShard),
				BlockStart:  blockStart,
				VolumeIndex: volume + 1,
			},
			FileSetType:      persist.FileSetFlushType,
			StreamingEnabled: true,
		}
		if err := dstReaders[i].Open(dstReadOpts); err != nil {
			return fmt.Errorf("unable to open dstReaders[%d]: %w", i, err)
		}
		dstEntries += dstReaders[i].Entries()
	}

	if srcReader.Entries() != dstEntries {
		return fmt.Errorf("entry count mismatch: src %d != dst %d", srcReader.Entries(), dstEntries)
	}

	for {
		srcEntry, err := srcReader.StreamingReadMetadata()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return fmt.Errorf("src read error: %w", err)
		}

		newShardID := hashFn(srcEntry.ID)
		if newShardID%srcNumShards != srcShard {
			return fmt.Errorf("mismatched shards, %d to %d", srcShard, newShardID)
		}
		dstReader := dstReaders[newShardID/srcNumShards]

		// Using StreamingRead() on destination filesets here because it also verifies data checksums.
		dstEntry, err := dstReader.StreamingRead()
		if err != nil {
			return fmt.Errorf("dst read error: %w", err)
		}

		if !bytes.Equal(srcEntry.ID, dstEntry.ID) {
			return fmt.Errorf("ID mismatch: %s != %s", srcEntry.ID, dstEntry.ID)
		}
		if !bytes.Equal(srcEntry.EncodedTags, dstEntry.EncodedTags) {
			return fmt.Errorf("EncodedTags mismatch: %s != %s", srcEntry.EncodedTags, dstEntry.EncodedTags)
		}
		if srcEntry.DataChecksum != dstEntry.DataChecksum {
			return fmt.Errorf("data checksum mismatch: %d != %d, id=%s",
				srcEntry.DataChecksum, dstEntry.DataChecksum, srcEntry.ID)
		}
	}

	for i := range dstReaders {
		dstReader := dstReaders[i]
		if _, err := dstReader.StreamingReadMetadata(); !errors.Is(err, io.EOF) {
			return fmt.Errorf("expected EOF on split shard %d, but got %w",
				dstReader.Status().Shard, err)
		}
		if err := dstReader.Close(); err != nil {
			return err
		}
	}

	return srcReader.Close()
}

func mapToDstShard(srcNumShards uint32, i int, srcShard uint32) uint32 {
	return srcNumShards*uint32(i) + srcShard
}

func dropDataSuffix(path string) string {
	dataIdx := strings.LastIndex(path, "/data")
	if dataIdx < 0 {
		return path
	}
	return path[:dataIdx]
}
