Batch insert observed messages using Duckdb appender (#998)

* Batch insert observed messages using Duckdb appender

For fast batch insertion of messages observed into duckdb, use the
native appender to batch the messages based on size or max wait time.
Whichever happens first.

Fixes #993

* Periodically attempt to flush in-flight messages out to the database

This is to ensure max batch delay is respected independent of message
arrival/discovery flow.

* Fix usage of timer instead of ticker for both rotation and flush

Author: masih

cmd/f3/observer.go

@@ -105,6 +105,16 @@ var observerCmd = cli.Command{
 			Usage: "The connectivity threshold below which peer discovery via Lotus Net Peers is engaged. Disabled if set to zero or no lotusDaemon endpoints are provided.",
 			Value: 100,
 		},
+		&cli.IntFlag{
+			Name:  "maxBatchSize",
+			Usage: "The maximum number of messages to batch together in a single insertion into database.",
+			Value: 1000,
+		},
+		&cli.DurationFlag{
+			Name:  "maxBatchDelay",
+			Usage: "The maximum time to wait before a batch is flushed to the database.",
+			Value: time.Minute,
+		},
 	},
 
 	Action: func(cctx *cli.Context) error {
@@ -116,6 +126,8 @@ var observerCmd = cli.Command{
 			observer.WithRetention(cctx.Duration("retention")),
 			observer.WithDataSourceName(cctx.String("dataSourceName")),
 			observer.WithMaxConcurrentConnectionAttempts(cctx.Int("reconnectConcurrency")),
+			observer.WithMaxBatchSize(cctx.Int("maxBatchSize")),
+			observer.WithMaxBatchDelay(cctx.Duration("maxBatchDelay")),
 		}
 		var identity crypto.PrivKey
 		if cctx.IsSet("identity") {

observer/observer.go

@@ -3,8 +3,8 @@ package observer
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	_ "embed"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -24,7 +24,7 @@ import (
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
 	record "github.com/libp2p/go-libp2p-record"
 	"github.com/libp2p/go-libp2p/core/peer"
-	_ "github.com/marcboeker/go-duckdb"
+	"github.com/marcboeker/go-duckdb"
 	"go.uber.org/multierr"
 	"golang.org/x/sync/errgroup"
 )
@@ -45,6 +45,12 @@ type Observer struct {
 
 	messageObserved chan *Message
 	msgEncoding     *encoding.ZSTD[*pmsg.PartialGMessage]
+
+	dbConnector           *duckdb.Connector
+	dbAppender            *duckdb.Appender
+	dbConnection          driver.Conn
+	unflushedMessageCount int
+	lastFlushedAt         time.Time
 }
 
 func New(o ...Option) (*Observer, error) {
@@ -75,14 +81,18 @@ func (o *Observer) Start(ctx context.Context) error {
 	eg.Go(o.listenAndServeQueries)
 	o.stop = func() error {
 		stop()
-		return eg.Wait()
+		return multierr.Combine(
+			o.dbAppender.Close(),
+			o.dbConnection.Close(),
+			o.dbConnector.Close(),
+			o.db.Close(),
+			eg.Wait())
 	}
 	return nil
 }
 
 func (o *Observer) initialize(ctx context.Context) error {
 	var err error
-
 	if o.pubSub == nil {
 		// Set up pubsub to listen for GPBFT messages.
 		if o.pubSub, err = pubsub.NewGossipSub(ctx, o.host,
@@ -115,9 +125,11 @@ func (o *Observer) initialize(ctx context.Context) error {
 	}
 
 	// Set up database connection.
-	if o.db, err = sql.Open("duckdb", o.dataSourceName); err != nil {
-		return err
+	o.dbConnector, err = duckdb.NewConnector(o.dataSourceName, nil)
+	if err != nil {
+		return fmt.Errorf("failed to create duckdb connector: %w", err)
 	}
+	o.db = sql.OpenDB(o.dbConnector)
 
 	// Create database schema.
 	if _, err := o.db.ExecContext(ctx, schema); err != nil {
@@ -138,6 +150,16 @@ func (o *Observer) initialize(ctx context.Context) error {
 		return err
 	}
 
+	// Set up appender used for batch insertion of messages observed.
+	o.dbConnection, err = o.dbConnector.Connect(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to connect to duckdb: %w", err)
+	}
+	o.dbAppender, err = duckdb.NewAppenderFromConn(o.dbConnection, "", "latest_messages")
+	if err != nil {
+		return fmt.Errorf("failed to create duckdb appender: %w", err)
+	}
+
 	// If connectivity check interval is enabled, repair connections once to avoid
 	// waiting for the ticker.
 	if o.connectivityCheckInterval > 0 {
@@ -169,14 +191,16 @@ func (o *Observer) createOrReplaceMessagesView(ctx context.Context, includeParqu
 }
 
 func (o *Observer) observe(ctx context.Context) error {
-	rotation := time.NewTimer(o.rotateInterval)
+	rotation := time.NewTicker(o.rotateInterval)
+	flush := time.NewTicker(o.maxBatchDelay)
 	stopObserverForNetwork, err := o.startObserverFor(ctx, o.networkName)
 	if err != nil {
 		return fmt.Errorf("failed to start observer for network %s: %w", o.networkName, err)
 	}
 
 	defer func() {
 		rotation.Stop()
+		flush.Stop()
 		stopObserverForNetwork()
 	}()
 
@@ -194,36 +218,49 @@ func (o *Observer) observe(ctx context.Context) error {
 			if err := o.rotateMessages(ctx); err != nil {
 				logger.Errorw("Failed to rotate latest messages", "err", err)
 			}
+		case <-flush.C:
+			if err := o.tryFlushMessages(); err != nil {
+				logger.Errorw("Failed to flush observed messages", "err", err)
+			}
 		}
 	}
 	return nil
 }
 
-func (o *Observer) storeMessage(ctx context.Context, om *Message) error {
-	const insertMessage = `INSERT INTO latest_messages VALUES(?,?,?,?::json,?,?,?::json,?);`
-	voteMarshaled, err := json.Marshal(om.Vote)
-	if err != nil {
-		return fmt.Errorf("failed to marshal vote: %w", err)
-	}
-	var justificationMarshaled any
+func (o *Observer) storeMessage(_ context.Context, om *Message) error {
+	var justification any
 	if om.Justification != nil {
-		v, err := json.Marshal(om.Justification)
-		if err != nil {
-			return fmt.Errorf("failed to marshal justification: %w", err)
-		}
-		justificationMarshaled = string(v)
+		// Dereference to get the go-duckdb reflection in appender to behave. Otherwise,
+		// it fails as it interprets the type as mismatch due to how nullable fields are
+		// handled.
+		justification = *om.Justification
 	}
-	if _, err := o.db.ExecContext(ctx, insertMessage,
+	if err := o.dbAppender.AppendRow(
 		om.Timestamp,
 		om.NetworkName,
-		om.Sender,
-		string(voteMarshaled),
+		int64(om.Sender),
+		om.Vote,
 		om.Signature,
 		om.Ticket,
-		justificationMarshaled,
+		justification,
 		om.VoteValueKey,
 	); err != nil {
-		return fmt.Errorf("failed to execute query: %w", err)
+		return fmt.Errorf("failed to append row: %w", err)
+	}
+	o.unflushedMessageCount++
+	return o.tryFlushMessages()
+}
+
+func (o *Observer) tryFlushMessages() error {
+	batchSizeReached := o.unflushedMessageCount >= o.maxBatchSize
+	batchDelayElapsed := !o.lastFlushedAt.IsZero() && time.Since(o.lastFlushedAt) >= o.maxBatchDelay
+	if batchSizeReached || batchDelayElapsed {
+		if err := o.dbAppender.Flush(); err != nil {
+			return fmt.Errorf("failed to flush appender: %w", err)
+		}
+		logger.Infow("Flushed messages to database", "count", o.unflushedMessageCount, "after", time.Since(o.lastFlushedAt))
+		o.unflushedMessageCount = 0
+		o.lastFlushedAt = time.Now()
 	}
 	return nil
 }

observer/options.go

@@ -41,6 +41,9 @@ type options struct {
 	pubSubValidatorDisabled bool
 
 	dataSourceName string
+
+	maxBatchSize  int
+	maxBatchDelay time.Duration
 }
 
 func newOptions(opts ...Option) (*options, error) {
@@ -54,6 +57,8 @@ func newOptions(opts ...Option) (*options, error) {
 		rotatePath:                ".",
 		rotateInterval:            10 * time.Minute,
 		retention:                 -1,
+		maxBatchSize:              1000,
+		maxBatchDelay:             time.Minute,
 	}
 	for _, apply := range opts {
 		if err := apply(&opt); err != nil {
@@ -256,3 +261,20 @@ func WithPubSub(ps *pubsub.PubSub) Option {
 		return nil
 	}
 }
+
+func WithMaxBatchSize(size int) Option {
+	return func(o *options) error {
+		o.maxBatchSize = size
+		return nil
+	}
+}
+
+func WithMaxBatchDelay(d time.Duration) Option {
+	return func(o *options) error {
+		if d < 0 {
+			return fmt.Errorf("max batch delay must be greater than or equal to 0")
+		}
+		o.maxBatchDelay = d
+		return nil
+	}
+}