github.com/ava-labs/avalanchego@v1.11.11/pubsub/filter_param.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package pubsub
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/ava-labs/avalanchego/pubsub/bloom"
    10  	"github.com/ava-labs/avalanchego/utils/set"
    11  )
    12  
    13  type FilterParam struct {
    14  	lock   sync.RWMutex
    15  	set    set.Set[string]
    16  	filter bloom.Filter
    17  }
    18  
    19  func NewFilterParam() *FilterParam {
    20  	return &FilterParam{
    21  		set: set.Set[string]{},
    22  	}
    23  }
    24  
    25  func (f *FilterParam) NewSet() {
    26  	f.lock.Lock()
    27  	defer f.lock.Unlock()
    28  
    29  	f.set = set.Set[string]{}
    30  	f.filter = nil
    31  }
    32  
    33  func (f *FilterParam) Filter() bloom.Filter {
    34  	f.lock.RLock()
    35  	defer f.lock.RUnlock()
    36  
    37  	return f.filter
    38  }
    39  
    40  func (f *FilterParam) SetFilter(filter bloom.Filter) bloom.Filter {
    41  	f.lock.Lock()
    42  	defer f.lock.Unlock()
    43  
    44  	f.filter = filter
    45  	f.set = nil
    46  	return f.filter
    47  }
    48  
    49  func (f *FilterParam) Check(addr []byte) bool {
    50  	f.lock.RLock()
    51  	defer f.lock.RUnlock()
    52  
    53  	if f.filter != nil && f.filter.Check(addr) {
    54  		return true
    55  	}
    56  	return f.set.Contains(string(addr))
    57  }
    58  
    59  func (f *FilterParam) Add(bl ...[]byte) error {
    60  	filter := f.Filter()
    61  	if filter != nil {
    62  		filter.Add(bl...)
    63  		return nil
    64  	}
    65  
    66  	f.lock.Lock()
    67  	defer f.lock.Unlock()
    68  
    69  	if f.set == nil {
    70  		return ErrFilterNotInitialized
    71  	}
    72  
    73  	if len(f.set)+len(bl) > MaxAddresses {
    74  		return ErrAddressLimit
    75  	}
    76  	for _, b := range bl {
    77  		f.set.Add(string(b))
    78  	}
    79  	return nil
    80  }
    81  
    82  func (f *FilterParam) Len() int {
    83  	f.lock.RLock()
    84  	defer f.lock.RUnlock()
    85  
    86  	return len(f.set)
    87  }