Merge pull request #4440 from carlaKC/4165-ratelimit

chanfitness: Rate limit in memory events based on peer flap rate

channeldb/db.go

@@ -164,6 +164,12 @@ var (
 			number:    17,
 			migration: mig.CreateTLB(closeSummaryBucket),
 		},
+		{
+			// Create a top level bucket which holds information
+			// about our peers.
+			number:    18,
+			migration: mig.CreateTLB(peersBucket),
+		},
 	}
 
 	// Big endian is the preferred byte order, due to cursor scans over
@@ -278,6 +284,7 @@ var topLevelBuckets = [][]byte{
 	invoiceBucket,
 	payAddrIndexBucket,
 	paymentsIndexBucket,
+	peersBucket,
 	nodeInfoBucket,
 	nodeBucket,
 	edgeBucket,

channeldb/peers.go

@@ -0,0 +1,121 @@
+package channeldb
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/lightningnetwork/lnd/routing/route"
+)
+
+var (
+	// peersBucket is the name of a top level bucket in which we store
+	// information about our peers. Information for different peers is
+	// stored in buckets keyed by their public key.
+	//
+	//
+	// peers-bucket
+	//      |
+	//      |-- <peer-pubkey>
+	//      |        |--flap-count-key: <ts><flap count>
+	//      |
+	//      |-- <peer-pubkey>
+	//      |        |--flap-count-key: <ts><flap count>
+	peersBucket = []byte("peers-bucket")
+
+	// flapCountKey is a key used in the peer pubkey sub-bucket that stores
+	// the timestamp of a peer's last flap count and its all time flap
+	// count.
+	flapCountKey = []byte("flap-count")
+)
+
+var (
+	// ErrNoPeerBucket is returned when we try to read entries for a peer
+	// that is not tracked.
+	ErrNoPeerBucket = errors.New("peer bucket not found")
+)
+
+// FlapCount contains information about a peer's flap count.
+type FlapCount struct {
+	// Count provides the total flap count for a peer.
+	Count uint32
+
+	// LastFlap is the timestamp of the last flap recorded for a peer.
+	LastFlap time.Time
+}
+
+// WriteFlapCounts writes the flap count for a set of peers to disk, creating a
+// bucket for the peer's pubkey if necessary. Note that this function overwrites
+// the current value.
+func (d *DB) WriteFlapCounts(flapCounts map[route.Vertex]*FlapCount) error {
+	return d.Update(func(tx walletdb.ReadWriteTx) error {
+		// Run through our set of flap counts and record them for
+		// each peer, creating a bucket for the peer pubkey if required.
+		for peer, flapCount := range flapCounts {
+			peers := tx.ReadWriteBucket(peersBucket)
+
+			peerBucket, err := peers.CreateBucketIfNotExists(
+				peer[:],
+			)
+			if err != nil {
+				return err
+			}
+
+			var b bytes.Buffer
+			err = serializeTime(&b, flapCount.LastFlap)
+			if err != nil {
+				return err
+			}
+
+			if err = WriteElement(&b, flapCount.Count); err != nil {
+				return err
+			}
+
+			err = peerBucket.Put(flapCountKey, b.Bytes())
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
+	})
+}
+
+// ReadFlapCount attempts to read the flap count for a peer, failing if the
+// peer is not found or we do not have flap count stored.
+func (d *DB) ReadFlapCount(pubkey route.Vertex) (*FlapCount, error) {
+	var flapCount FlapCount
+
+	if err := d.View(func(tx walletdb.ReadTx) error {
+		peers := tx.ReadBucket(peersBucket)
+
+		peerBucket := peers.NestedReadBucket(pubkey[:])
+		if peerBucket == nil {
+			return ErrNoPeerBucket
+		}
+
+		flapBytes := peerBucket.Get(flapCountKey)
+		if flapBytes == nil {
+			return fmt.Errorf("flap count not recorded for: %v",
+				pubkey)
+		}
+
+		var (
+			err error
+			r   = bytes.NewReader(flapBytes)
+		)
+
+		flapCount.LastFlap, err = deserializeTime(r)
+		if err != nil {
+			return err
+		}
+
+		return ReadElements(r, &flapCount.Count)
+	}); err != nil {
+		return nil, err
+	}
+
+	return &flapCount, nil
+}

channeldb/peers_test.go

@@ -0,0 +1,50 @@
+package channeldb
+
+import (
+	"testing"
+	"time"
+
+	"github.com/lightningnetwork/lnd/routing/route"
+	"github.com/stretchr/testify/require"
+)
+
+// TestFlapCount tests lookup and writing of flap count to disk.
+func TestFlapCount(t *testing.T) {
+	db, cleanup, err := MakeTestDB()
+	require.NoError(t, err)
+	defer cleanup()
+
+	// Try to read flap count for a peer that we have no records for.
+	_, err = db.ReadFlapCount(testPub)
+	require.Equal(t, ErrNoPeerBucket, err)
+
+	var (
+		testPub2       = route.Vertex{2, 2, 2}
+		peer1FlapCount = &FlapCount{
+			Count:    20,
+			LastFlap: time.Unix(100, 23),
+		}
+		peer2FlapCount = &FlapCount{
+			Count:    39,
+			LastFlap: time.Unix(200, 23),
+		}
+	)
+
+	peers := map[route.Vertex]*FlapCount{
+		testPub:  peer1FlapCount,
+		testPub2: peer2FlapCount,
+	}
+
+	err = db.WriteFlapCounts(peers)
+	require.NoError(t, err)
+
+	// Lookup flap count for our first pubkey.
+	count, err := db.ReadFlapCount(testPub)
+	require.NoError(t, err)
+	require.Equal(t, peer1FlapCount, count)
+
+	// Lookup our flap count for the second peer.
+	count, err = db.ReadFlapCount(testPub2)
+	require.NoError(t, err)
+	require.Equal(t, peer2FlapCount, count)
+}