Add PersistentDatastore interface and implement it on wrappers

This adds a PersistentDatastore interface which allows datastores
to report DiskUsage().

It implementes the interface on all wrapper types, which return
0 if the wrapped datastore does not provide this method.

Related: https://github.com/ipfs/go-ipfs/pull/4550

basic_ds.go

@@ -166,6 +166,13 @@ func (d *LogDatastore) Delete(key Key) (err error) {
 	return d.child.Delete(key)
 }
 
+// DiskUsage implements the PersistentDatastore interface.
+func (d *LogDatastore) DiskUsage() (uint64, error) {
+	du, err := DiskUsage(d.child)
+	log.Printf("%s: DiskUsage: %d\n", d.Name, du)
+	return du, err
+}
+
 // Query implements Datastore.Query
 func (d *LogDatastore) Query(q dsq.Query) (dsq.Results, error) {
 	log.Printf("%s: Query\n", d.Name)

callback/callback.go

@@ -40,3 +40,9 @@ func (c *Datastore) Query(q dsq.Query) (dsq.Results, error) {
 	c.F()
 	return c.D.Query(q)
 }
+
+// DiskUsage implements the PersistentDatastore interface.
+func (c *Datastore) DiskUsage() (uint64, error) {
+	c.F()
+	return ds.DiskUsage(c.D)
+}

coalesce/coalesce.go

@@ -138,3 +138,8 @@ func (d *datastore) Close() error {
 	}
 	return nil
 }
+
+// DiskUsage implements the PersistentDatastore interface.
+func (d *datastore) DiskUsage() (uint64, error) {
+	return ds.DiskUsage(d.child)
+}

datastore.go

