channeldb: add balance at height lookup

Add a balance at height lookup function which can be used to
obtain local/remote balance at a given height. The current in memory
commits and revocation log are used to source this information.

channeldb/channel_test.go

@@ -10,6 +10,8 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/lightningnetwork/lnd/channeldb/kvdb"
+
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
@@ -131,6 +133,25 @@ type testChannelParams struct {
 // default channel that is creates for testing.
 type testChannelOption func(params *testChannelParams)
 
+// channelCommitmentOption is an option which allows overwriting of the default
+// commitment height and balances. The local boolean can be used to set these
+// balances on the local or remote commit.
+func channelCommitmentOption(height uint64, localBalance,
+	remoteBalance lnwire.MilliSatoshi, local bool) testChannelOption {
+
+	return func(params *testChannelParams) {
+		if local {
+			params.channel.LocalCommitment.CommitHeight = height
+			params.channel.LocalCommitment.LocalBalance = localBalance
+			params.channel.LocalCommitment.RemoteBalance = remoteBalance
+		} else {
+			params.channel.RemoteCommitment.CommitHeight = height
+			params.channel.RemoteCommitment.LocalBalance = localBalance
+			params.channel.RemoteCommitment.RemoteBalance = remoteBalance
+		}
+	}
+}
+
 // pendingHeightOption is an option which can be used to set the height the
 // channel is marked as pending at.
 func pendingHeightOption(height uint32) testChannelOption {
@@ -1393,3 +1414,169 @@ func TestCloseChannelStatus(t *testing.T) {
 		t.Fatalf("channel should have status")
 	}
 }
+
+// TestBalanceAtHeight tests lookup of our local and remote balance at a given
+// height.
+func TestBalanceAtHeight(t *testing.T) {
+	const (
+		// Values that will be set on our current local commit in
+		// memory.
+		localHeight        = 2
+		localLocalBalance  = 1000
+		localRemoteBalance = 1500
+
+		// Values that will be set on our current remote commit in
+		// memory.
+		remoteHeight        = 3
+		remoteLocalBalance  = 2000
+		remoteRemoteBalance = 2500
+
+		// Values that will be written to disk in the revocation log.
+		oldHeight        = 0
+		oldLocalBalance  = 200
+		oldRemoteBalance = 300
+
+		// Heights to test error cases.
+		unknownHeight   = 1
+		unreachedHeight = 4
+	)
+
+	// putRevokedState is a helper function used to put commitments is
+	// the revocation log bucket to test lookup of balances at heights that
+	// are not our current height.
+	putRevokedState := func(c *OpenChannel, height uint64, local,
+		remote lnwire.MilliSatoshi) error {
+
+		err := kvdb.Update(c.Db, func(tx kvdb.RwTx) error {
+			chanBucket, err := fetchChanBucketRw(
+				tx, c.IdentityPub, &c.FundingOutpoint,
+				c.ChainHash,
+			)
+			if err != nil {
+				return err
+			}
+
+			logKey := revocationLogBucket
+			logBucket, err := chanBucket.CreateBucketIfNotExists(
+				logKey,
+			)
+			if err != nil {
+				return err
+			}
+
+			// Make a copy of our current commitment so we do not
+			// need to re-fill all the required fields and copy in
+			// our new desired values.
+			commit := c.LocalCommitment
+			commit.CommitHeight = height
+			commit.LocalBalance = local
+			commit.RemoteBalance = remote
+
+			return appendChannelLogEntry(logBucket, &commit)
+		})
+
+		return err
+	}
+
+	tests := []struct {
+		name                  string
+		targetHeight          uint64
+		expectedLocalBalance  lnwire.MilliSatoshi
+		expectedRemoteBalance lnwire.MilliSatoshi
+		expectedError         error
+	}{
+		{
+			name:                  "target is current local height",
+			targetHeight:          localHeight,
+			expectedLocalBalance:  localLocalBalance,
+			expectedRemoteBalance: localRemoteBalance,
+			expectedError:         nil,
+		},
+		{
+			name:                  "target is current remote height",
+			targetHeight:          remoteHeight,
+			expectedLocalBalance:  remoteLocalBalance,
+			expectedRemoteBalance: remoteRemoteBalance,
+			expectedError:         nil,
+		},
+		{
+			name:                  "need to lookup commit",
+			targetHeight:          oldHeight,
+			expectedLocalBalance:  oldLocalBalance,
+			expectedRemoteBalance: oldRemoteBalance,
+			expectedError:         nil,
+		},
+		{
+			name:                  "height not found",
+			targetHeight:          unknownHeight,
+			expectedLocalBalance:  0,
+			expectedRemoteBalance: 0,
+			expectedError:         errLogEntryNotFound,
+		},
+		{
+			name:                  "height not reached",
+			targetHeight:          unreachedHeight,
+			expectedLocalBalance:  0,
+			expectedRemoteBalance: 0,
+			expectedError:         errHeightNotReached,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			cdb, cleanUp, err := makeTestDB()
+			if err != nil {
+				t.Fatalf("unable to make test database: %v",
+					err)
+			}
+			defer cleanUp()
+
+			// Create options to set the heights and balances of
+			// our local and remote commitments.
+			localCommitOpt := channelCommitmentOption(
+				localHeight, localLocalBalance,
+				localRemoteBalance, true,
+			)
+
+			remoteCommitOpt := channelCommitmentOption(
+				remoteHeight, remoteLocalBalance,
+				remoteRemoteBalance, false,
+			)
+
+			// Create an open channel.
+			channel := createTestChannel(
+				t, cdb, openChannelOption(),
+				localCommitOpt, remoteCommitOpt,
+			)
+
+			// Write an older commit to disk.
+			err = putRevokedState(channel, oldHeight,
+				oldLocalBalance, oldRemoteBalance)
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+
+			local, remote, err := channel.BalancesAtHeight(
+				test.targetHeight,
+			)
+			if err != test.expectedError {
+				t.Fatalf("expected: %v, got: %v",
+					test.expectedError, err)
+			}
+
+			if local != test.expectedLocalBalance {
+				t.Fatalf("expected local: %v, got: %v",
+					test.expectedLocalBalance, local)
+			}
+
+			if remote != test.expectedRemoteBalance {
+				t.Fatalf("expected remote: %v, got: %v",
+					test.expectedRemoteBalance, remote)
+			}
+		})
+	}
+}