github.com/ethersphere/bee/v2@v2.2.0/pkg/skippeers/skippeers.go (about)

     1  // Copyright 2023 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package skippeers
     6  
     7  import (
     8  	"math"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/ethersphere/bee/v2/pkg/swarm"
    13  )
    14  
    15  const maxDuration time.Duration = math.MaxInt64
    16  
    17  type List struct {
    18  	mtx sync.Mutex
    19  
    20  	durC chan time.Duration
    21  	quit chan struct{}
    22  	// key is chunk address, value is map of peer address to expiration
    23  	skip map[string]map[string]int64
    24  
    25  	wg sync.WaitGroup
    26  }
    27  
    28  func NewList() *List {
    29  	l := &List{
    30  		skip: make(map[string]map[string]int64),
    31  		durC: make(chan time.Duration),
    32  		quit: make(chan struct{}),
    33  	}
    34  
    35  	l.wg.Add(1)
    36  	go l.worker()
    37  
    38  	return l
    39  }
    40  
    41  func (l *List) worker() {
    42  
    43  	defer l.wg.Done()
    44  
    45  	ticker := time.NewTicker(time.Minute)
    46  	defer ticker.Stop()
    47  
    48  	for {
    49  		select {
    50  		case <-ticker.C:
    51  			l.prune()
    52  		case <-l.quit:
    53  			return
    54  		}
    55  	}
    56  }
    57  
    58  func (l *List) Forever(chunk, peer swarm.Address) {
    59  	l.Add(chunk, peer, maxDuration)
    60  }
    61  
    62  func (l *List) Add(chunk, peer swarm.Address, expire time.Duration) {
    63  
    64  	l.mtx.Lock()
    65  	defer l.mtx.Unlock()
    66  
    67  	var t int64
    68  
    69  	if expire == maxDuration {
    70  		t = maxDuration.Nanoseconds()
    71  	} else {
    72  		t = time.Now().Add(expire).UnixNano()
    73  	}
    74  
    75  	if _, ok := l.skip[chunk.ByteString()]; !ok {
    76  		l.skip[chunk.ByteString()] = make(map[string]int64)
    77  	}
    78  
    79  	l.skip[chunk.ByteString()][peer.ByteString()] = t
    80  }
    81  
    82  func (l *List) ChunkPeers(ch swarm.Address) (peers []swarm.Address) {
    83  	l.mtx.Lock()
    84  	defer l.mtx.Unlock()
    85  
    86  	now := time.Now().UnixNano()
    87  
    88  	if p, ok := l.skip[ch.ByteString()]; ok {
    89  		for peer, exp := range p {
    90  			if exp > now {
    91  				peers = append(peers, swarm.NewAddress([]byte(peer)))
    92  			}
    93  		}
    94  	}
    95  
    96  	return peers
    97  }
    98  
    99  func (l *List) PruneExpiresAfter(ch swarm.Address, d time.Duration) int {
   100  	l.mtx.Lock()
   101  	defer l.mtx.Unlock()
   102  	return l.pruneChunk(ch.ByteString(), time.Now().Add(d).UnixNano())
   103  }
   104  
   105  func (l *List) prune() {
   106  	l.mtx.Lock()
   107  	defer l.mtx.Unlock()
   108  	expiresNano := time.Now().UnixNano()
   109  	for k := range l.skip {
   110  		l.pruneChunk(k, expiresNano)
   111  	}
   112  }
   113  
   114  // Must be called under lock
   115  func (l *List) pruneChunk(ch string, now int64) int {
   116  
   117  	count := 0
   118  
   119  	for peer, exp := range l.skip[ch] {
   120  		if exp <= now {
   121  			delete(l.skip[ch], peer)
   122  			count++
   123  		}
   124  	}
   125  	if len(l.skip[ch]) == 0 {
   126  		delete(l.skip, ch)
   127  	}
   128  
   129  	return count
   130  }
   131  
   132  func (l *List) Close() error {
   133  	close(l.quit)
   134  	l.wg.Wait()
   135  
   136  	l.mtx.Lock()
   137  	clear(l.skip)
   138  	l.mtx.Unlock()
   139  
   140  	return nil
   141  }