github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/aggregator/integration/placement.go (about)

     1  // Copyright (c) 2018 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 integration
    22  
    23  import (
    24  	"math"
    25  	"testing"
    26  
    27  	clusterclient "github.com/m3db/m3/src/cluster/client"
    28  	"github.com/m3db/m3/src/cluster/placement"
    29  	"github.com/m3db/m3/src/cluster/services"
    30  	"github.com/m3db/m3/src/cluster/shard"
    31  
    32  	"github.com/stretchr/testify/require"
    33  )
    34  
    35  type placementInstanceConfig struct {
    36  	instanceID          string
    37  	shardSetID          uint32
    38  	shardStartInclusive uint32
    39  	shardEndExclusive   uint32
    40  }
    41  
    42  func (c *placementInstanceConfig) newPlacementInstance() placement.Instance {
    43  	numShards := c.shardEndExclusive - c.shardStartInclusive
    44  	shardSet := make([]shard.Shard, 0, numShards)
    45  	for shardID := c.shardStartInclusive; shardID < c.shardEndExclusive; shardID++ {
    46  		shard := shard.NewShard(shardID).
    47  			SetState(shard.Available).
    48  			SetCutoverNanos(0).
    49  			SetCutoffNanos(math.MaxInt64)
    50  		shardSet = append(shardSet, shard)
    51  	}
    52  	shards := shard.NewShards(shardSet)
    53  	return placement.NewInstance().
    54  		SetID(c.instanceID).
    55  		SetShards(shards).
    56  		SetShardSetID(c.shardSetID).
    57  		SetEndpoint(c.instanceID).
    58  		SetWeight(1)
    59  }
    60  
    61  func newPlacement(numShards int, instances []placement.Instance) placement.Placement {
    62  	shards := make([]uint32, numShards)
    63  	for i := 0; i < numShards; i++ {
    64  		shards[i] = uint32(i)
    65  	}
    66  	var maxShardSetID uint32
    67  	for _, inst := range instances {
    68  		if maxShardSetID < inst.ShardSetID() {
    69  			maxShardSetID = inst.ShardSetID()
    70  		}
    71  	}
    72  	return placement.NewPlacement().
    73  		SetInstances(instances).
    74  		SetShards(shards).
    75  		SetIsSharded(true).
    76  		SetMaxShardSetID(maxShardSetID).
    77  		SetReplicaFactor(1)
    78  }
    79  
    80  func setPlacement(
    81  	t *testing.T,
    82  	key string,
    83  	clusterClient clusterclient.Client,
    84  	pl placement.Placement,
    85  ) {
    86  	svcs, err := clusterClient.Services(nil)
    87  	require.NoError(t, err)
    88  	ps, err := svcs.PlacementService(services.NewServiceID().SetName(defaultServiceName), placement.NewOptions())
    89  	require.NoError(t, err)
    90  	_, err = ps.Set(pl)
    91  	require.NoError(t, err)
    92  
    93  	store, err := clusterClient.KV()
    94  	require.NoError(t, err)
    95  
    96  	stagedPlacement, err := placement.NewPlacementsFromLatest(pl)
    97  	require.NoError(t, err)
    98  	stagedPlacementProto, err := stagedPlacement.Proto()
    99  	require.NoError(t, err)
   100  	_, err = store.Set(key, stagedPlacementProto)
   101  	require.NoError(t, err)
   102  }