ceremonyclient/go-libp2p/p2p/host/pstoremanager/pstoremanager.go
Cassie Heart ebb31fd792 v1.4.19
2024-06-08 11:32:45 +00:00

143 lines
3.5 KiB
Go

package pstoremanager
import (
"context"
"sync"
"time"
"github.com/libp2p/go-libp2p/core/event"
"github.com/libp2p/go-libp2p/core/network"
"github.com/libp2p/go-libp2p/core/peer"
"github.com/libp2p/go-libp2p/core/peerstore"
"github.com/libp2p/go-libp2p/p2p/host/eventbus"
logging "github.com/ipfs/go-log/v2"
)
var log = logging.Logger("pstoremanager")
type Option func(*PeerstoreManager) error
// WithGracePeriod sets the grace period.
// If a peer doesn't reconnect during the grace period, its data is removed.
// Default: 1 minute.
func WithGracePeriod(p time.Duration) Option {
return func(m *PeerstoreManager) error {
m.gracePeriod = p
return nil
}
}
// WithCleanupInterval set the clean up interval.
// During a clean up run peers that disconnected before the grace period are removed.
// If unset, the interval is set to half the grace period.
func WithCleanupInterval(t time.Duration) Option {
return func(m *PeerstoreManager) error {
m.cleanupInterval = t
return nil
}
}
type PeerstoreManager struct {
pstore peerstore.Peerstore
eventBus event.Bus
network network.Network
cancel context.CancelFunc
refCount sync.WaitGroup
gracePeriod time.Duration
cleanupInterval time.Duration
}
func NewPeerstoreManager(pstore peerstore.Peerstore, eventBus event.Bus, network network.Network, opts ...Option) (*PeerstoreManager, error) {
m := &PeerstoreManager{
pstore: pstore,
gracePeriod: time.Minute,
eventBus: eventBus,
network: network,
}
for _, opt := range opts {
if err := opt(m); err != nil {
return nil, err
}
}
if m.cleanupInterval == 0 {
m.cleanupInterval = m.gracePeriod / 2
}
return m, nil
}
func (m *PeerstoreManager) Start() {
ctx, cancel := context.WithCancel(context.Background())
m.cancel = cancel
sub, err := m.eventBus.Subscribe(&event.EvtPeerConnectednessChanged{}, eventbus.Name("pstoremanager"))
if err != nil {
log.Warnf("subscription failed. Peerstore manager not activated. Error: %s", err)
return
}
m.refCount.Add(1)
go m.background(ctx, sub)
}
func (m *PeerstoreManager) background(ctx context.Context, sub event.Subscription) {
defer m.refCount.Done()
defer sub.Close()
disconnected := make(map[peer.ID]time.Time)
ticker := time.NewTicker(m.cleanupInterval)
defer ticker.Stop()
defer func() {
for p := range disconnected {
m.pstore.RemovePeer(p)
}
}()
for {
select {
case e, ok := <-sub.Out():
if !ok {
return
}
ev := e.(event.EvtPeerConnectednessChanged)
p := ev.Peer
switch ev.Connectedness {
case network.Connected, network.Limited:
// If we reconnect to the peer before we've cleared the information,
// keep it. This is an optimization to keep the disconnected map
// small. We still need to check that a peer is actually
// disconnected before removing it from the peer store.
delete(disconnected, p)
default:
if _, ok := disconnected[p]; !ok {
disconnected[p] = time.Now()
}
}
case <-ticker.C:
now := time.Now()
for p, disconnectTime := range disconnected {
if disconnectTime.Add(m.gracePeriod).Before(now) {
// Check that the peer is actually not connected at this point.
// This avoids a race condition where the Connected notification
// is processed after this time has fired.
if m.network.Connectedness(p) != network.Connected {
m.pstore.RemovePeer(p)
delete(disconnected, p)
}
}
}
case <-ctx.Done():
return
}
}
}
func (m *PeerstoreManager) Close() error {
if m.cancel != nil {
m.cancel()
}
m.refCount.Wait()
return nil
}