@@ -95,7 +95,7 @@ type CheckedDatastore interface {
 
 // CheckedDatastore is an interface that should be implemented by datastores
 // which want to provide a mechanism to check data integrity and/or
-// error correction
+// error correction.
 type ScrubbedDatastore interface {
 	Datastore
 
@@ -103,13 +103,33 @@ type ScrubbedDatastore interface {
 }
 
 // GCDatastore is an interface that should be implemented by datastores which
-// don't free disk space by just removing data from them
+// don't free disk space by just removing data from them.
 type GCDatastore interface {
 	Datastore
 
 	CollectGarbage() error
 }
 
+// PersistentDatastore is an interface that should be implemented by datastores
+// which can report disk usage.
+type PersistentDatastore interface {
+	Datastore
+
+	// DiskUsage returns the space used by a datastore, in bytes.
+	DiskUsage() (uint64, error)
+}
+
+// DiskUsage checks if a Datastore is a
+// PersistentDatastore and returns its DiskUsage(),
+// otherwise returns 0.
+func DiskUsage(d Datastore) (uint64, error) {
+	persDs, ok := d.(PersistentDatastore)
+	if !ok {
+		return 0, nil
+	}
+	return persDs.DiskUsage()
+}
+
 // Errors
 
 // ErrNotFound is returned by Get, Has, and Delete when a datastore does not

fs/fs.go

 // Package fs is a simple Datastore implementation that stores keys
-// are directories and files, mirroring the key. That is, the key
+// as directories and files, mirroring the key. That is, the key
 // "/foo/bar" is stored as file "PATH/foo/bar/.dsobject".
 //
 // This means key some segments will not work. For example, the
@@ -20,6 +20,7 @@ package fs
 import (
 	"fmt"
 	"io/ioutil"
+	"log"
 	"os"
 	"path/filepath"
 	"strings"
@@ -157,3 +158,19 @@ func (d *Datastore) Close() error {
 func (d *Datastore) Batch() (ds.Batch, error) {
 	return ds.NewBasicBatch(d), nil
 }
+
+// DiskUsage returns the disk size used by the datastore in bytes.
+func (d *Datastore) DiskUsage() (uint64, error) {
+	var du uint64
+	err := filepath.Walk(d.path, func(p string, f os.FileInfo, err error) error {
+		if err != nil {
+			log.Println(err)
+			return err
+		}
+		if f != nil {
+			du += uint64(f.Size())
+		}
+		return nil
+	})
+	return du, err
+}

fs/fs_test.go

@@ -87,6 +87,30 @@ func (ks *DSSuite) TestBasic(c *C) {
 	}
 }
 
+func (ks *DSSuite) TestDiskUsage(c *C) {
+	keys := strsToKeys([]string{
+		"foo",
+		"foo/bar",
+		"foo/bar/baz",
+		"foo/barb",
+		"foo/bar/bazb",
+		"foo/bar/baz/barb",
+	})
+
+	for _, k := range keys {
+		err := ks.ds.Put(k, []byte(k.String()))
+		c.Check(err, Equals, nil)
+	}
+
+	if ps, ok := ks.ds.(ds.PersistentDatastore); ok {
+		if s, err := ps.DiskUsage(); s <= 100 || err != nil {
+			c.Error("unexpected size is: ", s)
+		}
+	} else {
+		c.Error("should implement PersistentDatastore")
+	}
+}
+
 func strsToKeys(strs []string) []ds.Key {
 	keys := make([]ds.Key, len(strs))
 	for i, s := range strs {

keytransform/keytransform.go

@@ -84,6 +84,11 @@ func (d *ktds) Close() error {
 	return nil
 }
 
+// DiskUsage implements the PersistentDatastore interface.
+func (d *ktds) DiskUsage() (uint64, error) {
+	return ds.DiskUsage(d.child)
+}
+
 func (d *ktds) Batch() (ds.Batch, error) {
 	bds, ok := d.child.(ds.Batching)
 	if !ok {

mount/mount.go

@@ -194,6 +194,19 @@ func (d *Datastore) Close() error {
 	return nil
 }
 
+// DiskUsage implements the PersistentDatastore interface.
+func (d *Datastore) DiskUsage() (uint64, error) {
+	var duTotal uint64 = 0
+	for _, d := range d.mounts {
+		du, err := datastore.DiskUsage(d.Datastore)
+		duTotal += du
+		if err != nil {
+			return duTotal, err
+		}
+	}
+	return duTotal, nil
+}
+
 type mountBatch struct {
 	mounts map[string]datastore.Batch
 

panic/panic.go

@@ -79,6 +79,15 @@ func (d *datastore) Close() error {
 	return nil
 }
 
+// DiskUsage implements the PersistentDatastore interface.
+func (d *datastore) DiskUsage() (uint64, error) {
+	du, err := ds.DiskUsage(d.child)
+	if err != nil {
+		panic(err)
+	}
+	return du, nil
+}
+
 func (d *datastore) Batch() (ds.Batch, error) {
 	b, err := d.child.(ds.Batching).Batch()
 	if err != nil {

sync/sync.go

@@ -92,6 +92,13 @@ func (d *MutexDatastore) Close() error {
 	return nil
 }
 
+// DiskUsage implements the PersistentDatastore interface.
+func (d *MutexDatastore) DiskUsage() (uint64, error) {
+	d.RLock()
+	defer d.RUnlock()
+	return ds.DiskUsage(d.child)
+}
+
 type syncBatch struct {
 	batch ds.Batch
 	mds   *MutexDatastore

syncmount/mount.go

@@ -203,6 +203,20 @@ func (d *Datastore) Close() error {
 	return nil
 }
 
+// DiskUsage returns the sum of DiskUsages for the mounted datastores.
+// Non PersistentDatastores will not be accounted.
+func (d *Datastore) DiskUsage() (uint64, error) {
+	var duTotal uint64 = 0
+	for _, d := range d.mounts {
+		du, err := ds.DiskUsage(d.Datastore)
+		duTotal += du
+		if err != nil {
+			return duTotal, err
+		}
+	}
+	return duTotal, nil
+}
+
 type mountBatch struct {
 	mounts map[string]ds.Batch
 	lk     sync.Mutex

tiered/tiered.go

 package tiered
 
 import (
+	"errors"
 	"fmt"
+	"strings"
 	"sync"
 
 	ds "github.com/ipfs/go-datastore"
@@ -87,6 +89,37 @@ func (d tiered) Delete(key ds.Key) (err error) {
 	return nil
 }
 
+// DiskUsage returns the sum of DiskUsages from the datastores.
+// Non PersistentDatastores are not accounted.
+func (d tiered) DiskUsage() (uint64, error) {
+	var duTotal uint64 = 0
+	errs := []string{}
+	var lock sync.Mutex
+
+	var wg sync.WaitGroup
+	for _, cd := range d {
+		wg.Add(1)
+		go func(cd ds.Datastore) {
+			defer wg.Done()
+			du, err := ds.DiskUsage(cd)
+			lock.Lock()
+			defer lock.Unlock()
+
+			if err != nil {
+				errs = append(errs, err.Error())
+			}
+			duTotal += du
+		}(cd)
+	}
+	wg.Wait()
+
+	if len(errs) > 0 {
+		return duTotal, errors.New(strings.Join(errs, "\n"))
+	}
+
+	return duTotal, nil
+}
+
 // Query returns a list of keys in the datastore
 func (d tiered) Query(q dsq.Query) (dsq.Results, error) {
 	// query always the last (most complete) one

tiered/tiered_test.go

@@ -64,6 +64,19 @@ func TestTiered(t *testing.T) {
 	testHas(t, td, ds.NewKey("foo"), "bar2")
 }
 
+func TestTieredDiskUsage(t *testing.T) {
+	d1 := ds.NewMapDatastore()
+	d2 := ds.NewMapDatastore()
+	d3 := ds.NewMapDatastore()
+	d4 := ds.NewMapDatastore()
+
+	td := New(d1, d2, d3, d4)
+	td.Put(ds.NewKey("foo"), "bar")
+	if du, err := td.DiskUsage(); du != 0 || err != nil {
+		t.Error("tiered datastore size should be 0")
+	}
+}
+
 func TestQueryCallsLast(t *testing.T) {
 	var d1n, d2n, d3n int
 	d1 := dscb.Wrap(ds.NewMapDatastore(), func() { d1n++ })

timecache/timecache.go

@@ -102,3 +102,8 @@ func (d *datastore) Close() error {
 	}
 	return nil
 }
+
+// DiskUsage implements the PersistentDatastore interface.
+func (d *datastore) DiskUsage() (uint64, error) {
+	return ds.DiskUsage(d.cache)
+}