Merge pull request #184 from libp2p/pstore-cleanup

when passed an event bus, automatically clean up disconnected peers

go.mod

@@ -6,12 +6,14 @@ retract v0.2.9 // Contains backwards-incompatible changes. Use v0.3.0 instead.
 
 require (
 	github.com/gogo/protobuf v1.3.2
+	github.com/golang/mock v1.6.0
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/ipfs/go-datastore v0.5.0
 	github.com/ipfs/go-ds-badger v0.3.0
 	github.com/ipfs/go-ds-leveldb v0.5.0
 	github.com/ipfs/go-log/v2 v2.3.0
 	github.com/libp2p/go-buffer-pool v0.0.2
+	github.com/libp2p/go-eventbus v0.2.1
 	github.com/libp2p/go-libp2p-core v0.12.0
 	github.com/multiformats/go-base32 v0.0.3
 	github.com/multiformats/go-multiaddr v0.3.3

pstoreds/peerstore.go

@@ -6,6 +6,10 @@ import (
 	"io"
 	"time"
 
+	"github.com/libp2p/go-libp2p-core/event"
+
+	"github.com/libp2p/go-libp2p-peerstore/pstoremanager"
+
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
 	pstore "github.com/libp2p/go-libp2p-peerstore"
@@ -23,6 +27,13 @@ type Options struct {
 	// MaxProtocols is the maximum number of protocols we store for one peer.
 	MaxProtocols int
 
+	// The EventBus that is used to subscribe to EvtPeerConnectednessChanged events.
+	// This allows the automatic clean up when a peer disconnect.
+	// This configuration option is optional. If no EventBus is set, it's the callers
+	// responsibility to call RemovePeer to ensure that memory consumption of the
+	// peerstore doesn't grow unboundedly.
+	EventBus event.Bus
+
 	// Sweep interval to purge expired addresses from the datastore. If this is a zero value, GC will not run
 	// automatically, but it'll be available on demand via explicit calls.
 	GCPurgeInterval time.Duration
@@ -60,9 +71,18 @@ type pstoreds struct {
 	*dsAddrBook
 	*dsProtoBook
 	*dsPeerMetadata
+
+	manager *pstoremanager.PeerstoreManager
 }
 
+var _ peerstore.Peerstore = &pstoreds{}
+
 // NewPeerstore creates a peerstore backed by the provided persistent datastore.
+// It is recommended to construct the peerstore with an event bus, using the WithEventBus option.
+// In that case, the peerstore will automatically perform cleanups when a peer disconnects
+// (see the pstoremanager package for details).
+// If constructed without an event bus, it's the caller's responsibility to call RemovePeer to ensure
+// that memory consumption of the peerstore doesn't grow unboundedly.
 func NewPeerstore(ctx context.Context, store ds.Batching, opts Options) (*pstoreds, error) {
 	addrBook, err := NewAddrBook(ctx, store, opts)
 	if err != nil {
@@ -91,6 +111,14 @@ func NewPeerstore(ctx context.Context, store ds.Batching, opts Options) (*pstore
 		dsPeerMetadata: peerMetadata,
 		dsProtoBook:    protoBook,
 	}
+	if opts.EventBus != nil {
+		manager, err := pstoremanager.NewPeerstoreManager(ps, opts.EventBus)
+		if err != nil {
+			ps.Close()
+			return nil, err
+		}
+		ps.manager = manager
+	}
 	return ps, nil
 }
 
@@ -128,6 +156,10 @@ func uniquePeerIds(ds ds.Datastore, prefix ds.Key, extractor func(result query.R
 	return ids, nil
 }
 
+func (ps *pstoreds) Start() {
+	ps.manager.Start()
+}
+
 func (ps *pstoreds) Close() (err error) {
 	var errs []error
 	weakClose := func(name string, c interface{}) {
@@ -137,7 +169,9 @@ func (ps *pstoreds) Close() (err error) {
 			}
 		}
 	}
-
+	if ps.manager != nil {
+		weakClose("manager", ps.manager)
+	}
 	weakClose("keybook", ps.dsKeyBook)
 	weakClose("addressbook", ps.dsAddrBook)
 	weakClose("protobook", ps.dsProtoBook)