go.chromium.org/luci@v0.0.0-20250314024836-d9a61d0730e6/tokenserver/appengine/impl/utils/shards/shards.go (about)

     1  // Copyright 2016 The LUCI Authors.
     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 shards provides a low level support for implementing sharded set of
    16  // []byte blobs.
    17  package shards
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/gob"
    22  	"hash/fnv"
    23  	"sort"
    24  )
    25  
    26  // Shard is a set of byte blobs.
    27  //
    28  // It represents s single shard of a sharded set.
    29  type Shard map[string]struct{}
    30  
    31  // ParseShard deserializes a shard (serialized with Serialize).
    32  func ParseShard(blob []byte) (Shard, error) {
    33  	var sorted []string
    34  	dec := gob.NewDecoder(bytes.NewReader(blob))
    35  	if err := dec.Decode(&sorted); err != nil {
    36  		return nil, err
    37  	}
    38  	shard := make(Shard, len(sorted))
    39  	for _, blob := range sorted {
    40  		shard[blob] = struct{}{}
    41  	}
    42  	return shard, nil
    43  }
    44  
    45  // Serialize serializes the shard to a byte buffer.
    46  func (s Shard) Serialize() []byte {
    47  	sorted := make([]string, 0, len(s))
    48  	for blob := range s {
    49  		sorted = append(sorted, blob)
    50  	}
    51  	sort.Strings(sorted)
    52  	out := bytes.Buffer{}
    53  	enc := gob.NewEncoder(&out)
    54  	err := enc.Encode(sorted)
    55  	if err != nil {
    56  		panic("impossible error when encoding []string")
    57  	}
    58  	return out.Bytes()
    59  }
    60  
    61  // Set is an array of shards (representing a single sharded set).
    62  //
    63  // The size of the array is number of shards in the set. Allocate it using
    64  // regular make(...).
    65  type Set []Shard
    66  
    67  // Insert adds a blob into the sharded set.
    68  func (s Set) Insert(blob []byte) {
    69  	shard := &s[ShardIndex(blob, len(s))]
    70  	if *shard == nil {
    71  		*shard = make(Shard)
    72  	}
    73  	(*shard)[string(blob)] = struct{}{}
    74  }
    75  
    76  // ShardIndex returns an index of a shard to use when storing given blob.
    77  func ShardIndex(member []byte, shardCount int) int {
    78  	hash := fnv.New32()
    79  	hash.Write(member)
    80  	return int(hash.Sum32() % uint32(shardCount))
    81  }