Add an optimized QueueLeaves implementation for single-leaf batches in the PostgreSQL storage backend (#3769)

* Add an optimized QueueLeaves implementation for single-leaf batches in the PostgreSQL storage backend

* Add tests

* Update CHANGELOG.md

* Add missing CHANGELOG.md entry for PR #3752

Author: robstradling

CHANGELOG.md

@@ -11,6 +11,8 @@
 ### Storage
 
 * For PostgreSQL, explicitly create index on SequencedLeafData(TreeId, LeafIdentityHash) by @robstradling in https://github.com/google/trillian/pull/3695
+* Fix error checking for QueueLeaves and AddSequencedLeaves in the PostgreSQL storage backend by @robstradling in https://github.com/google/trillian/pull/3752
+* Add an optimized QueueLeaves implementation for single-leaf batches in the PostgreSQL storage backend by @robstradling in https://github.com/google/trillian/pull/3769
 * Improve PostgreSQL functions by @robstradling in https://github.com/google/trillian/pull/3770
 
 ### Misc

storage/postgresql/log_storage.go

@@ -42,6 +42,15 @@ import (
 )
 
 const (
+	queueLeafSQL = "WITH insert_leaf AS (" +
+		"INSERT INTO LeafData (TreeId,LeafIdentityHash,LeafValue,ExtraData,QueueTimestampNanos) " +
+		"VALUES ($1,$2,$3,$4,$5) " +
+		"ON CONFLICT DO NOTHING " +
+		"RETURNING *" +
+		") " +
+		"INSERT INTO Unsequenced (TreeId,Bucket,LeafIdentityHash,MerkleLeafHash,QueueTimestampNanos,QueueID) " +
+		"SELECT TreeId,0,LeafIdentityHash,$6,QueueTimestampNanos,$7 " +
+		"FROM insert_leaf"
 	createTempQueueLeavesTable = "CREATE TEMP TABLE TempQueueLeaves (" +
 		" TreeId BIGINT," +
 		" LeafIdentityHash BYTEA," +
@@ -295,7 +304,14 @@ func (m *postgreSQLLogStorage) QueueLeaves(ctx context.Context, tree *trillian.T
 	if err != nil {
 		return nil, err
 	}
-	existing, err := tx.QueueLeaves(ctx, leaves, queueTimestamp)
+
+	// Queue leave(s), using a more efficient implementation when the batch size is 1.
+	var existing []*trillian.LogLeaf
+	if len(leaves) == 1 {
+		existing, err = tx.QueueLeaf(ctx, leaves[0], queueTimestamp)
+	} else {
+		existing, err = tx.QueueLeaves(ctx, leaves, queueTimestamp)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -389,6 +405,55 @@ func (t *logTreeTX) DequeueLeaves(ctx context.Context, limit int, cutoffTime tim
 	return leaves, nil
 }
 
+func (t *logTreeTX) QueueLeaf(ctx context.Context, leaf *trillian.LogLeaf, queueTimestamp time.Time) ([]*trillian.LogLeaf, error) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	// Prepare details to store, but don't accept leaf if invalid.
+	if len(leaf.LeafIdentityHash) != t.hashSizeBytes {
+		return nil, fmt.Errorf("queued leaf must have a leaf ID hash of length %d", t.hashSizeBytes)
+	}
+	leaf.QueueTimestamp = timestamppb.New(queueTimestamp)
+	if err := leaf.QueueTimestamp.CheckValid(); err != nil {
+		return nil, fmt.Errorf("got invalid queue timestamp: %w", err)
+	}
+	qTimestamp := leaf.QueueTimestamp.AsTime()
+	args := queueArgs(t.treeID, leaf.LeafIdentityHash, qTimestamp)
+	label := labelForTX(t)
+
+	// Create the leaf data record and work queue entry, unless the leaf already exists.
+	existingLeaves := make([]*trillian.LogLeaf, 1)
+	result, err := t.tx.Exec(ctx, queueLeafSQL, t.treeID, leaf.LeafIdentityHash, leaf.LeafValue, leaf.ExtraData, args[0], leaf.MerkleLeafHash, args[1])
+	if err != nil {
+		klog.Warningf("Failed to queue leaf: %s", err)
+		return nil, postgresqlToGRPC(err)
+	}
+	queuedCounter.Add(1, label)
+	if result.RowsAffected() > 0 {
+		// Leaf did not already exist.
+		return existingLeaves, nil
+	}
+
+	var toRetrieve [][]byte
+	toRetrieve = append(toRetrieve, leaf.LeafIdentityHash)
+	queuedDupCounter.Inc(label)
+	results, err := t.getLeafDataByIdentityHash(ctx, toRetrieve)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve existing leaf: %v", err)
+	}
+	if len(results) != len(toRetrieve) {
+		return nil, fmt.Errorf("failed to retrieve existing leaf: got %d, want %d", len(results), len(toRetrieve))
+	}
+	// Replace the requested leaf with the actual leaf.
+	if bytes.Equal(results[0].LeafIdentityHash, leaf.LeafIdentityHash) {
+		existingLeaves[0] = results[0]
+	} else {
+		return nil, fmt.Errorf("failed to find existing leaf for hash %x", leaf.LeafIdentityHash)
+	}
+
+	return existingLeaves, nil
+}
+
 func (t *logTreeTX) QueueLeaves(ctx context.Context, leaves []*trillian.LogLeaf, queueTimestamp time.Time) ([]*trillian.LogLeaf, error) {
 	t.mu.Lock()
 	defer t.mu.Unlock()

storage/postgresql/log_storage_test.go

@@ -137,6 +137,7 @@ func TestQueueDuplicateLeaf(t *testing.T) {
 	leaves2 := createTestLeaves(int64(count), 12)
 	leaves3 := createTestLeaves(3, 100)
 	leaves4 := createTestLeaves(3, 105)
+	leaves5 := createTestLeaves(1, 110)
 
 	// Note that tests accumulate queued leaves on top of each other.
 	tests := []struct {
@@ -165,6 +166,18 @@ func TestQueueDuplicateLeaf(t *testing.T) {
 			leaves: []*trillian.LogLeaf{leaves3[0], leaves3[0], leaves4[1], leaves3[1], leaves4[2]},
 			want:   []*trillian.LogLeaf{leaves3[0], leaves3[0], leaves4[1], leaves3[1], leaves4[2]},
 		},
+		{
+			// single leaf (dup)
+			desc:   "[100]",
+			leaves: []*trillian.LogLeaf{leaves3[0]},
+			want:   []*trillian.LogLeaf{leaves3[0]},
+		},
+		{
+			// single leaf (new)
+			desc:   "[110]",
+			leaves: []*trillian.LogLeaf{leaves5[0]},
+			want:   []*trillian.LogLeaf{leaves5[0]},
+		},
 	}
 
 	for _, test := range tests {