github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/client/write_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 client
    22  
    23  import (
    24  	"errors"
    25  	"sync"
    26  	"testing"
    27  
    28  	"github.com/m3db/m3/src/cluster/shard"
    29  	tterrors "github.com/m3db/m3/src/dbnode/network/server/tchannelthrift/errors"
    30  	"github.com/m3db/m3/src/dbnode/sharding"
    31  	"github.com/m3db/m3/src/dbnode/topology"
    32  	"github.com/m3db/m3/src/x/checked"
    33  	xerrors "github.com/m3db/m3/src/x/errors"
    34  
    35  	"github.com/golang/mock/gomock"
    36  	"github.com/stretchr/testify/assert"
    37  	"github.com/stretchr/testify/require"
    38  )
    39  
    40  // shard state tests
    41  
    42  func testWriteSuccess(t *testing.T, state shard.State, success bool) {
    43  	var writeWg sync.WaitGroup
    44  
    45  	wState, s, hosts := writeTestSetup(t, &writeWg)
    46  	setShardStates(t, s, hosts[0], state)
    47  	wState.completionFn(hosts[0], nil)
    48  
    49  	if success {
    50  		assert.Equal(t, int32(1), wState.success)
    51  	} else {
    52  		assert.Equal(t, int32(0), wState.success)
    53  	}
    54  
    55  	writeTestTeardown(wState, &writeWg)
    56  }
    57  
    58  func TestWriteToAvailableShards(t *testing.T) {
    59  	testWriteSuccess(t, shard.Available, true)
    60  }
    61  
    62  func TestWriteToInitializingShards(t *testing.T) {
    63  	testWriteSuccess(t, shard.Initializing, false)
    64  }
    65  
    66  func TestWriteToLeavingShards(t *testing.T) {
    67  	testWriteSuccess(t, shard.Leaving, false)
    68  }
    69  
    70  // retryability test
    71  
    72  type errTestFn func(error) bool
    73  
    74  func retryabilityCheck(t *testing.T, wState *writeState, testFn errTestFn) {
    75  	require.True(t, len(wState.errors) == 1)
    76  	assert.True(t, testFn(wState.errors[0]))
    77  }
    78  
    79  func simpleRetryableTest(t *testing.T, passedErr error, customHost topology.Host, testFn errTestFn) {
    80  	var writeWg sync.WaitGroup
    81  
    82  	wState, _, hosts := writeTestSetup(t, &writeWg)
    83  	if customHost != nil {
    84  		hosts[0] = customHost
    85  	}
    86  	wState.completionFn(hosts[0], passedErr)
    87  	retryabilityCheck(t, wState, testFn)
    88  	writeTestTeardown(wState, &writeWg)
    89  }
    90  
    91  func TestNonRetryableError(t *testing.T) {
    92  	simpleRetryableTest(t, xerrors.NewNonRetryableError(errors.New("")), nil, xerrors.IsNonRetryableError)
    93  }
    94  
    95  func TestBadRequestError(t *testing.T) {
    96  	simpleRetryableTest(t, tterrors.NewBadRequestError(errors.New("")), nil, IsBadRequestError)
    97  }
    98  
    99  func TestRetryableError(t *testing.T) {
   100  	simpleRetryableTest(t, xerrors.NewRetryableError(errors.New("")), nil, xerrors.IsRetryableError)
   101  }
   102  
   103  func TestBadHostID(t *testing.T) {
   104  	simpleRetryableTest(t, nil, fakeHost{id: "not a real host"}, xerrors.IsRetryableError)
   105  }
   106  
   107  func TestBadShardID(t *testing.T) {
   108  	var writeWg sync.WaitGroup
   109  
   110  	wState, _, hosts := writeTestSetup(t, &writeWg)
   111  	o := wState.op.(*writeOperation)
   112  	o.shardID = writeOperationZeroed.shardID
   113  	wState.completionFn(hosts[0], nil)
   114  	retryabilityCheck(t, wState, xerrors.IsRetryableError)
   115  	writeTestTeardown(wState, &writeWg)
   116  }
   117  
   118  func TestShardNotAvailable(t *testing.T) {
   119  	var writeWg sync.WaitGroup
   120  
   121  	wState, s, hosts := writeTestSetup(t, &writeWg)
   122  	setShardStates(t, s, hosts[0], shard.Initializing)
   123  	wState.completionFn(hosts[0], nil)
   124  	retryabilityCheck(t, wState, xerrors.IsRetryableError)
   125  	writeTestTeardown(wState, &writeWg)
   126  }
   127  
   128  func TestShardLeavingWithShardsLeavingCountTowardsConsistency(t *testing.T) {
   129  	var writeWg sync.WaitGroup
   130  
   131  	wState, s, hosts := writeTestSetup(t, &writeWg)
   132  	wState.shardsLeavingCountTowardsConsistency = true
   133  	setShardStates(t, s, hosts[0], shard.Leaving)
   134  	wState.completionFn(hosts[0], nil)
   135  	assert.Equal(t, int32(1), wState.success)
   136  	writeTestTeardown(wState, &writeWg)
   137  }
   138  
   139  func TestShardLeavingAndInitializingCountTowardsConsistencyWithTrueFlag(t *testing.T) {
   140  	var writeWg sync.WaitGroup
   141  
   142  	wState, s, hosts := writeTestSetup(t, &writeWg)
   143  
   144  	setupShardLeavingAndInitializingCountTowardsConsistency(t, wState, s, true)
   145  	wState.completionFn(hosts[1], nil)
   146  	wState.incRef()
   147  	assert.Equal(t, int32(0), wState.success)
   148  	wState.completionFn(hosts[0], nil)
   149  	assert.Equal(t, int32(1), wState.success)
   150  	writeTestTeardown(wState, &writeWg)
   151  }
   152  
   153  func TestShardLeavingAndInitializingCountTowardsConsistencyWithFalseFlag(t *testing.T) {
   154  	var writeWg sync.WaitGroup
   155  
   156  	wState, s, hosts := writeTestSetup(t, &writeWg)
   157  
   158  	setupShardLeavingAndInitializingCountTowardsConsistency(t, wState, s, false)
   159  	wState.completionFn(hosts[1], nil)
   160  	wState.incRef()
   161  	wState.completionFn(hosts[0], nil)
   162  	assert.Equal(t, int32(0), wState.success)
   163  	writeTestTeardown(wState, &writeWg)
   164  }
   165  
   166  func setupShardLeavingAndInitializingCountTowardsConsistency(
   167  	t *testing.T,
   168  	wState *writeState,
   169  	s *session,
   170  	leavingAndInitializingFlag bool) {
   171  	hostShardSets := []topology.HostShardSet{}
   172  	for _, host := range s.state.topoMap.Hosts() {
   173  		hostShard, _ := sharding.NewShardSet(
   174  			sharding.NewShards([]uint32{0, 1, 2}, shard.Available),
   175  			sharding.DefaultHashFn(3),
   176  		)
   177  		hostShardSet := topology.NewHostShardSet(host, hostShard)
   178  		hostShardSets = append(hostShardSets, hostShardSet)
   179  	}
   180  	opts := topology.NewStaticOptions().
   181  		SetShardSet(s.state.topoMap.ShardSet()).
   182  		SetReplicas(3).
   183  		SetHostShardSets(hostShardSets)
   184  	m := topology.NewStaticMap(opts)
   185  	s.state.topoMap = m
   186  	wState.topoMap = m // update topology with hostshards options
   187  
   188  	// mark leaving shards in host0 and init in host1
   189  	markHostReplacement(t, s, s.state.topoMap.Hosts()[0], s.state.topoMap.Hosts()[1])
   190  
   191  	opts = topology.NewStaticOptions().
   192  		SetShardSet(s.state.topoMap.ShardSet()).
   193  		SetReplicas(3).
   194  		SetHostShardSets(hostShardSets)
   195  	m = topology.NewStaticMap(opts)
   196  	wState.topoMap = m
   197  	s.state.topoMap = m // update the topology manually after replace node.
   198  
   199  	wState.shardsLeavingAndInitializingCountTowardsConsistency = leavingAndInitializingFlag
   200  }
   201  
   202  // utils
   203  
   204  func getWriteState(s *session, w writeStub) *writeState {
   205  	wState := s.pools.writeState.Get()
   206  	s.state.RLock()
   207  	wState.consistencyLevel = s.state.writeLevel
   208  	wState.topoMap = s.state.topoMap
   209  	s.state.RUnlock()
   210  	o := s.pools.writeOperation.Get()
   211  	o.shardID = 0 // Any valid shardID
   212  	wState.op = o
   213  	wState.nsID = w.ns
   214  	wState.tsID = w.id
   215  	var clonedAnnotation checked.Bytes
   216  	if len(w.annotation) > 0 {
   217  		clonedAnnotation = s.pools.checkedBytes.Get(len(w.annotation))
   218  		clonedAnnotation.IncRef()
   219  		clonedAnnotation.AppendAll(w.annotation)
   220  	}
   221  	wState.annotation = clonedAnnotation
   222  	return wState
   223  }
   224  
   225  func setShardStates(t *testing.T, s *session, host topology.Host, state shard.State) {
   226  	s.state.RLock()
   227  	hostShardSet, ok := s.state.topoMap.LookupHostShardSet(host.ID())
   228  	s.state.RUnlock()
   229  	require.True(t, ok)
   230  
   231  	for _, hostShard := range hostShardSet.ShardSet().All() {
   232  		hostShard.SetState(state)
   233  	}
   234  }
   235  
   236  func markHostReplacement(t *testing.T, s *session, leavingHost topology.Host, initializingHost topology.Host) {
   237  	s.state.RLock()
   238  	leavingHostShardSet, ok := s.state.topoMap.LookupHostShardSet(leavingHost.ID())
   239  	require.True(t, ok)
   240  	initializingHostShardSet, ok := s.state.topoMap.LookupHostShardSet(initializingHost.ID())
   241  	s.state.RUnlock()
   242  	require.True(t, ok)
   243  
   244  	for _, leavinghostShard := range leavingHostShardSet.ShardSet().All() {
   245  		leavinghostShard.SetState(shard.Leaving)
   246  	}
   247  	for _, initializinghostShard := range initializingHostShardSet.ShardSet().All() {
   248  		initializinghostShard.SetState(shard.Initializing)
   249  		initializinghostShard.SetSourceID(leavingHost.ID())
   250  	}
   251  }
   252  
   253  type fakeHost struct{ id string }
   254  
   255  func (f fakeHost) ID() string      { return f.id }
   256  func (f fakeHost) Address() string { return "" }
   257  func (f fakeHost) String() string  { return "" }
   258  
   259  func writeTestSetup(t *testing.T, writeWg *sync.WaitGroup) (*writeState, *session, []topology.Host) {
   260  	ctrl := gomock.NewController(t)
   261  	defer ctrl.Finish()
   262  
   263  	s := newDefaultTestSession(t).(*session)
   264  	w := newWriteStub()
   265  
   266  	var completionFn completionFn
   267  	enqueueWg := mockHostQueues(ctrl, s, sessionTestReplicas, []testEnqueueFn{func(idx int, op op) {
   268  		completionFn = op.CompletionFn()
   269  	}})
   270  
   271  	require.NoError(t, s.Open())
   272  	defer func() {
   273  		require.NoError(t, s.Close())
   274  	}()
   275  
   276  	hosts := s.state.topoMap.Hosts()
   277  
   278  	wState := getWriteState(s, w)
   279  	wState.incRef() // for the test
   280  	wState.incRef() // allow introspection
   281  	// Begin write
   282  	writeWg.Add(1)
   283  	go func() {
   284  		s.Write(w.ns, w.id, w.t, w.value, w.unit, w.annotation)
   285  		writeWg.Done()
   286  	}()
   287  
   288  	// Callbacks
   289  
   290  	enqueueWg.Wait()
   291  	require.True(t, s.state.topoMap.Replicas() == sessionTestReplicas)
   292  	for i := 0; i < s.state.topoMap.Replicas(); i++ {
   293  		completionFn(hosts[i], nil) // maintain session state
   294  	}
   295  
   296  	return wState, s, hosts
   297  }
   298  
   299  func writeTestTeardown(wState *writeState, writeWg *sync.WaitGroup) {
   300  	wState.decRef() // end introspection
   301  	writeWg.Wait()  // wait for write to complete
   302  }