github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/memoryengine/shard.go (about)

     1  // Copyright 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package memoryengine
    16  
    17  import (
    18  	"context"
    19  	"sort"
    20  	"sync"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/common/mpool"
    23  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    24  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    25  	"github.com/matrixorigin/matrixone/pkg/pb/metadata"
    26  	"github.com/matrixorigin/matrixone/pkg/vm/engine"
    27  )
    28  
    29  type Shard = metadata.TNShard
    30  
    31  type shardsFunc = func() ([]Shard, error)
    32  
    33  type getDefsFunc = func(context.Context) ([]engine.TableDef, error)
    34  
    35  func NewDefaultShardPolicy(mp *mpool.MPool) ShardPolicy {
    36  	return FallbackShard{
    37  		NewHashShard(mp),
    38  		new(NoShard),
    39  	}
    40  }
    41  
    42  type ShardPolicy interface {
    43  	Vector(
    44  		ctx context.Context,
    45  		tableID ID,
    46  		getDefs getDefsFunc,
    47  		colName string,
    48  		vec *vector.Vector,
    49  		nodes []metadata.TNService,
    50  	) (
    51  		sharded []*ShardedVector,
    52  		err error,
    53  	)
    54  
    55  	Batch(
    56  		ctx context.Context,
    57  		tableID ID,
    58  		getDefs getDefsFunc,
    59  		batch *batch.Batch,
    60  		nodes []metadata.TNService,
    61  	) (
    62  		sharded []*ShardedBatch,
    63  		err error,
    64  	)
    65  }
    66  
    67  type ShardedVector struct {
    68  	Shard  Shard
    69  	Vector *vector.Vector
    70  }
    71  
    72  type ShardedBatch struct {
    73  	Shard Shard
    74  	Batch *batch.Batch
    75  }
    76  
    77  func (e *Engine) allShards() (shards []Shard, err error) {
    78  	for _, store := range getTNServices(e.cluster) {
    79  		for _, shard := range store.Shards {
    80  			shards = append(shards, Shard{
    81  				TNShardRecord: metadata.TNShardRecord{
    82  					ShardID: shard.ShardID,
    83  				},
    84  				ReplicaID: shard.ReplicaID,
    85  				Address:   store.TxnServiceAddress,
    86  			})
    87  		}
    88  	}
    89  	return
    90  }
    91  
    92  func (e *Engine) anyShard() (shards []Shard, err error) {
    93  	for _, store := range getTNServices(e.cluster) {
    94  		for _, shard := range store.Shards {
    95  			shards = append(shards, Shard{
    96  				TNShardRecord: metadata.TNShardRecord{
    97  					ShardID: shard.ShardID,
    98  				},
    99  				ReplicaID: shard.ReplicaID,
   100  				Address:   store.TxnServiceAddress,
   101  			})
   102  			return
   103  		}
   104  	}
   105  	return
   106  }
   107  
   108  func thisShard(shard Shard) func() ([]Shard, error) {
   109  	shards := []Shard{shard}
   110  	return func() ([]Shard, error) {
   111  		return shards, nil
   112  	}
   113  }
   114  
   115  func theseShards(shards []Shard) func() ([]Shard, error) {
   116  	return func() ([]Shard, error) {
   117  		return shards, nil
   118  	}
   119  }
   120  
   121  // NoShard doesn't do sharding at all
   122  type NoShard struct {
   123  	setOnce sync.Once
   124  	shard   Shard
   125  }
   126  
   127  func (s *NoShard) setShard(stores []metadata.TNService) {
   128  	s.setOnce.Do(func() {
   129  		type ShardInfo struct {
   130  			Store metadata.TNService
   131  			Shard metadata.TNShard
   132  		}
   133  		infos := make([]ShardInfo, 0, len(stores))
   134  		for _, store := range stores {
   135  			for _, info := range store.Shards {
   136  				infos = append(infos, ShardInfo{
   137  					Store: store,
   138  					Shard: info,
   139  				})
   140  			}
   141  		}
   142  		if len(infos) == 0 {
   143  			panic("no shard")
   144  		}
   145  		sort.Slice(infos, func(i, j int) bool {
   146  			return infos[i].Shard.ShardID < infos[j].Shard.ShardID
   147  		})
   148  		info := infos[0]
   149  		s.shard = Shard{
   150  			TNShardRecord: metadata.TNShardRecord{
   151  				ShardID: info.Shard.ShardID,
   152  			},
   153  			ReplicaID: info.Shard.ReplicaID,
   154  			Address:   info.Store.TxnServiceAddress,
   155  		}
   156  	})
   157  }
   158  
   159  var _ ShardPolicy = new(NoShard)
   160  
   161  func (s *NoShard) Vector(
   162  	ctx context.Context,
   163  	tableID ID,
   164  	getDefs getDefsFunc,
   165  	colName string,
   166  	vec *vector.Vector,
   167  	nodes []metadata.TNService,
   168  ) (
   169  	sharded []*ShardedVector,
   170  	err error,
   171  ) {
   172  	s.setShard(nodes)
   173  	sharded = append(sharded, &ShardedVector{
   174  		Shard:  s.shard,
   175  		Vector: vec,
   176  	})
   177  	return
   178  }
   179  
   180  func (s *NoShard) Batch(
   181  	ctx context.Context,
   182  	tableID ID,
   183  	getDefs getDefsFunc,
   184  	bat *batch.Batch,
   185  	nodes []metadata.TNService,
   186  ) (
   187  	sharded []*ShardedBatch,
   188  	err error,
   189  ) {
   190  	s.setShard(nodes)
   191  	sharded = append(sharded, &ShardedBatch{
   192  		Shard: s.shard,
   193  		Batch: bat,
   194  	})
   195  	return
   196  }