From 9827df70e040f6ee5fa9b615d5796d0af5c55cbc Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Fri, 21 Feb 2014 09:49:15 -0700
Subject: [PATCH] Add DB.Stat().

---
 db.go      | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++
 db_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/db.go b/db.go
index b7338c3..c9b0e92 100644
--- a/db.go
+++ b/db.go
@@ -506,6 +506,34 @@ func (db *DB) CopyFile(path string, mode os.FileMode) error {
 	return db.Copy(f)
 }
 
+// Stat retrieves stats on the database and its page usage.
+// Returns an error if the database is not open.
+func (db *DB) Stat() (*Stat, error) {
+	// Obtain meta & mmap locks.
+	db.metalock.Lock()
+	db.mmaplock.RLock()
+
+	var s = &Stat{
+		MmapSize:         len(db.data),
+		TransactionCount: len(db.transactions),
+	}
+
+	// Release locks.
+	db.mmaplock.RUnlock()
+	db.metalock.Unlock()
+
+	err := db.Do(func(t *RWTransaction) error {
+		s.PageCount = int(t.meta.pgid)
+		s.FreePageCount = len(db.freelist.all())
+		s.PageSize = db.pageSize
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return s, nil
+}
+
 // page retrieves a page reference from the mmap based on the current page size.
 func (db *DB) page(id pgid) *page {
 	return (*page)(unsafe.Pointer(&db.data[id*pgid(db.pageSize)]))
@@ -550,3 +578,26 @@ func (db *DB) allocate(count int) (*page, error) {
 
 	return p, nil
 }
+
+// Stat represents stats on the database such as free pages and sizes.
+type Stat struct {
+	// PageCount is the total number of allocated pages. This is a high water
+	// mark in the database that represents how many pages have actually been
+	// used. This will be smaller than the MmapSize / PageSize.
+	PageCount int
+
+	// FreePageCount is the total number of pages which have been previously
+	// allocated but are no longer used.
+	FreePageCount int
+
+	// PageSize is the size, in bytes, of individual database pages.
+	PageSize int
+
+	// MmapSize is the mmap-allocated size of the data file. When the data file
+	// grows beyond this size, the database will obtain a lock on the mmap and
+	// resize it.
+	MmapSize int
+
+	// TransactionCount is the total number of reader transactions.
+	TransactionCount int
+}
diff --git a/db_test.go b/db_test.go
index ecf2290..3d0b987 100644
--- a/db_test.go
+++ b/db_test.go
@@ -5,6 +5,7 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"strconv"
 	"syscall"
 	"testing"
 	"time"
@@ -329,6 +330,51 @@ func TestDBCopyFile(t *testing.T) {
 	})
 }
 
+// Ensure the database can return stats about itself.
+func TestDBStat(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		db.Do(func(txn *RWTransaction) error {
+			txn.CreateBucket("widgets")
+			for i := 0; i < 10000; i++ {
+				txn.Put("widgets", []byte(strconv.Itoa(i)), []byte(strconv.Itoa(i)))
+			}
+			return nil
+		})
+
+		// Delete some keys.
+		db.Delete("widgets", []byte("10"))
+		db.Delete("widgets", []byte("1000"))
+
+		// Open some readers.
+		t0, _ := db.Transaction()
+		t1, _ := db.Transaction()
+		t2, _ := db.Transaction()
+		t2.Close()
+
+		// Obtain stats.
+		stat, err := db.Stat()
+		assert.NoError(t, err)
+		assert.Equal(t, stat.PageCount, 128)
+		assert.Equal(t, stat.FreePageCount, 2)
+		assert.Equal(t, stat.PageSize, 4096)
+		assert.Equal(t, stat.MmapSize, 4194304)
+		assert.Equal(t, stat.TransactionCount, 2)
+
+		// Close readers.
+		t0.Close()
+		t1.Close()
+	})
+}
+
+// Ensure the getting stats on a closed database returns an error.
+func TestDBStatWhileClosed(t *testing.T) {
+	withDB(func(db *DB, path string) {
+		stat, err := db.Stat()
+		assert.Equal(t, err, ErrDatabaseNotOpen)
+		assert.Nil(t, stat)
+	})
+}
+
 // Ensure that an error is returned when a database write fails.
 func TestDBWriteFail(t *testing.T) {
 	t.Skip("pending") // TODO(benbjohnson)