github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/cache/subnet_ids.go (about)

     1  package cache
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	lru "github.com/hashicorp/golang-lru"
     8  	"github.com/patrickmn/go-cache"
     9  	types "github.com/prysmaticlabs/eth2-types"
    10  	"github.com/prysmaticlabs/prysm/shared/params"
    11  	"github.com/prysmaticlabs/prysm/shared/sliceutil"
    12  )
    13  
    14  type subnetIDs struct {
    15  	attester          *lru.Cache
    16  	attesterLock      sync.RWMutex
    17  	aggregator        *lru.Cache
    18  	aggregatorLock    sync.RWMutex
    19  	persistentSubnets *cache.Cache
    20  	subnetsLock       sync.RWMutex
    21  }
    22  
    23  // SubnetIDs for attester and aggregator.
    24  var SubnetIDs = newSubnetIDs()
    25  
    26  func newSubnetIDs() *subnetIDs {
    27  	// Given a node can calculate committee assignments of current epoch and next epoch.
    28  	// Max size is set to 2 epoch length.
    29  	cacheSize := int(params.BeaconConfig().SlotsPerEpoch.Mul(params.BeaconConfig().MaxCommitteesPerSlot * 2))
    30  	attesterCache, err := lru.New(cacheSize)
    31  	if err != nil {
    32  		panic(err)
    33  	}
    34  	aggregatorCache, err := lru.New(cacheSize)
    35  	if err != nil {
    36  		panic(err)
    37  	}
    38  	epochDuration := time.Duration(params.BeaconConfig().SlotsPerEpoch.Mul(params.BeaconConfig().SecondsPerSlot))
    39  	subLength := epochDuration * time.Duration(params.BeaconConfig().EpochsPerRandomSubnetSubscription)
    40  	persistentCache := cache.New(subLength*time.Second, epochDuration*time.Second)
    41  	return &subnetIDs{attester: attesterCache, aggregator: aggregatorCache, persistentSubnets: persistentCache}
    42  }
    43  
    44  // AddAttesterSubnetID adds the subnet index for subscribing subnet for the attester of a given slot.
    45  func (s *subnetIDs) AddAttesterSubnetID(slot types.Slot, subnetID uint64) {
    46  	s.attesterLock.Lock()
    47  	defer s.attesterLock.Unlock()
    48  
    49  	ids := []uint64{subnetID}
    50  	val, exists := s.attester.Get(slot)
    51  	if exists {
    52  		ids = sliceutil.UnionUint64(append(val.([]uint64), ids...))
    53  	}
    54  	s.attester.Add(slot, ids)
    55  }
    56  
    57  // GetAttesterSubnetIDs gets the subnet IDs for subscribed subnets for attesters of the slot.
    58  func (s *subnetIDs) GetAttesterSubnetIDs(slot types.Slot) []uint64 {
    59  	s.attesterLock.RLock()
    60  	defer s.attesterLock.RUnlock()
    61  
    62  	val, exists := s.attester.Get(slot)
    63  	if !exists {
    64  		return nil
    65  	}
    66  	if v, ok := val.([]uint64); ok {
    67  		return v
    68  	}
    69  	return nil
    70  }
    71  
    72  // AddAggregatorSubnetID adds the subnet ID for subscribing subnet for the aggregator of a given slot.
    73  func (s *subnetIDs) AddAggregatorSubnetID(slot types.Slot, subnetID uint64) {
    74  	s.aggregatorLock.Lock()
    75  	defer s.aggregatorLock.Unlock()
    76  
    77  	ids := []uint64{subnetID}
    78  	val, exists := s.aggregator.Get(slot)
    79  	if exists {
    80  		ids = sliceutil.UnionUint64(append(val.([]uint64), ids...))
    81  	}
    82  	s.aggregator.Add(slot, ids)
    83  }
    84  
    85  // GetAggregatorSubnetIDs gets the subnet IDs for subscribing subnet for aggregator of the slot.
    86  func (s *subnetIDs) GetAggregatorSubnetIDs(slot types.Slot) []uint64 {
    87  	s.aggregatorLock.RLock()
    88  	defer s.aggregatorLock.RUnlock()
    89  
    90  	val, exists := s.aggregator.Get(slot)
    91  	if !exists {
    92  		return []uint64{}
    93  	}
    94  	return val.([]uint64)
    95  }
    96  
    97  // GetPersistentSubnets retrieves the persistent subnet and expiration time of that validator's
    98  // subscription.
    99  func (s *subnetIDs) GetPersistentSubnets(pubkey []byte) ([]uint64, bool, time.Time) {
   100  	s.subnetsLock.RLock()
   101  	defer s.subnetsLock.RUnlock()
   102  
   103  	id, duration, ok := s.persistentSubnets.GetWithExpiration(string(pubkey))
   104  	if !ok {
   105  		return []uint64{}, ok, time.Time{}
   106  	}
   107  	return id.([]uint64), ok, duration
   108  }
   109  
   110  // GetAllSubnets retrieves all the non-expired subscribed subnets of all the validators
   111  // in the cache.
   112  func (s *subnetIDs) GetAllSubnets() []uint64 {
   113  	s.subnetsLock.RLock()
   114  	defer s.subnetsLock.RUnlock()
   115  
   116  	itemsMap := s.persistentSubnets.Items()
   117  	var committees []uint64
   118  
   119  	for _, v := range itemsMap {
   120  		if v.Expired() {
   121  			continue
   122  		}
   123  		committees = append(committees, v.Object.([]uint64)...)
   124  	}
   125  	return sliceutil.SetUint64(committees)
   126  }
   127  
   128  // AddPersistentCommittee adds the relevant committee for that particular validator along with its
   129  // expiration period.
   130  func (s *subnetIDs) AddPersistentCommittee(pubkey []byte, comIndex []uint64, duration time.Duration) {
   131  	s.subnetsLock.Lock()
   132  	defer s.subnetsLock.Unlock()
   133  
   134  	s.persistentSubnets.Set(string(pubkey), comIndex, duration)
   135  }
   136  
   137  // EmptyAllCaches empties out all the related caches and flushes any stored
   138  // entries on them. This should only ever be used for testing, in normal
   139  // production, handling of the relevant subnets for each role is done
   140  // separately.
   141  func (s *subnetIDs) EmptyAllCaches() {
   142  	// Clear the caches.
   143  	s.attesterLock.Lock()
   144  	s.attester.Purge()
   145  	s.attesterLock.Unlock()
   146  
   147  	s.aggregatorLock.Lock()
   148  	s.aggregator.Purge()
   149  	s.aggregatorLock.Unlock()
   150  
   151  	s.subnetsLock.Lock()
   152  	s.persistentSubnets.Flush()
   153  	s.subnetsLock.Unlock()
   154  }