github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/storage/namespace_test.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package storage
    22  
    23  import (
    24  	stdlibctx "context"
    25  	"errors"
    26  	"fmt"
    27  	"sync"
    28  	"testing"
    29  	"time"
    30  
    31  	"github.com/m3db/m3/src/cluster/shard"
    32  	"github.com/m3db/m3/src/dbnode/namespace"
    33  	"github.com/m3db/m3/src/dbnode/persist"
    34  	"github.com/m3db/m3/src/dbnode/persist/fs"
    35  	"github.com/m3db/m3/src/dbnode/retention"
    36  	"github.com/m3db/m3/src/dbnode/runtime"
    37  	"github.com/m3db/m3/src/dbnode/sharding"
    38  	"github.com/m3db/m3/src/dbnode/storage/bootstrap"
    39  	"github.com/m3db/m3/src/dbnode/storage/bootstrap/result"
    40  	"github.com/m3db/m3/src/dbnode/storage/index"
    41  	"github.com/m3db/m3/src/dbnode/storage/index/convert"
    42  	"github.com/m3db/m3/src/dbnode/storage/repair"
    43  	"github.com/m3db/m3/src/dbnode/storage/series"
    44  	"github.com/m3db/m3/src/dbnode/tracepoint"
    45  	xmetrics "github.com/m3db/m3/src/dbnode/x/metrics"
    46  	xidx "github.com/m3db/m3/src/m3ninx/idx"
    47  	idxpersist "github.com/m3db/m3/src/m3ninx/persist"
    48  	"github.com/m3db/m3/src/x/context"
    49  	xerrors "github.com/m3db/m3/src/x/errors"
    50  	"github.com/m3db/m3/src/x/ident"
    51  	"github.com/m3db/m3/src/x/instrument"
    52  	xtest "github.com/m3db/m3/src/x/test"
    53  	xtime "github.com/m3db/m3/src/x/time"
    54  
    55  	"github.com/fortytw2/leaktest"
    56  	"github.com/golang/mock/gomock"
    57  	"github.com/opentracing/opentracing-go"
    58  	"github.com/opentracing/opentracing-go/mocktracer"
    59  	"github.com/stretchr/testify/assert"
    60  	"github.com/stretchr/testify/require"
    61  	"github.com/uber-go/tally"
    62  )
    63  
    64  var (
    65  	sourceNsID = ident.StringID("source")
    66  	targetNsID = ident.StringID("target")
    67  
    68  	sourceBlockSize = time.Hour
    69  	targetBlockSize = 2 * time.Hour
    70  
    71  	insOpts = instrument.NewOptions()
    72  
    73  	testShardIDs = sharding.NewShards([]uint32{0, 1}, shard.Available)
    74  )
    75  
    76  type closerFn func()
    77  
    78  func newTestNamespace(t *testing.T) (*dbNamespace, closerFn) {
    79  	return newTestNamespaceWithIDOpts(t, defaultTestNs1ID, defaultTestNs1Opts)
    80  }
    81  
    82  func newTestNamespaceMetadata(t *testing.T) namespace.Metadata {
    83  	return newTestNamespaceMetadataWithIDOpts(t, defaultTestNs1ID, defaultTestNs1Opts)
    84  }
    85  
    86  func newTestNamespaceMetadataWithIDOpts(
    87  	t *testing.T,
    88  	nsID ident.ID,
    89  	opts namespace.Options,
    90  ) namespace.Metadata {
    91  	metadata, err := namespace.NewMetadata(nsID, opts)
    92  	require.NoError(t, err)
    93  	return metadata
    94  }
    95  
    96  func newTestNamespaceWithIDOpts(
    97  	t *testing.T,
    98  	nsID ident.ID,
    99  	opts namespace.Options,
   100  ) (*dbNamespace, closerFn) {
   101  	metadata := newTestNamespaceMetadataWithIDOpts(t, nsID, opts)
   102  	hashFn := func(identifier ident.ID) uint32 { return testShardIDs[0].ID() }
   103  	shardSet, err := sharding.NewShardSet(testShardIDs, hashFn)
   104  	require.NoError(t, err)
   105  	dopts := DefaultTestOptions().SetRuntimeOptionsManager(runtime.NewOptionsManager())
   106  	ns, err := newDatabaseNamespace(metadata,
   107  		namespace.NewRuntimeOptionsManager(metadata.ID().String()),
   108  		shardSet, nil, nil, nil, dopts)
   109  	require.NoError(t, err)
   110  	closer := dopts.RuntimeOptionsManager().Close
   111  	return ns.(*dbNamespace), closer
   112  }
   113  
   114  func newTestNamespaceWithOpts(
   115  	t *testing.T,
   116  	opts namespace.Options,
   117  	dopts Options,
   118  ) (*dbNamespace, closerFn) {
   119  	nsID := defaultTestNs1ID
   120  	metadata := newTestNamespaceMetadataWithIDOpts(t, nsID, opts)
   121  	hashFn := func(identifier ident.ID) uint32 { return testShardIDs[0].ID() }
   122  	shardSet, err := sharding.NewShardSet(testShardIDs, hashFn)
   123  	require.NoError(t, err)
   124  	ns, err := newDatabaseNamespace(metadata,
   125  		namespace.NewRuntimeOptionsManager(metadata.ID().String()),
   126  		shardSet, nil, nil, nil, dopts)
   127  	require.NoError(t, err)
   128  	closer := dopts.RuntimeOptionsManager().Close
   129  	return ns.(*dbNamespace), closer
   130  }
   131  
   132  func newTestNamespaceWithIndex(
   133  	t *testing.T,
   134  	index NamespaceIndex,
   135  ) (*dbNamespace, closerFn) {
   136  	ns, closer := newTestNamespace(t)
   137  	if index != nil {
   138  		ns.reverseIndex = index
   139  	}
   140  	return ns, closer
   141  }
   142  
   143  func newTestNamespaceWithTruncateType(
   144  	t *testing.T,
   145  	index NamespaceIndex,
   146  	truncateType series.TruncateType,
   147  ) (*dbNamespace, closerFn) {
   148  	opts := DefaultTestOptions().
   149  		SetRuntimeOptionsManager(runtime.NewOptionsManager()).
   150  		SetTruncateType(truncateType)
   151  
   152  	ns, closer := newTestNamespaceWithOpts(t, defaultTestNs1Opts, opts)
   153  	ns.reverseIndex = index
   154  	return ns, closer
   155  }
   156  
   157  func TestNamespaceName(t *testing.T) {
   158  	ns, closer := newTestNamespace(t)
   159  	defer closer()
   160  	require.True(t, defaultTestNs1ID.Equal(ns.ID()))
   161  }
   162  
   163  func TestNamespaceTick(t *testing.T) {
   164  	ctrl := xtest.NewController(t)
   165  	defer ctrl.Finish()
   166  
   167  	ns, closer := newTestNamespace(t)
   168  	defer closer()
   169  	for i := range testShardIDs {
   170  		shard := NewMockdatabaseShard(ctrl)
   171  		shard.EXPECT().Tick(context.NewNoOpCanncellable(), gomock.Any(), gomock.Any()).Return(tickResult{}, nil)
   172  		ns.shards[testShardIDs[i].ID()] = shard
   173  	}
   174  
   175  	// Only asserting the expected methods are called
   176  	require.NoError(t, ns.Tick(context.NewNoOpCanncellable(), xtime.Now()))
   177  }
   178  
   179  func TestNamespaceTickError(t *testing.T) {
   180  	ctrl := xtest.NewController(t)
   181  	defer ctrl.Finish()
   182  
   183  	fakeErr := errors.New("fake error")
   184  	ns, closer := newTestNamespace(t)
   185  	defer closer()
   186  
   187  	for i := range testShardIDs {
   188  		shard := NewMockdatabaseShard(ctrl)
   189  		if i == 0 {
   190  			shard.EXPECT().Tick(context.NewNoOpCanncellable(), gomock.Any(), gomock.Any()).Return(tickResult{}, fakeErr)
   191  		} else {
   192  			shard.EXPECT().Tick(context.NewNoOpCanncellable(), gomock.Any(), gomock.Any()).Return(tickResult{}, nil)
   193  		}
   194  		ns.shards[testShardIDs[i].ID()] = shard
   195  	}
   196  
   197  	err := ns.Tick(context.NewNoOpCanncellable(), xtime.Now())
   198  	require.NotNil(t, err)
   199  	require.Equal(t, fakeErr.Error(), err.Error())
   200  }
   201  
   202  func TestNamespaceWriteShardNotOwned(t *testing.T) {
   203  	ctx := context.NewBackground()
   204  	defer ctx.Close()
   205  
   206  	ns, closer := newTestNamespace(t)
   207  	defer closer()
   208  	for i := range ns.shards {
   209  		ns.shards[i] = nil
   210  	}
   211  	now := xtime.Now()
   212  	seriesWrite, err := ns.Write(ctx, ident.StringID("foo"), now, 0.0, xtime.Second, nil)
   213  	require.Error(t, err)
   214  	require.True(t, xerrors.IsRetryableError(err))
   215  	require.Equal(t, "not responsible for shard 0", err.Error())
   216  	require.False(t, seriesWrite.WasWritten)
   217  }
   218  
   219  func TestNamespaceReadOnlyRejectWrites(t *testing.T) {
   220  	ctx := context.NewBackground()
   221  	defer ctx.Close()
   222  
   223  	ns, closer := newTestNamespace(t)
   224  	defer closer()
   225  
   226  	ns.SetReadOnly(true)
   227  
   228  	id := ident.StringID("foo")
   229  	now := xtime.Now()
   230  
   231  	seriesWrite, err := ns.Write(ctx, id, now, 0, xtime.Second, nil)
   232  	require.EqualError(t, err, errNamespaceReadOnly.Error())
   233  	require.False(t, seriesWrite.WasWritten)
   234  
   235  	seriesWrite, err = ns.WriteTagged(ctx, id, convert.EmptyTagMetadataResolver, now, 0, xtime.Second, nil)
   236  	require.EqualError(t, err, errNamespaceReadOnly.Error())
   237  	require.False(t, seriesWrite.WasWritten)
   238  }
   239  
   240  func TestNamespaceWriteShardOwned(t *testing.T) {
   241  	ctrl := xtest.NewController(t)
   242  	defer ctrl.Finish()
   243  
   244  	ctx := context.NewBackground()
   245  	defer ctx.Close()
   246  
   247  	id := ident.StringID("foo")
   248  	now := xtime.Now()
   249  	val := 0.0
   250  	unit := xtime.Second
   251  	ant := []byte(nil)
   252  
   253  	truncateTypes := []series.TruncateType{series.TypeBlock, series.TypeNone}
   254  	for _, truncateType := range truncateTypes {
   255  		ns, closer := newTestNamespaceWithTruncateType(t, nil, truncateType)
   256  		defer closer()
   257  		shard := NewMockdatabaseShard(ctrl)
   258  		opts := series.WriteOptions{
   259  			TruncateType: truncateType,
   260  		}
   261  		shard.EXPECT().Write(ctx, id, now, val, unit, ant, opts).
   262  			Return(SeriesWrite{WasWritten: true}, nil).Times(1)
   263  		shard.EXPECT().Write(ctx, id, now, val, unit, ant, opts).
   264  			Return(SeriesWrite{WasWritten: false}, nil).Times(1)
   265  
   266  		ns.shards[testShardIDs[0].ID()] = shard
   267  
   268  		seriesWrite, err := ns.Write(ctx, id, now, val, unit, ant)
   269  		require.NoError(t, err)
   270  		require.True(t, seriesWrite.WasWritten)
   271  
   272  		seriesWrite, err = ns.Write(ctx, id, now, val, unit, ant)
   273  		require.NoError(t, err)
   274  		require.False(t, seriesWrite.WasWritten)
   275  	}
   276  }
   277  
   278  func TestNamespaceReadEncodedShardNotOwned(t *testing.T) {
   279  	ctx := context.NewBackground()
   280  	defer ctx.Close()
   281  
   282  	ns, closer := newTestNamespace(t)
   283  	defer closer()
   284  
   285  	for i := range ns.shards {
   286  		ns.shards[i] = nil
   287  	}
   288  
   289  	now := xtime.Now()
   290  	_, err := ns.ReadEncoded(ctx, ident.StringID("foo"), now, now)
   291  	require.Error(t, err)
   292  }
   293  
   294  func TestNamespaceReadEncodedShardOwned(t *testing.T) {
   295  	ctrl := xtest.NewController(t)
   296  	defer ctrl.Finish()
   297  
   298  	ctx := context.NewBackground()
   299  	defer ctx.Close()
   300  
   301  	id := ident.StringID("foo")
   302  	start := xtime.Now()
   303  	end := start.Add(time.Second)
   304  
   305  	ns, closer := newTestNamespace(t)
   306  	defer closer()
   307  
   308  	shard := NewMockdatabaseShard(ctrl)
   309  	shard.EXPECT().ReadEncoded(ctx, id, start, end, gomock.Any()).Return(nil, nil)
   310  	ns.shards[testShardIDs[0].ID()] = shard
   311  
   312  	shard.EXPECT().IsBootstrapped().Return(true)
   313  	_, err := ns.ReadEncoded(ctx, id, start, end)
   314  	require.NoError(t, err)
   315  
   316  	shard.EXPECT().IsBootstrapped().Return(false)
   317  	_, err = ns.ReadEncoded(ctx, id, start, end)
   318  	require.Error(t, err)
   319  	require.True(t, xerrors.IsRetryableError(err))
   320  	require.Equal(t, errShardNotBootstrappedToRead, xerrors.GetInnerRetryableError(err))
   321  }
   322  
   323  func TestNamespaceFetchBlocksShardNotOwned(t *testing.T) {
   324  	ctx := context.NewBackground()
   325  	defer ctx.Close()
   326  
   327  	ns, closer := newTestNamespace(t)
   328  	defer closer()
   329  
   330  	for i := range ns.shards {
   331  		ns.shards[i] = nil
   332  	}
   333  	_, err := ns.FetchBlocks(ctx, testShardIDs[0].ID(), ident.StringID("foo"), nil)
   334  	require.True(t, xerrors.IsRetryableError(err))
   335  	require.Equal(t, "not responsible for shard 0", err.Error())
   336  }
   337  
   338  func TestNamespaceFetchBlocksShardOwned(t *testing.T) {
   339  	ctrl := xtest.NewController(t)
   340  	defer ctrl.Finish()
   341  
   342  	ctx := context.NewBackground()
   343  	defer ctx.Close()
   344  
   345  	ns, closer := newTestNamespace(t)
   346  	defer closer()
   347  	shard := NewMockdatabaseShard(ctrl)
   348  	shard.EXPECT().FetchBlocks(ctx, ident.NewIDMatcher("foo"), nil, gomock.Any()).Return(nil, nil)
   349  	ns.shards[testShardIDs[0].ID()] = shard
   350  
   351  	shard.EXPECT().IsBootstrapped().Return(true)
   352  	res, err := ns.FetchBlocks(ctx, testShardIDs[0].ID(), ident.StringID("foo"), nil)
   353  	require.NoError(t, err)
   354  	require.Nil(t, res)
   355  
   356  	shard.EXPECT().IsBootstrapped().Return(false)
   357  	_, err = ns.FetchBlocks(ctx, testShardIDs[0].ID(), ident.StringID("foo"), nil)
   358  	require.Error(t, err)
   359  	require.True(t, xerrors.IsRetryableError(err))
   360  	require.Equal(t, errShardNotBootstrappedToRead, xerrors.GetInnerRetryableError(err))
   361  }
   362  
   363  func TestNamespaceBootstrapBootstrapping(t *testing.T) {
   364  	ns, closer := newTestNamespace(t)
   365  	defer closer()
   366  
   367  	ns.bootstrapState = Bootstrapping
   368  
   369  	ctx := context.NewBackground()
   370  	defer ctx.Close()
   371  
   372  	err := ns.Bootstrap(ctx, bootstrap.NamespaceResult{})
   373  	require.Equal(t, errNamespaceIsBootstrapping, err)
   374  }
   375  
   376  func TestNamespaceBootstrapDontNeedBootstrap(t *testing.T) {
   377  	ns, closer := newTestNamespaceWithIDOpts(t, defaultTestNs1ID,
   378  		namespace.NewOptions().SetBootstrapEnabled(false))
   379  	defer closer()
   380  
   381  	ctx := context.NewBackground()
   382  	defer ctx.Close()
   383  
   384  	require.NoError(t, ns.Bootstrap(ctx, bootstrap.NamespaceResult{}))
   385  	require.Equal(t, Bootstrapped, ns.bootstrapState)
   386  }
   387  
   388  func TestNamespaceBootstrapAllShards(t *testing.T) {
   389  	ctrl := xtest.NewController(t)
   390  	defer ctrl.Finish()
   391  
   392  	ns, closer := newTestNamespace(t)
   393  	defer closer()
   394  
   395  	errs := []error{nil, errors.New("foo")}
   396  	shardIDs := make([]uint32, 0, len(errs))
   397  	for i := range errs {
   398  		shardID := uint32(i)
   399  		shard := NewMockdatabaseShard(ctrl)
   400  		shard.EXPECT().IsBootstrapped().Return(false)
   401  		shard.EXPECT().ID().Return(shardID)
   402  		shard.EXPECT().Bootstrap(gomock.Any(), gomock.Any()).Return(errs[i])
   403  		ns.shards[testShardIDs[i].ID()] = shard
   404  		shardIDs = append(shardIDs, shardID)
   405  	}
   406  
   407  	nsResult := bootstrap.NamespaceResult{
   408  		DataResult: result.NewDataBootstrapResult(),
   409  		Shards:     shardIDs,
   410  	}
   411  
   412  	ctx := context.NewBackground()
   413  	defer ctx.Close()
   414  
   415  	require.Equal(t, "foo", ns.Bootstrap(ctx, nsResult).Error())
   416  	require.Equal(t, BootstrapNotStarted, ns.bootstrapState)
   417  }
   418  
   419  func TestNamespaceBootstrapOnlyNonBootstrappedShards(t *testing.T) {
   420  	ctrl := xtest.NewController(t)
   421  	defer ctrl.Finish()
   422  
   423  	var (
   424  		needsBootstrap, alreadyBootstrapped []shard.Shard
   425  		needsBootstrapShardIDs              []uint32
   426  	)
   427  	for i, shard := range testShardIDs {
   428  		if i%2 == 0 {
   429  			needsBootstrap = append(needsBootstrap, shard)
   430  			needsBootstrapShardIDs = append(needsBootstrapShardIDs, shard.ID())
   431  		} else {
   432  			alreadyBootstrapped = append(alreadyBootstrapped, shard)
   433  		}
   434  	}
   435  
   436  	require.True(t, len(needsBootstrap) > 0)
   437  	require.True(t, len(alreadyBootstrapped) > 0)
   438  
   439  	ns, closer := newTestNamespace(t)
   440  	defer closer()
   441  
   442  	shardIDs := make([]uint32, 0, len(needsBootstrap))
   443  	for _, testShard := range needsBootstrap {
   444  		shard := NewMockdatabaseShard(ctrl)
   445  		shard.EXPECT().IsBootstrapped().Return(false)
   446  		shard.EXPECT().ID().Return(testShard.ID())
   447  		shard.EXPECT().Bootstrap(gomock.Any(), gomock.Any()).Return(nil)
   448  		ns.shards[testShard.ID()] = shard
   449  		shardIDs = append(shardIDs, testShard.ID())
   450  	}
   451  
   452  	for _, testShard := range alreadyBootstrapped {
   453  		shard := NewMockdatabaseShard(ctrl)
   454  		shard.EXPECT().IsBootstrapped().Return(true)
   455  		shard.EXPECT().ID().Return(testShard.ID())
   456  		ns.shards[testShard.ID()] = shard
   457  		shardIDs = append(shardIDs, testShard.ID())
   458  	}
   459  
   460  	nsResult := bootstrap.NamespaceResult{
   461  		DataResult: result.NewDataBootstrapResult(),
   462  		Shards:     shardIDs,
   463  	}
   464  
   465  	ctx := context.NewBackground()
   466  	defer ctx.Close()
   467  
   468  	// do not panic for invariant violation to test some shards are still bootstrapped.
   469  	defer instrument.SetShouldPanicEnvironmentVariable(false)()
   470  	require.Error(t, ns.Bootstrap(ctx, nsResult))
   471  	require.Equal(t, BootstrapNotStarted, ns.bootstrapState)
   472  }
   473  
   474  func TestNamespaceBootstrapUnfulfilledShards(t *testing.T) {
   475  	unfulfilledRangeForShards := func(ids ...uint32) result.IndexBootstrapResult {
   476  		var (
   477  			unfulfilledRange = result.NewIndexBootstrapResult()
   478  			unfulfilledTo    = xtime.Now().Truncate(time.Hour)
   479  			unfulfilledFrom  = unfulfilledTo.Add(-time.Hour)
   480  		)
   481  		unfulfilledRange.SetUnfulfilled(result.NewShardTimeRangesFromRange(
   482  			unfulfilledFrom, unfulfilledTo, ids...))
   483  		return unfulfilledRange
   484  	}
   485  
   486  	shardIDs := []uint32{0, 1}
   487  	tests := []struct {
   488  		name                string
   489  		withIndex           bool
   490  		unfulfilledShardIDs []uint32
   491  		nsResult            bootstrap.NamespaceResult
   492  	}{
   493  		{
   494  			name:                "no index, unfulfilled data",
   495  			withIndex:           false,
   496  			unfulfilledShardIDs: []uint32{0},
   497  			nsResult: bootstrap.NamespaceResult{
   498  				DataResult: unfulfilledRangeForShards(0),
   499  				Shards:     shardIDs,
   500  			},
   501  		},
   502  		{
   503  			name:                "with index, unfulfilled data",
   504  			withIndex:           true,
   505  			unfulfilledShardIDs: []uint32{0, 1},
   506  			nsResult: bootstrap.NamespaceResult{
   507  				DataResult:  unfulfilledRangeForShards(0, 1),
   508  				IndexResult: unfulfilledRangeForShards(),
   509  				Shards:      shardIDs,
   510  			},
   511  		},
   512  		{
   513  			name:                "with index, unfulfilled index",
   514  			withIndex:           true,
   515  			unfulfilledShardIDs: []uint32{1},
   516  			nsResult: bootstrap.NamespaceResult{
   517  				DataResult:  unfulfilledRangeForShards(),
   518  				IndexResult: unfulfilledRangeForShards(1),
   519  				Shards:      shardIDs,
   520  			},
   521  		},
   522  		{
   523  			name:                "with index, unfulfilled data and index",
   524  			withIndex:           true,
   525  			unfulfilledShardIDs: []uint32{0, 1},
   526  			nsResult: bootstrap.NamespaceResult{
   527  				DataResult:  unfulfilledRangeForShards(0),
   528  				IndexResult: unfulfilledRangeForShards(1),
   529  				Shards:      shardIDs,
   530  			},
   531  		},
   532  	}
   533  
   534  	for _, tt := range tests {
   535  		t.Run(tt.name, func(t *testing.T) {
   536  			testNamespaceBootstrapUnfulfilledShards(t, shardIDs, tt.unfulfilledShardIDs,
   537  				tt.withIndex, tt.nsResult)
   538  		})
   539  	}
   540  }
   541  
   542  func testNamespaceBootstrapUnfulfilledShards(
   543  	t *testing.T,
   544  	shardIDs, unfulfilledShardIDs []uint32,
   545  	withIndex bool,
   546  	nsResult bootstrap.NamespaceResult,
   547  ) {
   548  	ctrl := xtest.NewController(t)
   549  	defer ctrl.Finish()
   550  	ctx := context.NewBackground()
   551  	defer ctx.Close()
   552  
   553  	var (
   554  		ns     *dbNamespace
   555  		closer func()
   556  	)
   557  	if withIndex {
   558  		idx := NewMockNamespaceIndex(ctrl)
   559  		idx.EXPECT().Bootstrap(gomock.Any()).Return(nil)
   560  		ns, closer = newTestNamespaceWithIndex(t, idx)
   561  	} else {
   562  		ns, closer = newTestNamespace(t)
   563  	}
   564  	defer closer()
   565  
   566  	for _, id := range shardIDs {
   567  		shard := NewMockdatabaseShard(ctrl)
   568  		shard.EXPECT().ID().Return(id)
   569  		shard.EXPECT().IsBootstrapped().Return(false)
   570  		if !contains(unfulfilledShardIDs, id) {
   571  			shard.EXPECT().Bootstrap(gomock.Any(), gomock.Any()).Return(nil)
   572  		}
   573  		ns.shards[id] = shard
   574  	}
   575  
   576  	require.Error(t, ns.Bootstrap(ctx, nsResult))
   577  }
   578  
   579  func TestNamespaceFlushNotBootstrapped(t *testing.T) {
   580  	ns, closer := newTestNamespace(t)
   581  	defer closer()
   582  	err := ns.WarmFlush(xtime.Now(), nil)
   583  	require.Equal(t, errNamespaceNotBootstrapped, err)
   584  	require.Equal(t, errNamespaceNotBootstrapped, ns.ColdFlush(nil))
   585  }
   586  
   587  func TestNamespaceFlushDontNeedFlush(t *testing.T) {
   588  	ns, close := newTestNamespaceWithIDOpts(t, defaultTestNs1ID,
   589  		namespace.NewOptions().SetFlushEnabled(false))
   590  	defer close()
   591  
   592  	ns.bootstrapState = Bootstrapped
   593  	err := ns.WarmFlush(xtime.Now(), nil)
   594  	require.NoError(t, err)
   595  	require.NoError(t, ns.ColdFlush(nil))
   596  }
   597  
   598  func TestNamespaceSkipFlushIfReadOnly(t *testing.T) {
   599  	ctrl := xtest.NewController(t)
   600  	defer ctrl.Finish()
   601  
   602  	// Configure the namespace so that flushing would only be enabled/disabled by read-only property.
   603  	indexOpts := namespace.NewIndexOptions().
   604  		SetEnabled(true)
   605  	nsOpts := namespace.NewOptions().
   606  		SetFlushEnabled(true).
   607  		SetColdWritesEnabled(true).
   608  		SetIndexOptions(indexOpts)
   609  
   610  	// Set mocked 'OnColdFlush' so that the test would fail if cold flush would happen.
   611  	opts := DefaultTestOptions().
   612  		SetOnColdFlush(NewMockOnColdFlush(ctrl)).
   613  		SetRuntimeOptionsManager(runtime.NewOptionsManager())
   614  
   615  	ns, closer := newTestNamespaceWithOpts(t, nsOpts, opts)
   616  	defer closer()
   617  
   618  	ns.bootstrapState = Bootstrapped
   619  	ns.SetReadOnly(true)
   620  	require.NoError(t, ns.WarmFlush(xtime.Now(), nil))
   621  	require.NoError(t, ns.ColdFlush(nil))
   622  	require.NoError(t, ns.FlushIndex(nil))
   623  }
   624  
   625  func TestNamespaceFlushSkipFlushed(t *testing.T) {
   626  	ctrl := xtest.NewController(t)
   627  	defer ctrl.Finish()
   628  
   629  	ctx := context.NewBackground()
   630  	defer ctx.Close()
   631  
   632  	ns, closer := newTestNamespace(t)
   633  	defer closer()
   634  
   635  	ns.bootstrapState = Bootstrapped
   636  	blockStart := xtime.Now().Truncate(ns.Options().RetentionOptions().BlockSize())
   637  
   638  	states := []fileOpState{
   639  		{WarmStatus: warmStatus{DataFlushed: fileOpNotStarted}},
   640  		{WarmStatus: warmStatus{DataFlushed: fileOpSuccess}},
   641  	}
   642  	for i, s := range states {
   643  		shard := NewMockdatabaseShard(ctrl)
   644  		shard.EXPECT().IsBootstrapped().Return(true).AnyTimes()
   645  		shard.EXPECT().FlushState(blockStart).Return(s, nil)
   646  		if s.WarmStatus.DataFlushed != fileOpSuccess {
   647  			shard.EXPECT().WarmFlush(blockStart, gomock.Any(), gomock.Any()).Return(nil)
   648  		}
   649  		ns.shards[testShardIDs[i].ID()] = shard
   650  	}
   651  
   652  	err := ns.WarmFlush(blockStart, nil)
   653  	require.NoError(t, err)
   654  }
   655  
   656  func TestNamespaceFlushSkipShardNotBootstrapped(t *testing.T) {
   657  	ctrl := xtest.NewController(t)
   658  	defer ctrl.Finish()
   659  
   660  	ctx := context.NewBackground()
   661  	defer ctx.Close()
   662  
   663  	ns, closer := newTestNamespace(t)
   664  	defer closer()
   665  
   666  	ns.bootstrapState = Bootstrapped
   667  	blockStart := xtime.Now().Truncate(ns.Options().RetentionOptions().BlockSize())
   668  
   669  	shard := NewMockdatabaseShard(ctrl)
   670  	shard.EXPECT().ID().Return(testShardIDs[0].ID()).AnyTimes()
   671  	shard.EXPECT().IsBootstrapped().Return(false)
   672  	ns.shards[testShardIDs[0].ID()] = shard
   673  
   674  	err := ns.WarmFlush(blockStart, nil)
   675  	require.NoError(t, err)
   676  	require.NoError(t, ns.ColdFlush(nil))
   677  }
   678  
   679  type snapshotTestCase struct {
   680  	isSnapshotting                bool
   681  	expectSnapshot                bool
   682  	shardBootstrapStateBeforeTick BootstrapState
   683  	lastSnapshotTime              func(blockStart time.Time, blockSize time.Duration) time.Time
   684  	shardSnapshotErr              error
   685  	isBootstrapped                bool
   686  }
   687  
   688  func TestNamespaceSnapshotNotBootstrapped(t *testing.T) {
   689  	ctrl := xtest.NewController(t)
   690  	defer ctrl.Finish()
   691  
   692  	ctx := context.NewBackground()
   693  	defer ctx.Close()
   694  
   695  	ns, close := newTestNamespace(t)
   696  	defer close()
   697  
   698  	ns.bootstrapState = Bootstrapping
   699  
   700  	blockSize := ns.Options().RetentionOptions().BlockSize()
   701  	blockStart := xtime.Now().Truncate(blockSize)
   702  	require.Equal(t, errNamespaceNotBootstrapped, ns.Snapshot([]xtime.UnixNano{blockStart}, blockStart, nil))
   703  }
   704  
   705  func TestNamespaceSnapshotAllShardsSuccess(t *testing.T) {
   706  	shardMethodResults := []snapshotTestCase{
   707  		{
   708  			isSnapshotting:                false,
   709  			expectSnapshot:                true,
   710  			shardBootstrapStateBeforeTick: Bootstrapped,
   711  			shardSnapshotErr:              nil,
   712  			isBootstrapped:                true,
   713  		},
   714  		{
   715  			isSnapshotting:                false,
   716  			expectSnapshot:                true,
   717  			shardBootstrapStateBeforeTick: Bootstrapped,
   718  			shardSnapshotErr:              nil,
   719  			isBootstrapped:                true,
   720  		},
   721  	}
   722  	require.NoError(t, testSnapshotWithShardSnapshotErrs(t, shardMethodResults))
   723  }
   724  
   725  func TestNamespaceSnapshotShardError(t *testing.T) {
   726  	shardMethodResults := []snapshotTestCase{
   727  		{
   728  			isSnapshotting:                false,
   729  			expectSnapshot:                true,
   730  			shardBootstrapStateBeforeTick: Bootstrapped,
   731  			shardSnapshotErr:              nil,
   732  			isBootstrapped:                true,
   733  		},
   734  		{
   735  			isSnapshotting:                false,
   736  			expectSnapshot:                true,
   737  			shardBootstrapStateBeforeTick: Bootstrapped,
   738  			shardSnapshotErr:              errors.New("err"),
   739  			isBootstrapped:                true,
   740  		},
   741  	}
   742  	require.Error(t, testSnapshotWithShardSnapshotErrs(t, shardMethodResults))
   743  }
   744  
   745  func TestNamespaceSnapshotShardSkipNotBootstrapped(t *testing.T) {
   746  	shardMethodResults := []snapshotTestCase{
   747  		{
   748  			isSnapshotting:                false,
   749  			expectSnapshot:                true,
   750  			shardBootstrapStateBeforeTick: Bootstrapped,
   751  			shardSnapshotErr:              nil,
   752  			isBootstrapped:                true,
   753  		},
   754  		{
   755  			isSnapshotting:                false,
   756  			expectSnapshot:                true,
   757  			shardBootstrapStateBeforeTick: Bootstrapped,
   758  			// Skip this shard (not bootstrapped) so we do not see this error.
   759  			shardSnapshotErr: errors.New("shard not bootstrapped"),
   760  			isBootstrapped:   false,
   761  		},
   762  	}
   763  	require.NoError(t, testSnapshotWithShardSnapshotErrs(t, shardMethodResults))
   764  }
   765  
   766  func TestNamespaceSnapshotShardBlockFiltered(t *testing.T) {
   767  	var (
   768  		ctrl       = xtest.NewController(t)
   769  		ctx        = context.NewBackground()
   770  		now        = xtime.Now()
   771  		ns, closer = newTestNamespaceWithIDOpts(t, defaultTestNs1ID,
   772  			namespace.NewOptions().SetSnapshotEnabled(true))
   773  	)
   774  	defer func() {
   775  		ctrl.Finish()
   776  		ctx.Close()
   777  		closer()
   778  	}()
   779  
   780  	var (
   781  		shardBootstrapStates = ShardBootstrapStates{}
   782  		blockSize            = ns.Options().RetentionOptions().BlockSize()
   783  		block1               = now.Truncate(blockSize)
   784  		block2               = block1.Truncate(blockSize)
   785  		blocks               = []xtime.UnixNano{block2, block1}
   786  		filteredBlocks       = []xtime.UnixNano{block1}
   787  	)
   788  
   789  	ns.bootstrapState = Bootstrapped
   790  	ns.nowFn = func() time.Time { return now.ToTime() }
   791  	shard := newTestShard(ctrl)
   792  	ns.shards[shard.ID()] = shard
   793  	shardBootstrapStates[shard.ID()] = Bootstrapped
   794  	shard.EXPECT().FilterBlocksNeedSnapshot(blocks).Return(filteredBlocks)
   795  	shard.EXPECT().
   796  		Snapshot(block1, now, gomock.Any(), gomock.Any()).
   797  		Return(ShardSnapshotResult{}, nil)
   798  
   799  	err := ns.Snapshot(blocks, now, nil)
   800  	require.NoError(t, err)
   801  }
   802  
   803  func TestNamespaceSnapshotShardBlockAllShardsFiltered(t *testing.T) {
   804  	var (
   805  		ctrl       = xtest.NewController(t)
   806  		ctx        = context.NewBackground()
   807  		now        = xtime.Now()
   808  		ns, closer = newTestNamespaceWithIDOpts(t, defaultTestNs1ID,
   809  			namespace.NewOptions().SetSnapshotEnabled(true))
   810  	)
   811  	defer func() {
   812  		ctrl.Finish()
   813  		ctx.Close()
   814  		closer()
   815  	}()
   816  
   817  	var (
   818  		shardBootstrapStates = ShardBootstrapStates{}
   819  		blockSize            = ns.Options().RetentionOptions().BlockSize()
   820  		blocks               = []xtime.UnixNano{now.Truncate(blockSize)}
   821  	)
   822  
   823  	ns.bootstrapState = Bootstrapped
   824  	ns.nowFn = func() time.Time { return now.ToTime() }
   825  	shard := newTestShard(ctrl)
   826  	ns.shards[shard.ID()] = shard
   827  	shardBootstrapStates[shard.ID()] = Bootstrapped
   828  	shard.EXPECT().FilterBlocksNeedSnapshot(blocks).Return([]xtime.UnixNano{})
   829  
   830  	err := ns.Snapshot(blocks, now, nil)
   831  	require.NoError(t, err)
   832  }
   833  
   834  func testSnapshotWithShardSnapshotErrs(
   835  	t *testing.T,
   836  	shardMethodResults []snapshotTestCase,
   837  ) error {
   838  	ctrl := xtest.NewController(t)
   839  	defer ctrl.Finish()
   840  
   841  	ctx := context.NewBackground()
   842  	defer ctx.Close()
   843  
   844  	ns, closer := newTestNamespaceWithIDOpts(t, defaultTestNs1ID,
   845  		namespace.NewOptions().SetSnapshotEnabled(true))
   846  	defer closer()
   847  	ns.bootstrapState = Bootstrapped
   848  	now := xtime.Now()
   849  	ns.nowFn = func() time.Time {
   850  		return now.ToTime()
   851  	}
   852  
   853  	var (
   854  		shardBootstrapStates = ShardBootstrapStates{}
   855  		blockSize            = ns.Options().RetentionOptions().BlockSize()
   856  		blockStart           = now.Truncate(blockSize)
   857  	)
   858  
   859  	for i, tc := range shardMethodResults {
   860  		shard := NewMockdatabaseShard(ctrl)
   861  		shardID := uint32(i)
   862  		shard.EXPECT().ID().Return(uint32(i)).AnyTimes()
   863  		shard.EXPECT().IsBootstrapped().Return(tc.isBootstrapped).AnyTimes()
   864  		if !tc.isBootstrapped {
   865  			continue
   866  		}
   867  		shard.EXPECT().
   868  			FilterBlocksNeedSnapshot([]xtime.UnixNano{blockStart}).
   869  			Return([]xtime.UnixNano{blockStart})
   870  		if tc.expectSnapshot {
   871  			shard.EXPECT().
   872  				Snapshot(blockStart, now, gomock.Any(), gomock.Any()).
   873  				Return(ShardSnapshotResult{}, tc.shardSnapshotErr)
   874  		}
   875  		ns.shards[testShardIDs[i].ID()] = shard
   876  		shardBootstrapStates[shardID] = tc.shardBootstrapStateBeforeTick
   877  	}
   878  
   879  	return ns.Snapshot([]xtime.UnixNano{blockStart}, now, nil)
   880  }
   881  
   882  func TestNamespaceTruncate(t *testing.T) {
   883  	ctrl := xtest.NewController(t)
   884  	defer ctrl.Finish()
   885  
   886  	ns, closer := newTestNamespace(t)
   887  	defer closer()
   888  	for _, shard := range testShardIDs {
   889  		mockShard := NewMockdatabaseShard(ctrl)
   890  		mockShard.EXPECT().NumSeries().Return(int64(shard.ID()))
   891  		mockShard.EXPECT().ID().Return(shard.ID())
   892  		ns.shards[shard.ID()] = mockShard
   893  	}
   894  
   895  	res, err := ns.Truncate()
   896  	require.NoError(t, err)
   897  	require.Equal(t, int64(1), res)
   898  	require.NotNil(t, ns.shards[testShardIDs[0].ID()])
   899  	require.True(t, ns.shards[testShardIDs[0].ID()].IsBootstrapped())
   900  }
   901  
   902  func TestNamespaceRepair(t *testing.T) {
   903  	ctrl := xtest.NewController(t)
   904  	defer ctrl.Finish()
   905  
   906  	ns, closer := newTestNamespaceWithIDOpts(t, defaultTestNs1ID,
   907  		namespace.NewOptions().SetRepairEnabled(true))
   908  	defer closer()
   909  	now := xtime.Now()
   910  	repairTimeRange := xtime.Range{Start: now, End: now.Add(time.Hour)}
   911  	opts := repair.NewOptions().SetRepairThrottle(time.Duration(0))
   912  	repairer := NewMockdatabaseShardRepairer(ctrl)
   913  	repairer.EXPECT().Options().Return(opts).AnyTimes()
   914  
   915  	errs := []error{nil, errors.New("foo")}
   916  	for i := range errs {
   917  		shard := NewMockdatabaseShard(ctrl)
   918  		var res repair.MetadataComparisonResult
   919  		if errs[i] == nil {
   920  			res = repair.MetadataComparisonResult{
   921  				NumSeries:           1,
   922  				NumBlocks:           2,
   923  				SizeDifferences:     repair.NewReplicaSeriesMetadata(),
   924  				ChecksumDifferences: repair.NewReplicaSeriesMetadata(),
   925  			}
   926  		}
   927  		shard.EXPECT().
   928  			Repair(gomock.Any(), gomock.Any(), gomock.Any(), repairTimeRange, repairer).
   929  			Return(res, errs[i])
   930  		ns.shards[testShardIDs[i].ID()] = shard
   931  	}
   932  
   933  	require.Equal(t, "foo", ns.Repair(repairer, repairTimeRange, NamespaceRepairOptions{}).Error())
   934  }
   935  
   936  func TestNamespaceShardAt(t *testing.T) {
   937  	ctrl := xtest.NewController(t)
   938  	defer ctrl.Finish()
   939  
   940  	ns, closer := newTestNamespace(t)
   941  	defer closer()
   942  
   943  	s0 := NewMockdatabaseShard(ctrl)
   944  	s0.EXPECT().IsBootstrapped().Return(true)
   945  	ns.shards[0] = s0
   946  
   947  	s1 := NewMockdatabaseShard(ctrl)
   948  	s1.EXPECT().IsBootstrapped().Return(false)
   949  	ns.shards[1] = s1
   950  
   951  	_, _, err := ns.ReadableShardAt(0)
   952  	require.NoError(t, err)
   953  	_, _, err = ns.ReadableShardAt(1)
   954  	require.Error(t, err)
   955  	require.True(t, xerrors.IsRetryableError(err))
   956  	require.Equal(t, errShardNotBootstrappedToRead.Error(), err.Error())
   957  	_, _, err = ns.ReadableShardAt(2)
   958  	require.Error(t, err)
   959  	require.True(t, xerrors.IsRetryableError(err))
   960  	require.Equal(t, "not responsible for shard 2", err.Error())
   961  }
   962  
   963  func TestNamespaceAssignShardSet(t *testing.T) {
   964  	ctrl := xtest.NewController(t)
   965  	defer ctrl.Finish()
   966  
   967  	shards := sharding.NewShards([]uint32{0, 1, 2, 3, 4}, shard.Available)
   968  	prevAssignment := shard.NewShards([]shard.Shard{shards[0], shards[2], shards[3]})
   969  	nextAssignment := shard.NewShards([]shard.Shard{shards[0], shards[4]})
   970  	closing := shard.NewShards([]shard.Shard{shards[2], shards[3]})
   971  	closingErrors := shard.NewShards([]shard.Shard{shards[3]})
   972  	adding := shard.NewShards([]shard.Shard{shards[4]})
   973  
   974  	metadata, err := namespace.NewMetadata(defaultTestNs1ID, namespace.NewOptions())
   975  	require.NoError(t, err)
   976  	hashFn := func(identifier ident.ID) uint32 { return shards[0].ID() }
   977  	shardSet, err := sharding.NewShardSet(prevAssignment.All(), hashFn)
   978  	require.NoError(t, err)
   979  	dopts := DefaultTestOptions()
   980  
   981  	reporter := xmetrics.NewTestStatsReporter(xmetrics.NewTestStatsReporterOptions())
   982  	scope, closer := tally.NewRootScope(tally.ScopeOptions{Reporter: reporter}, time.Millisecond)
   983  	defer closer.Close()
   984  
   985  	dopts = dopts.SetInstrumentOptions(dopts.InstrumentOptions().
   986  		SetMetricsScope(scope))
   987  	oNs, err := newDatabaseNamespace(metadata,
   988  		namespace.NewRuntimeOptionsManager(metadata.ID().String()),
   989  		shardSet, nil, nil, nil, dopts)
   990  	require.NoError(t, err)
   991  	ns := oNs.(*dbNamespace)
   992  
   993  	prevMockShards := make(map[uint32]*MockdatabaseShard)
   994  	for _, testShard := range prevAssignment.All() {
   995  		shard := NewMockdatabaseShard(ctrl)
   996  		shard.EXPECT().ID().Return(testShard.ID()).AnyTimes()
   997  		if closing.Contains(testShard.ID()) {
   998  			if closingErrors.Contains(testShard.ID()) {
   999  				shard.EXPECT().Close().Return(fmt.Errorf("an error"))
  1000  			} else {
  1001  				shard.EXPECT().Close().Return(nil)
  1002  			}
  1003  		}
  1004  		ns.shards[testShard.ID()] = shard
  1005  		prevMockShards[testShard.ID()] = shard
  1006  	}
  1007  
  1008  	nextShardSet, err := sharding.NewShardSet(nextAssignment.All(), hashFn)
  1009  	require.NoError(t, err)
  1010  
  1011  	ns.AssignShardSet(nextShardSet)
  1012  
  1013  	waitForStats(reporter, func(r xmetrics.TestStatsReporter) bool {
  1014  		var (
  1015  			counts       = r.Counters()
  1016  			adds         = int64(adding.NumShards())
  1017  			closeSuccess = int64(closing.NumShards() - closingErrors.NumShards())
  1018  			closeErrors  = int64(closingErrors.NumShards())
  1019  		)
  1020  		return counts["database.dbnamespace.shards.add"] == adds &&
  1021  			counts["database.dbnamespace.shards.close"] == closeSuccess &&
  1022  			counts["database.dbnamespace.shards.close-errors"] == closeErrors
  1023  	})
  1024  
  1025  	for _, shard := range shards {
  1026  		if nextAssignment.Contains(shard.ID()) {
  1027  			assert.NotNil(t, ns.shards[shard.ID()])
  1028  			if prevAssignment.Contains(shard.ID()) {
  1029  				assert.Equal(t, prevMockShards[shard.ID()], ns.shards[shard.ID()])
  1030  			} else {
  1031  				assert.True(t, adding.Contains(shard.ID()))
  1032  			}
  1033  		} else {
  1034  			assert.Nil(t, ns.shards[shard.ID()])
  1035  		}
  1036  	}
  1037  }
  1038  
  1039  type needsFlushTestCase struct {
  1040  	shardNum   uint32
  1041  	needsFlush map[xtime.UnixNano]bool
  1042  }
  1043  
  1044  func newNeedsFlushNamespace(t *testing.T, shardNumbers []uint32) *dbNamespace {
  1045  	shards := sharding.NewShards(shardNumbers, shard.Available)
  1046  	dopts := DefaultTestOptions()
  1047  
  1048  	hashFn := func(identifier ident.ID) uint32 { return shards[0].ID() }
  1049  	metadata, err := namespace.NewMetadata(defaultTestNs1ID, defaultTestNs1Opts)
  1050  	require.NoError(t, err)
  1051  	ropts := metadata.Options().RetentionOptions()
  1052  	shardSet, err := sharding.NewShardSet(shards, hashFn)
  1053  	require.NoError(t, err)
  1054  
  1055  	at := time.Unix(0, 0).Add(2 * ropts.RetentionPeriod())
  1056  	dopts = dopts.SetClockOptions(dopts.ClockOptions().SetNowFn(func() time.Time {
  1057  		return at
  1058  	}))
  1059  
  1060  	ns, err := newDatabaseNamespace(metadata,
  1061  		namespace.NewRuntimeOptionsManager(metadata.ID().String()),
  1062  		shardSet, nil, nil, nil, dopts)
  1063  	require.NoError(t, err)
  1064  	return ns.(*dbNamespace)
  1065  }
  1066  
  1067  func setShardExpects(ns *dbNamespace, ctrl *gomock.Controller, cases []needsFlushTestCase) {
  1068  	for _, cs := range cases {
  1069  		shard := NewMockdatabaseShard(ctrl)
  1070  		shard.EXPECT().ID().Return(cs.shardNum).AnyTimes()
  1071  		for t, needFlush := range cs.needsFlush {
  1072  			if needFlush {
  1073  				shard.EXPECT().FlushState(t).Return(fileOpState{
  1074  					WarmStatus: warmStatus{
  1075  						DataFlushed: fileOpNotStarted,
  1076  					},
  1077  				}, nil).AnyTimes()
  1078  			} else {
  1079  				shard.EXPECT().FlushState(t).Return(fileOpState{
  1080  					WarmStatus: warmStatus{
  1081  						DataFlushed: fileOpSuccess,
  1082  					},
  1083  				}, nil).AnyTimes()
  1084  			}
  1085  		}
  1086  		ns.shards[cs.shardNum] = shard
  1087  	}
  1088  }
  1089  
  1090  func TestNamespaceNeedsFlushRange(t *testing.T) {
  1091  	ctrl := xtest.NewController(t)
  1092  	defer ctrl.Finish()
  1093  
  1094  	var (
  1095  		shards    = []uint32{0, 2, 4}
  1096  		ns        = newNeedsFlushNamespace(t, shards)
  1097  		ropts     = ns.Options().RetentionOptions()
  1098  		blockSize = ropts.BlockSize()
  1099  		t1        = retention.FlushTimeEnd(ropts, xtime.ToUnixNano(ns.opts.ClockOptions().NowFn()()))
  1100  		t0        = t1.Add(-blockSize)
  1101  	)
  1102  
  1103  	inputCases := []needsFlushTestCase{
  1104  		{0, map[xtime.UnixNano]bool{t0: false, t1: true}},
  1105  		{2, map[xtime.UnixNano]bool{t0: false, t1: true}},
  1106  		{4, map[xtime.UnixNano]bool{t0: false, t1: true}},
  1107  	}
  1108  
  1109  	setShardExpects(ns, ctrl, inputCases)
  1110  
  1111  	assertNeedsFlush(t, ns, t0, t0, false)
  1112  	assertNeedsFlush(t, ns, t0, t1, true)
  1113  	assertNeedsFlush(t, ns, t1, t1, true)
  1114  	assertNeedsFlush(t, ns, t1, t0, false)
  1115  }
  1116  
  1117  func TestNamespaceNeedsFlushRangeMultipleShardConflict(t *testing.T) {
  1118  	ctrl := xtest.NewController(t)
  1119  	defer ctrl.Finish()
  1120  
  1121  	var (
  1122  		shards    = []uint32{0, 2, 4}
  1123  		ns        = newNeedsFlushNamespace(t, shards)
  1124  		ropts     = ns.Options().RetentionOptions()
  1125  		blockSize = ropts.BlockSize()
  1126  		t2        = retention.FlushTimeEnd(ropts, xtime.ToUnixNano(ns.opts.ClockOptions().NowFn()()))
  1127  		t1        = t2.Add(-blockSize)
  1128  		t0        = t1.Add(-blockSize)
  1129  	)
  1130  
  1131  	inputCases := []needsFlushTestCase{
  1132  		{0, map[xtime.UnixNano]bool{t0: false, t1: true, t2: true}},
  1133  		{2, map[xtime.UnixNano]bool{t0: true, t1: false, t2: true}},
  1134  		{4, map[xtime.UnixNano]bool{t0: false, t1: true, t2: true}},
  1135  	}
  1136  
  1137  	setShardExpects(ns, ctrl, inputCases)
  1138  	assertNeedsFlush(t, ns, t0, t0, true)
  1139  	assertNeedsFlush(t, ns, t1, t1, true)
  1140  	assertNeedsFlush(t, ns, t2, t2, true)
  1141  	assertNeedsFlush(t, ns, t0, t1, true)
  1142  	assertNeedsFlush(t, ns, t0, t2, true)
  1143  	assertNeedsFlush(t, ns, t1, t2, true)
  1144  	assertNeedsFlush(t, ns, t2, t1, false)
  1145  	assertNeedsFlush(t, ns, t2, t0, false)
  1146  }
  1147  
  1148  func TestNamespaceNeedsFlushRangeSingleShardConflict(t *testing.T) {
  1149  	ctrl := xtest.NewController(t)
  1150  	defer ctrl.Finish()
  1151  
  1152  	var (
  1153  		shards    = []uint32{0, 2, 4}
  1154  		ns        = newNeedsFlushNamespace(t, shards)
  1155  		ropts     = ns.Options().RetentionOptions()
  1156  		blockSize = ropts.BlockSize()
  1157  		t2        = retention.FlushTimeEnd(ropts, xtime.ToUnixNano(ns.opts.ClockOptions().NowFn()()))
  1158  		t1        = t2.Add(-blockSize)
  1159  		t0        = t1.Add(-blockSize)
  1160  	)
  1161  	inputCases := []needsFlushTestCase{
  1162  		{0, map[xtime.UnixNano]bool{t0: false, t1: false, t2: true}},
  1163  		{2, map[xtime.UnixNano]bool{t0: true, t1: false, t2: true}},
  1164  		{4, map[xtime.UnixNano]bool{t0: false, t1: false, t2: true}},
  1165  	}
  1166  
  1167  	setShardExpects(ns, ctrl, inputCases)
  1168  	assertNeedsFlush(t, ns, t0, t0, true)
  1169  	assertNeedsFlush(t, ns, t1, t1, false)
  1170  	assertNeedsFlush(t, ns, t2, t2, true)
  1171  	assertNeedsFlush(t, ns, t0, t1, true)
  1172  	assertNeedsFlush(t, ns, t0, t2, true)
  1173  	assertNeedsFlush(t, ns, t1, t2, true)
  1174  	assertNeedsFlush(t, ns, t2, t1, false)
  1175  	assertNeedsFlush(t, ns, t2, t0, false)
  1176  }
  1177  
  1178  func TestNamespaceNeedsFlushAllSuccess(t *testing.T) {
  1179  	ctrl := xtest.NewController(t)
  1180  	defer ctrl.Finish()
  1181  
  1182  	var (
  1183  		shards = sharding.NewShards([]uint32{0, 2, 4}, shard.Available)
  1184  		dopts  = DefaultTestOptions()
  1185  		hashFn = func(identifier ident.ID) uint32 { return shards[0].ID() }
  1186  	)
  1187  
  1188  	metadata, err := namespace.NewMetadata(defaultTestNs1ID, defaultTestNs1Opts)
  1189  	require.NoError(t, err)
  1190  	shardSet, err := sharding.NewShardSet(shards, hashFn)
  1191  	require.NoError(t, err)
  1192  
  1193  	ropts := metadata.Options().RetentionOptions()
  1194  	at := xtime.UnixNano(2 * ropts.RetentionPeriod())
  1195  	dopts = dopts.SetClockOptions(dopts.ClockOptions().SetNowFn(func() time.Time {
  1196  		return at.ToTime()
  1197  	}))
  1198  
  1199  	blockStart := retention.FlushTimeEnd(ropts, at)
  1200  
  1201  	oNs, err := newDatabaseNamespace(metadata,
  1202  		namespace.NewRuntimeOptionsManager(metadata.ID().String()),
  1203  		shardSet, nil, nil, nil, dopts)
  1204  	require.NoError(t, err)
  1205  	ns := oNs.(*dbNamespace)
  1206  
  1207  	for _, s := range shards {
  1208  		shard := NewMockdatabaseShard(ctrl)
  1209  		shard.EXPECT().ID().Return(s.ID()).AnyTimes()
  1210  		shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1211  			WarmStatus: warmStatus{
  1212  				DataFlushed: fileOpSuccess,
  1213  			},
  1214  		}, nil).AnyTimes()
  1215  		ns.shards[s.ID()] = shard
  1216  	}
  1217  
  1218  	assertNeedsFlush(t, ns, blockStart, blockStart, false)
  1219  }
  1220  
  1221  func TestNamespaceNeedsFlushAnyFailed(t *testing.T) {
  1222  	ctrl := xtest.NewController(t)
  1223  	defer ctrl.Finish()
  1224  
  1225  	var (
  1226  		shards = sharding.NewShards([]uint32{0, 2, 4}, shard.Available)
  1227  		dopts  = DefaultTestOptions()
  1228  	)
  1229  	testNs, err := namespace.NewMetadata(defaultTestNs1ID, defaultTestNs1Opts)
  1230  	require.NoError(t, err)
  1231  
  1232  	var (
  1233  		ropts  = testNs.Options().RetentionOptions()
  1234  		hashFn = func(identifier ident.ID) uint32 { return shards[0].ID() }
  1235  	)
  1236  	shardSet, err := sharding.NewShardSet(shards, hashFn)
  1237  	require.NoError(t, err)
  1238  
  1239  	at := xtime.UnixNano(2 * ropts.RetentionPeriod())
  1240  	dopts = dopts.SetClockOptions(dopts.ClockOptions().SetNowFn(func() time.Time {
  1241  		return at.ToTime()
  1242  	}))
  1243  
  1244  	blockStart := retention.FlushTimeEnd(ropts, at)
  1245  
  1246  	oNs, err := newDatabaseNamespace(testNs,
  1247  		namespace.NewRuntimeOptionsManager(testNs.ID().String()),
  1248  		shardSet, nil, nil, nil, dopts)
  1249  	require.NoError(t, err)
  1250  	ns := oNs.(*dbNamespace)
  1251  	for _, s := range shards {
  1252  		shard := NewMockdatabaseShard(ctrl)
  1253  		shard.EXPECT().ID().Return(s.ID()).AnyTimes()
  1254  		switch shard.ID() {
  1255  		case shards[0].ID():
  1256  			shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1257  				WarmStatus: warmStatus{
  1258  					DataFlushed: fileOpSuccess,
  1259  				},
  1260  			}, nil).AnyTimes()
  1261  		case shards[1].ID():
  1262  			shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1263  				WarmStatus: warmStatus{
  1264  					DataFlushed: fileOpSuccess,
  1265  				},
  1266  			}, nil).AnyTimes()
  1267  		case shards[2].ID():
  1268  			shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1269  				WarmStatus: warmStatus{
  1270  					DataFlushed: fileOpFailed,
  1271  				},
  1272  				NumFailures: 999,
  1273  			}, nil).AnyTimes()
  1274  		}
  1275  		ns.shards[s.ID()] = shard
  1276  	}
  1277  
  1278  	assertNeedsFlush(t, ns, blockStart, blockStart, true)
  1279  }
  1280  
  1281  func TestNamespaceNeedsFlushAnyNotStarted(t *testing.T) {
  1282  	ctrl := xtest.NewController(t)
  1283  	defer ctrl.Finish()
  1284  
  1285  	var (
  1286  		shards = sharding.NewShards([]uint32{0, 2, 4}, shard.Available)
  1287  		dopts  = DefaultTestOptions()
  1288  	)
  1289  	testNs, err := namespace.NewMetadata(defaultTestNs1ID, defaultTestNs1Opts)
  1290  	require.NoError(t, err)
  1291  
  1292  	var (
  1293  		ropts  = testNs.Options().RetentionOptions()
  1294  		hashFn = func(identifier ident.ID) uint32 { return shards[0].ID() }
  1295  	)
  1296  	shardSet, err := sharding.NewShardSet(shards, hashFn)
  1297  	require.NoError(t, err)
  1298  
  1299  	at := xtime.UnixNano(2 * ropts.RetentionPeriod())
  1300  	dopts = dopts.SetClockOptions(dopts.ClockOptions().SetNowFn(func() time.Time {
  1301  		return at.ToTime()
  1302  	}))
  1303  
  1304  	blockStart := retention.FlushTimeEnd(ropts, at)
  1305  
  1306  	oNs, err := newDatabaseNamespace(testNs,
  1307  		namespace.NewRuntimeOptionsManager(testNs.ID().String()),
  1308  		shardSet, nil, nil, nil, dopts)
  1309  	require.NoError(t, err)
  1310  	ns := oNs.(*dbNamespace)
  1311  	for _, s := range shards {
  1312  		shard := NewMockdatabaseShard(ctrl)
  1313  		shard.EXPECT().ID().Return(s.ID()).AnyTimes()
  1314  		switch shard.ID() {
  1315  		case shards[0].ID():
  1316  			shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1317  				WarmStatus: warmStatus{
  1318  					DataFlushed: fileOpSuccess,
  1319  				},
  1320  			}, nil).AnyTimes()
  1321  		case shards[1].ID():
  1322  			shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1323  				WarmStatus: warmStatus{
  1324  					DataFlushed: fileOpNotStarted,
  1325  				},
  1326  			}, nil).AnyTimes()
  1327  		case shards[2].ID():
  1328  			shard.EXPECT().FlushState(blockStart).Return(fileOpState{
  1329  				WarmStatus: warmStatus{
  1330  					DataFlushed: fileOpSuccess,
  1331  				},
  1332  			}, nil).AnyTimes()
  1333  		}
  1334  		ns.shards[s.ID()] = shard
  1335  	}
  1336  
  1337  	assertNeedsFlush(t, ns, blockStart, blockStart, true)
  1338  }
  1339  
  1340  func TestNamespaceCloseWillCloseShard(t *testing.T) {
  1341  	ctrl := xtest.NewController(t)
  1342  	defer ctrl.Finish()
  1343  
  1344  	ctx := context.NewBackground()
  1345  	defer ctx.Close()
  1346  
  1347  	// mock namespace + 1 shard
  1348  	ns, closer := newTestNamespace(t)
  1349  	defer closer()
  1350  
  1351  	// specify a mock shard to test being closed
  1352  	shard := NewMockdatabaseShard(ctrl)
  1353  	shard.EXPECT().Close().Return(nil)
  1354  	ns.Lock()
  1355  	ns.shards[testShardIDs[0].ID()] = shard
  1356  	ns.Unlock()
  1357  
  1358  	// Close the namespace
  1359  	require.NoError(t, ns.Close())
  1360  
  1361  	// Check the namespace no long owns any shards
  1362  	require.Empty(t, ns.OwnedShards())
  1363  }
  1364  
  1365  func TestNamespaceCloseDoesNotLeak(t *testing.T) {
  1366  	// Need to generate leaktest at top of test as that is when
  1367  	// goroutines that are interesting are captured
  1368  	leakCheck := leaktest.Check(t)
  1369  	defer leakCheck()
  1370  
  1371  	ctrl := xtest.NewController(t)
  1372  	defer ctrl.Finish()
  1373  
  1374  	ctx := context.NewBackground()
  1375  	defer ctx.Close()
  1376  
  1377  	// new namespace
  1378  	ns, closer := newTestNamespace(t)
  1379  	defer closer()
  1380  
  1381  	// verify has shards it will need to close
  1382  	ns.RLock()
  1383  	assert.True(t, len(ns.shards) > 0)
  1384  	ns.RUnlock()
  1385  
  1386  	// Close the namespace
  1387  	require.NoError(t, ns.Close())
  1388  
  1389  	// Check the namespace no long owns any shards
  1390  	require.Empty(t, ns.OwnedShards())
  1391  }
  1392  
  1393  func TestNamespaceIndexInsert(t *testing.T) {
  1394  	ctrl := xtest.NewController(t)
  1395  	defer ctrl.Finish()
  1396  
  1397  	truncateTypes := []series.TruncateType{series.TypeBlock, series.TypeNone}
  1398  	for _, truncateType := range truncateTypes {
  1399  		idx := NewMockNamespaceIndex(ctrl)
  1400  
  1401  		ns, closer := newTestNamespaceWithTruncateType(t, idx, truncateType)
  1402  		ns.reverseIndex = idx
  1403  		defer closer()
  1404  
  1405  		ctx := context.NewBackground()
  1406  		now := xtime.Now()
  1407  
  1408  		shard := NewMockdatabaseShard(ctrl)
  1409  
  1410  		opts := series.WriteOptions{
  1411  			TruncateType: truncateType,
  1412  		}
  1413  		shard.EXPECT().
  1414  			WriteTagged(ctx, ident.NewIDMatcher("a"), convert.EmptyTagMetadataResolver,
  1415  				now, 1.0, xtime.Second, nil, opts).
  1416  			Return(SeriesWrite{WasWritten: true}, nil)
  1417  		shard.EXPECT().
  1418  			WriteTagged(ctx, ident.NewIDMatcher("a"), convert.EmptyTagMetadataResolver,
  1419  				now, 1.0, xtime.Second, nil, opts).
  1420  			Return(SeriesWrite{WasWritten: false}, nil)
  1421  
  1422  		ns.shards[testShardIDs[0].ID()] = shard
  1423  
  1424  		seriesWrite, err := ns.WriteTagged(ctx, ident.StringID("a"),
  1425  			convert.EmptyTagMetadataResolver, now, 1.0, xtime.Second, nil)
  1426  		require.NoError(t, err)
  1427  		require.True(t, seriesWrite.WasWritten)
  1428  
  1429  		seriesWrite, err = ns.WriteTagged(ctx, ident.StringID("a"),
  1430  			convert.EmptyTagMetadataResolver, now, 1.0, xtime.Second, nil)
  1431  		require.NoError(t, err)
  1432  		require.False(t, seriesWrite.WasWritten)
  1433  
  1434  		shard.EXPECT().Close()
  1435  		idx.EXPECT().Close().Return(nil)
  1436  		require.NoError(t, ns.Close())
  1437  	}
  1438  }
  1439  
  1440  func TestNamespaceIndexQuery(t *testing.T) {
  1441  	ctrl := xtest.NewController(t)
  1442  	defer ctrl.Finish()
  1443  
  1444  	idx := NewMockNamespaceIndex(ctrl)
  1445  	idx.EXPECT().Bootstrapped().Return(true)
  1446  
  1447  	ns, closer := newTestNamespaceWithIndex(t, idx)
  1448  	defer closer()
  1449  
  1450  	ctx := context.NewBackground()
  1451  	mtr := mocktracer.New()
  1452  	sp := mtr.StartSpan("root")
  1453  	ctx.SetGoContext(opentracing.ContextWithSpan(stdlibctx.Background(), sp))
  1454  
  1455  	query := index.Query{
  1456  		Query: xidx.NewTermQuery([]byte("foo"), []byte("bar")),
  1457  	}
  1458  	opts := index.QueryOptions{}
  1459  
  1460  	idx.EXPECT().Query(gomock.Any(), query, opts)
  1461  	_, err := ns.QueryIDs(ctx, query, opts)
  1462  	require.NoError(t, err)
  1463  
  1464  	idx.EXPECT().Close().Return(nil)
  1465  	require.NoError(t, ns.Close())
  1466  
  1467  	sp.Finish()
  1468  	spans := mtr.FinishedSpans()
  1469  	require.Len(t, spans, 2)
  1470  	assert.Equal(t, tracepoint.NSQueryIDs, spans[0].OperationName)
  1471  	assert.Equal(t, "root", spans[1].OperationName)
  1472  }
  1473  
  1474  func TestNamespaceAggregateQuery(t *testing.T) {
  1475  	ctrl := xtest.NewController(t)
  1476  	defer ctrl.Finish()
  1477  
  1478  	idx := NewMockNamespaceIndex(ctrl)
  1479  	idx.EXPECT().Bootstrapped().Return(true)
  1480  
  1481  	ns, closer := newTestNamespaceWithIndex(t, idx)
  1482  	defer closer()
  1483  
  1484  	ctx := context.NewBackground()
  1485  	query := index.Query{
  1486  		Query: xidx.NewTermQuery([]byte("foo"), []byte("bar")),
  1487  	}
  1488  	aggOpts := index.AggregationOptions{}
  1489  
  1490  	idx.EXPECT().AggregateQuery(ctx, query, aggOpts)
  1491  	_, err := ns.AggregateQuery(ctx, query, aggOpts)
  1492  	require.NoError(t, err)
  1493  
  1494  	idx.EXPECT().Close().Return(nil)
  1495  	require.NoError(t, ns.Close())
  1496  }
  1497  
  1498  func TestNamespaceTicksIndex(t *testing.T) {
  1499  	ctrl := xtest.NewController(t)
  1500  	defer ctrl.Finish()
  1501  
  1502  	idx := NewMockNamespaceIndex(ctrl)
  1503  	ns, closer := newTestNamespaceWithIndex(t, idx)
  1504  	defer closer()
  1505  
  1506  	ns.RLock()
  1507  	nsCtx := ns.nsContextWithRLock()
  1508  	ns.RUnlock()
  1509  
  1510  	ctx := context.NewBackground()
  1511  	defer ctx.Close()
  1512  
  1513  	for _, s := range ns.shards {
  1514  		if s != nil {
  1515  			s.Bootstrap(ctx, nsCtx)
  1516  		}
  1517  	}
  1518  
  1519  	cancel := context.NewCancellable()
  1520  	idx.EXPECT().Tick(cancel, gomock.Any()).Return(namespaceIndexTickResult{}, nil)
  1521  	err := ns.Tick(cancel, xtime.Now())
  1522  	require.NoError(t, err)
  1523  }
  1524  
  1525  func TestNamespaceIndexDisabledQuery(t *testing.T) {
  1526  	ns, closer := newTestNamespace(t)
  1527  	defer closer()
  1528  
  1529  	ctx := context.NewBackground()
  1530  	query := index.Query{
  1531  		Query: xidx.NewTermQuery([]byte("foo"), []byte("bar")),
  1532  	}
  1533  	opts := index.QueryOptions{}
  1534  
  1535  	_, err := ns.QueryIDs(ctx, query, opts)
  1536  	require.Error(t, err)
  1537  
  1538  	require.NoError(t, ns.Close())
  1539  }
  1540  
  1541  func TestNamespaceBootstrapState(t *testing.T) {
  1542  	ns, closer := newTestNamespace(t)
  1543  	defer closer()
  1544  
  1545  	require.Equal(t, BootstrapNotStarted, ns.BootstrapState())
  1546  
  1547  	ns.bootstrapState = Bootstrapped
  1548  	require.Equal(t, Bootstrapped, ns.BootstrapState())
  1549  }
  1550  
  1551  func TestNamespaceShardBootstrapState(t *testing.T) {
  1552  	ctrl := xtest.NewController(t)
  1553  	defer ctrl.Finish()
  1554  
  1555  	ns, closer := newTestNamespace(t)
  1556  	defer closer()
  1557  
  1558  	shard0 := NewMockdatabaseShard(ctrl)
  1559  	shard0.EXPECT().ID().Return(uint32(0))
  1560  	shard0.EXPECT().BootstrapState().Return(Bootstrapped)
  1561  	ns.shards[0] = shard0
  1562  
  1563  	shard1 := NewMockdatabaseShard(ctrl)
  1564  	shard1.EXPECT().ID().Return(uint32(1))
  1565  	shard1.EXPECT().BootstrapState().Return(Bootstrapping)
  1566  	ns.shards[1] = shard1
  1567  
  1568  	require.Equal(t, ShardBootstrapStates{
  1569  		0: Bootstrapped,
  1570  		1: Bootstrapping,
  1571  	}, ns.ShardBootstrapState())
  1572  }
  1573  
  1574  func TestNamespaceFlushState(t *testing.T) {
  1575  	ctrl := xtest.NewController(t)
  1576  	defer ctrl.Finish()
  1577  
  1578  	ns, closer := newTestNamespace(t)
  1579  	defer closer()
  1580  
  1581  	var (
  1582  		blockStart         = xtime.Now().Truncate(2 * time.Hour)
  1583  		expectedFlushState = fileOpState{
  1584  			ColdVersionRetrievable: 2,
  1585  		}
  1586  		shard0 = NewMockdatabaseShard(ctrl)
  1587  	)
  1588  	shard0.EXPECT().FlushState(blockStart).Return(expectedFlushState, nil)
  1589  	ns.shards[0] = shard0
  1590  
  1591  	flushState, err := ns.FlushState(0, blockStart)
  1592  	require.NoError(t, err)
  1593  	require.Equal(t, expectedFlushState, flushState)
  1594  }
  1595  
  1596  func TestNamespaceAggregateTilesFailUntilBootstrapped(t *testing.T) {
  1597  	ctx := context.NewBackground()
  1598  	defer ctx.Close()
  1599  
  1600  	var (
  1601  		start = xtime.Now().Truncate(time.Hour)
  1602  		opts  = AggregateTilesOptions{Start: start, End: start.Add(time.Hour), InsOptions: insOpts}
  1603  	)
  1604  
  1605  	sourceNs, sourceCloser := newTestNamespaceWithIDOpts(t, sourceNsID, namespace.NewOptions())
  1606  	defer sourceCloser()
  1607  
  1608  	targetNs, targetCloser := newTestNamespaceWithIDOpts(t, targetNsID, namespace.NewOptions())
  1609  	defer targetCloser()
  1610  
  1611  	_, err := targetNs.AggregateTiles(ctx, sourceNs, opts)
  1612  	require.Equal(t, errNamespaceNotBootstrapped, err)
  1613  
  1614  	sourceNs.bootstrapState = Bootstrapped
  1615  
  1616  	_, err = targetNs.AggregateTiles(ctx, sourceNs, opts)
  1617  	require.Equal(t, errNamespaceNotBootstrapped, err)
  1618  }
  1619  
  1620  func TestNamespaceAggregateTiles(t *testing.T) {
  1621  	ctrl := xtest.NewController(t)
  1622  	defer ctrl.Finish()
  1623  
  1624  	ctx := context.NewBackground()
  1625  	defer ctx.Close()
  1626  
  1627  	var (
  1628  		start    = xtime.Now().Truncate(targetBlockSize)
  1629  		shard0ID = uint32(10)
  1630  		shard1ID = uint32(20)
  1631  
  1632  		createdWarmIndexForBlockStart xtime.UnixNano
  1633  	)
  1634  
  1635  	opts, err := NewAggregateTilesOptions(
  1636  		start, start.Add(targetBlockSize), time.Second, targetNsID, AggregateTilesRegular,
  1637  		false, false, nil, insOpts)
  1638  	require.NoError(t, err)
  1639  
  1640  	sourceNs, sourceCloser := newTestNamespaceWithIDOpts(t, sourceNsID, namespace.NewOptions())
  1641  	defer sourceCloser()
  1642  	sourceNs.bootstrapState = Bootstrapped
  1643  	sourceRetentionOpts := sourceNs.nopts.RetentionOptions().SetBlockSize(sourceBlockSize)
  1644  	sourceNs.nopts = sourceNs.nopts.SetRetentionOptions(sourceRetentionOpts)
  1645  
  1646  	targetNs, targetCloser := newTestNamespaceWithIDOpts(t, targetNsID, namespace.NewOptions())
  1647  	defer targetCloser()
  1648  	targetNs.bootstrapState = Bootstrapped
  1649  	targetNs.createEmptyWarmIndexIfNotExistsFn = func(blockStart xtime.UnixNano) error {
  1650  		createdWarmIndexForBlockStart = blockStart
  1651  		return nil
  1652  	}
  1653  	targetRetentionOpts := targetNs.nopts.RetentionOptions().SetBlockSize(targetBlockSize)
  1654  	targetNs.nopts = targetNs.nopts.SetColdWritesEnabled(true).SetRetentionOptions(targetRetentionOpts)
  1655  
  1656  	// Pass in mock cold flusher and expect the cold flush ns process to finish.
  1657  	mockOnColdFlushNs := NewMockOnColdFlushNamespace(ctrl)
  1658  	mockOnColdFlushNs.EXPECT().Done().Return(nil)
  1659  	mockOnColdFlush := NewMockOnColdFlush(ctrl)
  1660  	cfOpts := NewColdFlushNsOpts(false)
  1661  	mockOnColdFlush.EXPECT().ColdFlushNamespace(gomock.Any(), cfOpts).Return(mockOnColdFlushNs, nil)
  1662  	targetNs.opts = targetNs.opts.SetOnColdFlush(mockOnColdFlush)
  1663  
  1664  	targetShard0 := NewMockdatabaseShard(ctrl)
  1665  	targetShard1 := NewMockdatabaseShard(ctrl)
  1666  	targetNs.shards[0] = targetShard0
  1667  	targetNs.shards[1] = targetShard1
  1668  
  1669  	targetShard0.EXPECT().IsBootstrapped().Return(true)
  1670  	targetShard1.EXPECT().IsBootstrapped().Return(true)
  1671  
  1672  	targetShard0.EXPECT().ID().Return(shard0ID)
  1673  	targetShard1.EXPECT().ID().Return(shard1ID)
  1674  
  1675  	targetShard0.EXPECT().
  1676  		AggregateTiles(ctx, sourceNs, targetNs, shard0ID, mockOnColdFlushNs, opts).
  1677  		Return(int64(3), nil)
  1678  
  1679  	targetShard1.EXPECT().
  1680  		AggregateTiles(ctx, sourceNs, targetNs, shard1ID, mockOnColdFlushNs, opts).
  1681  		Return(int64(2), nil)
  1682  
  1683  	processedTileCount, err := targetNs.AggregateTiles(ctx, sourceNs, opts)
  1684  
  1685  	require.NoError(t, err)
  1686  	assert.Equal(t, int64(3+2), processedTileCount)
  1687  	assert.Equal(t, start, createdWarmIndexForBlockStart)
  1688  }
  1689  
  1690  func TestNamespaceAggregateTilesSkipBootstrappingShards(t *testing.T) {
  1691  	ctrl := xtest.NewController(t)
  1692  	defer ctrl.Finish()
  1693  
  1694  	ctx := context.NewBackground()
  1695  	defer ctx.Close()
  1696  
  1697  	start := xtime.Now().Truncate(targetBlockSize)
  1698  
  1699  	opts, err := NewAggregateTilesOptions(
  1700  		start, start.Add(targetBlockSize), time.Second, targetNsID, AggregateTilesRegular,
  1701  		false, false, nil, insOpts)
  1702  	require.NoError(t, err)
  1703  
  1704  	sourceNs, sourceCloser := newTestNamespaceWithIDOpts(t, sourceNsID, namespace.NewOptions())
  1705  	defer sourceCloser()
  1706  	sourceNs.bootstrapState = Bootstrapped
  1707  	sourceRetentionOpts := sourceNs.nopts.RetentionOptions().SetBlockSize(sourceBlockSize)
  1708  	sourceNs.nopts = sourceNs.nopts.SetRetentionOptions(sourceRetentionOpts)
  1709  
  1710  	targetNs, targetCloser := newTestNamespaceWithIDOpts(t, targetNsID, namespace.NewOptions())
  1711  	defer targetCloser()
  1712  	targetNs.bootstrapState = Bootstrapped
  1713  	targetNs.createEmptyWarmIndexIfNotExistsFn = func(blockStart xtime.UnixNano) error {
  1714  		return nil
  1715  	}
  1716  	targetRetentionOpts := targetNs.nopts.RetentionOptions().SetBlockSize(targetBlockSize)
  1717  	targetNs.nopts = targetNs.nopts.SetColdWritesEnabled(true).SetRetentionOptions(targetRetentionOpts)
  1718  
  1719  	targetShard0 := NewMockdatabaseShard(ctrl)
  1720  	targetShard1 := NewMockdatabaseShard(ctrl)
  1721  	targetNs.shards[0] = targetShard0
  1722  	targetNs.shards[1] = targetShard1
  1723  
  1724  	targetShard0.EXPECT().IsBootstrapped().Return(false)
  1725  	targetShard1.EXPECT().IsBootstrapped().Return(false)
  1726  
  1727  	targetShard0.EXPECT().ID().Return(uint32(10))
  1728  	targetShard1.EXPECT().ID().Return(uint32(11))
  1729  
  1730  	processedTileCount, err := targetNs.AggregateTiles(ctx, sourceNs, opts)
  1731  
  1732  	require.NoError(t, err)
  1733  	assert.Zero(t, processedTileCount)
  1734  }
  1735  
  1736  func TestNamespaceAggregateTilesSkipIfNoShardsOwned(t *testing.T) {
  1737  	ctrl := xtest.NewController(t)
  1738  	defer ctrl.Finish()
  1739  
  1740  	ctx := context.NewBackground()
  1741  	defer ctx.Close()
  1742  
  1743  	start := xtime.Now().Truncate(targetBlockSize)
  1744  
  1745  	opts, err := NewAggregateTilesOptions(
  1746  		start, start.Add(targetBlockSize), time.Second, targetNsID, AggregateTilesRegular,
  1747  		false, false, nil, insOpts)
  1748  	require.NoError(t, err)
  1749  
  1750  	sourceNs, sourceCloser := newTestNamespaceWithIDOpts(t, sourceNsID, namespace.NewOptions())
  1751  	defer sourceCloser()
  1752  	sourceNs.bootstrapState = Bootstrapped
  1753  	sourceRetentionOpts := sourceNs.nopts.RetentionOptions().SetBlockSize(sourceBlockSize)
  1754  	sourceNs.nopts = sourceNs.nopts.SetRetentionOptions(sourceRetentionOpts)
  1755  
  1756  	targetNs, targetCloser := newTestNamespaceWithIDOpts(t, targetNsID, namespace.NewOptions())
  1757  	defer targetCloser()
  1758  	targetNs.bootstrapState = Bootstrapped
  1759  	targetNs.createEmptyWarmIndexIfNotExistsFn = func(blockStart xtime.UnixNano) error {
  1760  		return nil
  1761  	}
  1762  	targetRetentionOpts := targetNs.nopts.RetentionOptions().SetBlockSize(targetBlockSize)
  1763  	targetNs.nopts = targetNs.nopts.SetColdWritesEnabled(true).SetRetentionOptions(targetRetentionOpts)
  1764  
  1765  	noShards := sharding.NewEmptyShardSet(sharding.DefaultHashFn(1))
  1766  	targetNs.AssignShardSet(noShards)
  1767  
  1768  	processedTileCount, err := targetNs.AggregateTiles(ctx, sourceNs, opts)
  1769  
  1770  	require.NoError(t, err)
  1771  	assert.Zero(t, processedTileCount)
  1772  }
  1773  
  1774  func TestCreateEmptyWarmIndexIfNotExists(t *testing.T) {
  1775  	ctrl := xtest.NewController(t)
  1776  	defer ctrl.Finish()
  1777  
  1778  	var (
  1779  		nsID       = ident.StringID("warm")
  1780  		nsOpts     = namespace.NewOptions()
  1781  		blockStart = xtime.Now().Truncate(nsOpts.IndexOptions().BlockSize())
  1782  	)
  1783  
  1784  	ns, nsCloser := newTestNamespaceWithIDOpts(t, nsID, nsOpts)
  1785  	defer nsCloser()
  1786  
  1787  	shard0 := NewMockdatabaseShard(ctrl)
  1788  	shard1 := NewMockdatabaseShard(ctrl)
  1789  	ns.shards[0] = shard0
  1790  	ns.shards[1] = shard1
  1791  
  1792  	shard0.EXPECT().IsBootstrapped().Return(false).Times(2)
  1793  
  1794  	shard1.EXPECT().IsBootstrapped().Return(true).Times(2)
  1795  	shard1.EXPECT().ID().Return(uint32(5)).Times(2)
  1796  
  1797  	bootstrappedShardIds := map[uint32]struct{}{5: {}}
  1798  
  1799  	err := ns.createEmptyWarmIndexIfNotExists(blockStart)
  1800  	require.NoError(t, err)
  1801  
  1802  	err = ns.createEmptyWarmIndexIfNotExists(blockStart)
  1803  	require.NoError(t, err)
  1804  
  1805  	reader, err := fs.NewIndexReader(ns.StorageOptions().CommitLogOptions().FilesystemOptions())
  1806  	require.NoError(t, err)
  1807  
  1808  	openOpts := fs.IndexReaderOpenOptions{
  1809  		Identifier: fs.FileSetFileIdentifier{
  1810  			FileSetContentType: persist.FileSetIndexContentType,
  1811  			Namespace:          nsID,
  1812  			BlockStart:         blockStart,
  1813  			VolumeIndex:        0,
  1814  		},
  1815  	}
  1816  
  1817  	idx, err := reader.Open(openOpts)
  1818  	require.NoError(t, err)
  1819  
  1820  	assert.Equal(t, bootstrappedShardIds, idx.Shards)
  1821  	assert.Equal(t, idxpersist.DefaultIndexVolumeType, reader.IndexVolumeType())
  1822  }
  1823  
  1824  func waitForStats(
  1825  	reporter xmetrics.TestStatsReporter,
  1826  	check func(xmetrics.TestStatsReporter) bool,
  1827  ) {
  1828  	var wg sync.WaitGroup
  1829  	wg.Add(1)
  1830  	go func() {
  1831  		for !check(reporter) {
  1832  			time.Sleep(100 * time.Millisecond)
  1833  		}
  1834  		wg.Done()
  1835  	}()
  1836  
  1837  	wg.Wait()
  1838  }
  1839  
  1840  func assertNeedsFlush(t *testing.T, ns *dbNamespace, t0, t1 xtime.UnixNano, assertTrue bool) {
  1841  	needsFlush, err := ns.NeedsFlush(t0, t1)
  1842  	require.NoError(t, err)
  1843  	require.Equal(t, assertTrue, needsFlush)
  1844  }
  1845  
  1846  func contains(c []uint32, x uint32) bool {
  1847  	for _, y := range c {
  1848  		if x == y {
  1849  			return true
  1850  		}
  1851  	}
  1852  	return false
  1853  }
  1854  
  1855  func newTestShard(ctrl *gomock.Controller) *MockdatabaseShard {
  1856  	s := NewMockdatabaseShard(ctrl)
  1857  	shardID := testShardIDs[0].ID()
  1858  	s.EXPECT().ID().Return(shardID).AnyTimes()
  1859  	s.EXPECT().IsBootstrapped().Return(true)
  1860  	return s
  1861  }