go.temporal.io/server@v1.23.0/common/persistence/tests/shard.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package tests
    26  
    27  import (
    28  	"context"
    29  	"math/rand"
    30  	"testing"
    31  	"time"
    32  
    33  	"github.com/stretchr/testify/require"
    34  	"github.com/stretchr/testify/suite"
    35  
    36  	"go.temporal.io/server/common/debug"
    37  	"go.temporal.io/server/common/log"
    38  	p "go.temporal.io/server/common/persistence"
    39  	"go.temporal.io/server/common/persistence/serialization"
    40  	"go.temporal.io/server/common/testing/protorequire"
    41  )
    42  
    43  type (
    44  	ShardSuite struct {
    45  		suite.Suite
    46  		*require.Assertions
    47  		protorequire.ProtoAssertions
    48  
    49  		ShardID int32
    50  
    51  		ShardManager p.ShardManager
    52  		Logger       log.Logger
    53  
    54  		Ctx    context.Context
    55  		Cancel context.CancelFunc
    56  	}
    57  )
    58  
    59  func NewShardSuite(
    60  	t *testing.T,
    61  	shardStore p.ShardStore,
    62  	serializer serialization.Serializer,
    63  	logger log.Logger,
    64  ) *ShardSuite {
    65  	return &ShardSuite{
    66  		Assertions:      require.New(t),
    67  		ProtoAssertions: protorequire.New(t),
    68  		ShardManager: p.NewShardManager(
    69  			shardStore,
    70  			serializer,
    71  		),
    72  		Logger: logger,
    73  	}
    74  }
    75  
    76  func (s *ShardSuite) SetupSuite() {
    77  }
    78  
    79  func (s *ShardSuite) TearDownSuite() {
    80  }
    81  
    82  func (s *ShardSuite) SetupTest() {
    83  	s.Assertions = require.New(s.T())
    84  	s.ProtoAssertions = protorequire.New(s.T())
    85  	s.Ctx, s.Cancel = context.WithTimeout(context.Background(), 30*time.Second*debug.TimeoutMultiplier)
    86  
    87  	s.ShardID++
    88  }
    89  
    90  func (s *ShardSuite) TearDownTest() {
    91  	s.Cancel()
    92  }
    93  
    94  func (s *ShardSuite) TestGetOrCreateShard_Create() {
    95  	rangeID := rand.Int63()
    96  	shardInfo := RandomShardInfo(s.ShardID, rangeID)
    97  
    98  	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
    99  		ShardID:          s.ShardID,
   100  		InitialShardInfo: shardInfo,
   101  	})
   102  	s.NoError(err)
   103  	s.ProtoEqual(shardInfo, resp.ShardInfo)
   104  
   105  }
   106  
   107  func (s *ShardSuite) TestGetOrCreateShard_Get() {
   108  	rangeID := rand.Int63()
   109  	shardInfo := RandomShardInfo(s.ShardID, rangeID)
   110  
   111  	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
   112  		ShardID:          s.ShardID,
   113  		InitialShardInfo: shardInfo,
   114  	})
   115  	s.NoError(err)
   116  	s.ProtoEqual(shardInfo, resp.ShardInfo)
   117  
   118  	resp, err = s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
   119  		ShardID:          s.ShardID,
   120  		InitialShardInfo: RandomShardInfo(s.ShardID, rand.Int63()),
   121  	})
   122  	s.NoError(err)
   123  	s.ProtoEqual(shardInfo, resp.ShardInfo)
   124  }
   125  
   126  func (s *ShardSuite) TestUpdateShard_OwnershipLost() {
   127  	rangeID := rand.Int63()
   128  	shardInfo := RandomShardInfo(s.ShardID, rangeID)
   129  
   130  	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
   131  		ShardID:          s.ShardID,
   132  		InitialShardInfo: shardInfo,
   133  	})
   134  	s.NoError(err)
   135  	s.ProtoEqual(shardInfo, resp.ShardInfo)
   136  
   137  	updateRangeID := rand.Int63()
   138  	updateShardInfo := RandomShardInfo(s.ShardID, rand.Int63())
   139  	err = s.ShardManager.UpdateShard(s.Ctx, &p.UpdateShardRequest{
   140  		ShardInfo:       updateShardInfo,
   141  		PreviousRangeID: updateRangeID,
   142  	})
   143  	s.IsType(&p.ShardOwnershipLostError{}, err)
   144  
   145  	resp, err = s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
   146  		ShardID:          s.ShardID,
   147  		InitialShardInfo: shardInfo,
   148  	})
   149  	s.NoError(err)
   150  	s.ProtoEqual(shardInfo, resp.ShardInfo)
   151  }
   152  
   153  func (s *ShardSuite) TestUpdateShard_Success() {
   154  	rangeID := rand.Int63()
   155  	shardInfo := RandomShardInfo(s.ShardID, rangeID)
   156  
   157  	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
   158  		ShardID:          s.ShardID,
   159  		InitialShardInfo: shardInfo,
   160  	})
   161  	s.NoError(err)
   162  	s.ProtoEqual(shardInfo, resp.ShardInfo)
   163  
   164  	updateShardInfo := RandomShardInfo(s.ShardID, rangeID+1)
   165  	err = s.ShardManager.UpdateShard(s.Ctx, &p.UpdateShardRequest{
   166  		ShardInfo:       updateShardInfo,
   167  		PreviousRangeID: rangeID,
   168  	})
   169  	s.NoError(err)
   170  
   171  	resp, err = s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
   172  		ShardID:          s.ShardID,
   173  		InitialShardInfo: shardInfo,
   174  	})
   175  	s.NoError(err)
   176  	s.ProtoEqual(updateShardInfo, resp.ShardInfo)
   177  }