github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/spanz/hash_map.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package spanz
    15  
    16  import (
    17  	"encoding/binary"
    18  
    19  	"blainsmith.com/go/seahash"
    20  	"github.com/pingcap/tiflow/cdc/processor/tablepb"
    21  )
    22  
    23  // HashMap is a specialized hash map that map a Span to a value.
    24  type HashMap[T any] struct {
    25  	hashMap map[hashableSpan]T
    26  }
    27  
    28  // NewHashMap returns a new HashMap.
    29  func NewHashMap[T any]() *HashMap[T] {
    30  	return &HashMap[T]{
    31  		hashMap: make(map[hashableSpan]T),
    32  	}
    33  }
    34  
    35  // Len returns the number of items currently in the map.
    36  func (m *HashMap[T]) Len() int {
    37  	return len(m.hashMap)
    38  }
    39  
    40  // Has returns true if the given key is in the map.
    41  func (m *HashMap[T]) Has(span tablepb.Span) bool {
    42  	_, ok := m.hashMap[toHashableSpan(span)]
    43  	return ok
    44  }
    45  
    46  // Get looks for the key item in the map, returning it.
    47  // It returns (zeroValue, false) if unable to find that item.
    48  func (m *HashMap[T]) Get(span tablepb.Span) (T, bool) {
    49  	item, ok := m.hashMap[toHashableSpan(span)]
    50  	return item, ok
    51  }
    52  
    53  // GetV looks for the key item in the map, returning it.
    54  // It returns zeroValue if unable to find that item.
    55  func (m *HashMap[T]) GetV(span tablepb.Span) T {
    56  	item := m.hashMap[toHashableSpan(span)]
    57  	return item
    58  }
    59  
    60  // Delete removes an item whose key equals to the span.
    61  func (m *HashMap[T]) Delete(span tablepb.Span) {
    62  	delete(m.hashMap, toHashableSpan(span))
    63  }
    64  
    65  // ReplaceOrInsert adds the given item to the map.
    66  func (m *HashMap[T]) ReplaceOrInsert(span tablepb.Span, value T) {
    67  	m.hashMap[toHashableSpan(span)] = value
    68  }
    69  
    70  // Range calls the iterator for every value in the map until iterator returns
    71  // false.
    72  func (m *HashMap[T]) Range(iterator ItemIterator[T]) {
    73  	for k, v := range m.hashMap {
    74  		ok := iterator(k.toSpan(), v)
    75  		if !ok {
    76  			break
    77  		}
    78  	}
    79  }
    80  
    81  // HashTableSpan hashes the given span to a slot offset.
    82  func HashTableSpan(span tablepb.Span, slots int) int {
    83  	b := make([]byte, 8+len(span.StartKey))
    84  	binary.LittleEndian.PutUint64(b[0:8], uint64(span.TableID))
    85  	copy(b[8:], span.StartKey)
    86  	return int(seahash.Sum64(b) % uint64(slots))
    87  }