From d323c0f42b13a5dc47ce022488b8c3a010d6f2ff Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Tue, 24 Mar 2020 20:05:34 +0100
Subject: [PATCH] Add compactdb command

---
 README.md                  |  24 +++++
 cmd/chantools/compactdb.go | 174 +++++++++++++++++++++++++++++++++++++
 cmd/chantools/main.go      |   8 +-
 go.mod                     |   1 +
 4 files changed, 206 insertions(+), 1 deletion(-)
 create mode 100644 cmd/chantools/compactdb.go

diff --git a/README.md b/README.md
index 14ce91c..fd4b20b 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,7 @@
 * [Overview](#overview)
 * [Commands](#commands)
   + [chanbackup](#chanbackup)
+  + [compactdb](#compactdb)
   + [derivekey](#derivekey)
   + [dumpbackup](#dumpbackup)
   + [dumpchannels](#dumpchannels)
@@ -59,6 +60,7 @@ Help Options:
 
 Available commands:
   chanbackup       Create a channel.backup file from a channel database.
+  compactdb        Open a source channel.db database file in safe/read-only mode and copy it to a fresh database, compacting it in the process.
   derivekey        Derive a key with a specific derivation path from the BIP32 HD root key.
   dumpbackup       Dump the content of a channel.backup file.
   dumpchannels     Dump all channel information from lnd's channel database.
@@ -97,6 +99,28 @@ chantools chanbackup --rootkey xprvxxxxxxxxxx \
   --multi_file new_channel_backup.backup 
 ```
 
+### compactdb
+
+```text
+Usage:
+  chantools [OPTIONS] compactdb [compactdb-OPTIONS]
+
+[compactdb command options]
+          --txmaxsize=   Maximum transaction size. (default 65536)
+          --sourcedb=    The lnd channel.db file to create the database backup from.
+          --destdb=      The lnd new channel.db file to copy the compacted database to.
+```
+
+This command opens a database in read-only mode and tries to create a copy of it
+to a destination file, compacting it in the process.
+
+Example command:
+
+```bash
+chantools compactdb --sourcedb ~/.lnd/data/graph/mainnet/channel.db \
+  --destdb ./results/compacted.db
+```
+
 ### derivekey
 
 ```text
diff --git a/cmd/chantools/compactdb.go b/cmd/chantools/compactdb.go
new file mode 100644
index 0000000..fffa435
--- /dev/null
+++ b/cmd/chantools/compactdb.go
@@ -0,0 +1,174 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/coreos/bbolt"
+)
+
+const (
+	dbFilePermission = 0600
+	defaultTxMaxSize = 65536
+)
+
+type compactDBCommand struct {
+	TxMaxSize int64  `long:"txmaxsize" description:"Maximum transaction size. (default 65536)"`
+	SourceDB  string `long:"sourcedb" description:"The lnd channel.db file to create the database backup from."`
+	DestDB    string `long:"destdb" description:"The lnd new channel.db file to copy the compacted database to."`
+}
+
+func (c *compactDBCommand) Execute(_ []string) error {
+	// Check that we have a source and destination channel DB.
+	if c.SourceDB == "" {
+		return fmt.Errorf("source channel DB is required")
+	}
+	if c.DestDB == "" {
+		return fmt.Errorf("destination channel DB is required")
+	}
+	if c.TxMaxSize <= 0 {
+		c.TxMaxSize = defaultTxMaxSize
+	}
+	src, err := c.openDB(c.SourceDB, true)
+	if err != nil {
+		return fmt.Errorf("error opening source DB: %v", err)
+	}
+	dst, err := c.openDB(c.DestDB, false)
+	if err != nil {
+		return fmt.Errorf("error opening destination DB: %v", err)
+	}
+	err = c.compact(dst, src)
+	if err != nil {
+		return fmt.Errorf("error compacting DB: %v", err)
+	}
+	return nil
+}
+
+func (c *compactDBCommand) openDB(path string, ro bool) (*bbolt.DB, error) {
+	options := &bbolt.Options{
+		NoFreelistSync: false,
+		FreelistType:   bbolt.FreelistMapType,
+		ReadOnly:       ro,
+	}
+
+	bdb, err := bbolt.Open(path, dbFilePermission, options)
+	if err != nil {
+		return nil, err
+	}
+	return bdb, nil
+}
+
+func (c *compactDBCommand) compact(dst, src *bbolt.DB) error {
+	// commit regularly, or we'll run out of memory for large datasets if
+	// using one transaction.
+	var size int64
+	tx, err := dst.Begin(true)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	if err := c.walk(src, func(keys [][]byte, k, v []byte, seq uint64) error {
+		// On each key/value, check if we have exceeded tx size.
+		sz := int64(len(k) + len(v))
+		if size+sz > c.TxMaxSize && c.TxMaxSize != 0 {
+			// Commit previous transaction.
+			if err := tx.Commit(); err != nil {
+				return err
+			}
+
+			// Start new transaction.
+			tx, err = dst.Begin(true)
+			if err != nil {
+				return err
+			}
+			size = 0
+		}
+		size += sz
+
+		// Create bucket on the root transaction if this is the first
+		// level.
+		nk := len(keys)
+		if nk == 0 {
+			bkt, err := tx.CreateBucket(k)
+			if err != nil {
+				return err
+			}
+			if err := bkt.SetSequence(seq); err != nil {
+				return err
+			}
+			return nil
+		}
+
+		// Create buckets on subsequent levels, if necessary.
+		b := tx.Bucket(keys[0])
+		if nk > 1 {
+			for _, k := range keys[1:] {
+				b = b.Bucket(k)
+			}
+		}
+
+		// Fill the entire page for best compaction.
+		b.FillPercent = 1.0
+
+		// If there is no value then this is a bucket call.
+		if v == nil {
+			bkt, err := b.CreateBucket(k)
+			if err != nil {
+				return err
+			}
+			if err := bkt.SetSequence(seq); err != nil {
+				return err
+			}
+			return nil
+		}
+
+		// Otherwise treat it as a key/value pair.
+		return b.Put(k, v)
+	}); err != nil {
+		return err
+	}
+
+	return tx.Commit()
+}
+
+// walkFunc is the type of the function called for keys (buckets and "normal"
+// values) discovered by Walk. keys is the list of keys to descend to the bucket
+// owning the discovered key/value pair k/v.
+type walkFunc func(keys [][]byte, k, v []byte, seq uint64) error
+
+// walk walks recursively the bolt database db, calling walkFn for each key it
+// finds.
+func (c *compactDBCommand) walk(db *bbolt.DB, walkFn walkFunc) error {
+	return db.View(func(tx *bbolt.Tx) error {
+		return tx.ForEach(func(name []byte, b *bbolt.Bucket) error {
+			return c.walkBucket(
+				b, nil, name, nil, b.Sequence(), walkFn,
+			)
+		})
+	})
+}
+
+func (c *compactDBCommand) walkBucket(b *bbolt.Bucket, keypath [][]byte,
+	k, v []byte, seq uint64, fn walkFunc) error {
+	// Execute callback.
+	if err := fn(keypath, k, v, seq); err != nil {
+		return err
+	}
+
+	// If this is not a bucket then stop.
+	if v != nil {
+		return nil
+	}
+
+	// Iterate over each child key/value.
+	keypath = append(keypath, k)
+	return b.ForEach(func(k, v []byte) error {
+		if v == nil {
+			bkt := b.Bucket(k)
+			return c.walkBucket(
+				bkt, keypath, k, nil, bkt.Sequence(), fn,
+			)
+		}
+		return c.walkBucket(b, keypath, k, v, b.Sequence(), fn)
+	})
+}
diff --git a/cmd/chantools/main.go b/cmd/chantools/main.go
index 30b9c3c..04f19d2 100644
--- a/cmd/chantools/main.go
+++ b/cmd/chantools/main.go
@@ -5,7 +5,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"github.com/lightningnetwork/lnd/chanbackup"
 	"io/ioutil"
 	"os"
 	"path"
@@ -20,6 +19,7 @@ import (
 	"github.com/jessevdk/go-flags"
 	"github.com/lightningnetwork/lnd/aezeed"
 	"github.com/lightningnetwork/lnd/build"
+	"github.com/lightningnetwork/lnd/chanbackup"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"golang.org/x/crypto/ssh/terminal"
 )
@@ -122,6 +122,12 @@ func runCommandParser() error {
 		"chanbackup", "Create a channel.backup file from a channel "+
 			"database.", "", &chanBackupCommand{},
 	)
+	_, _ = parser.AddCommand(
+		"compactdb", "Open a source channel.db database file in safe/"+
+			"read-only mode and copy it to a fresh database, "+
+			"compacting it in the process.", "",
+		&compactDBCommand{},
+	)
 
 	_, err := parser.Parse()
 	return err
diff --git a/go.mod b/go.mod
index c618951..0d7099c 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/btcsuite/btcutil v0.0.0-20191219182022-e17c9730c422
 	github.com/btcsuite/btcwallet v0.11.1-0.20200219004649-ae9416ad7623
 	github.com/btcsuite/btcwallet/walletdb v1.2.0
+	github.com/coreos/bbolt v1.3.3
 	github.com/davecgh/go-spew v1.1.1
 	github.com/golang/protobuf v1.3.2 // indirect
 	github.com/jessevdk/go-flags v1.4.0