add io.Closer interface and implementation to Peerstore.

interface.go

@@ -3,6 +3,7 @@ package peerstore
 import (
 	"context"
 	"errors"
+	"io"
 	"math"
 	"time"
 
@@ -49,6 +50,8 @@ const (
 // Peerstore provides a threadsafe store of Peer related
 // information.
 type Peerstore interface {
+	io.Closer
+
 	AddrBook
 	KeyBook
 	PeerMetadata

peerstore.go

@@ -2,6 +2,7 @@ package peerstore
 
 import (
 	"fmt"
+	"io"
 	"sync"
 
 	peer "github.com/libp2p/go-libp2p-peer"
@@ -31,6 +32,19 @@ func NewPeerstore(kb KeyBook, ab AddrBook, md PeerMetadata) Peerstore {
 	}
 }
 
+func (ps *peerstore) Close() (err error) {
+	if cl, ok := ps.KeyBook.(io.Closer); ok {
+		cl.Close()
+	}
+	if cl, ok := ps.AddrBook.(io.Closer); ok {
+		cl.Close()
+	}
+	if cl, ok := ps.PeerMetadata.(io.Closer); ok {
+		cl.Close()
+	}
+	return nil
+}
+
 func (ps *peerstore) Peers() peer.IDSlice {
 	set := map[peer.ID]struct{}{}
 	for _, p := range ps.PeersWithKeys() {

pstoreds/addr_book.go

@@ -175,9 +175,10 @@ func NewAddrBook(ctx context.Context, store ds.Batching, opts Options) (ab *dsAd
 	return ab, nil
 }
 
-func (ab *dsAddrBook) Close() {
+func (ab *dsAddrBook) Close() error {
 	ab.cancelFn()
 	ab.childrenDone.Wait()
+	return nil
 }
 
 // loadRecord is a read-through fetch. It fetches a record from cache, falling back to the

pstoreds/ds_test.go

@@ -19,7 +19,6 @@ type datastoreFactory func(tb testing.TB) (ds.Batching, func())
 
 var dstores = map[string]datastoreFactory{
 	"Badger": badgerStore,
-	// TODO: Enable once go-ds-leveldb supports TTL via a shim.
 	// "Leveldb": leveldbStore,
 }
 
@@ -122,42 +121,41 @@ func leveldbStore(tb testing.TB) (ds.TxnDatastore, func()) {
 
 func peerstoreFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.PeerstoreFactory {
 	return func() (pstore.Peerstore, func()) {
-		store, closeFunc := storeFactory(tb)
-
+		store, storeCloseFn := storeFactory(tb)
 		ps, err := NewPeerstore(context.Background(), store, opts)
 		if err != nil {
 			tb.Fatal(err)
 		}
-
-		return ps, closeFunc
+		closer := func() {
+			ps.Close()
+			storeCloseFn()
+		}
+		return ps, closer
 	}
 }
 
 func addressBookFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.AddrBookFactory {
 	return func() (pstore.AddrBook, func()) {
 		store, closeFunc := storeFactory(tb)
-
 		ab, err := NewAddrBook(context.Background(), store, opts)
 		if err != nil {
 			tb.Fatal(err)
 		}
-
-		return ab, func() {
+		closer := func() {
 			ab.Close()
 			closeFunc()
 		}
+		return ab, closer
 	}
 }
 
 func keyBookFactory(tb testing.TB, storeFactory datastoreFactory, opts Options) pt.KeyBookFactory {
 	return func() (pstore.KeyBook, func()) {
-		store, closeFunc := storeFactory(tb)
-
+		store, storeCloseFn := storeFactory(tb)
 		kb, err := NewKeyBook(context.Background(), store, opts)
 		if err != nil {
 			tb.Fatal(err)
 		}
-
-		return kb, closeFunc
+		return kb, storeCloseFn
 	}
 }