optimize s3 insertion

Author: jakeloo

internal/storage/block_buffer.go

@@ -1,6 +1,8 @@
 package storage
 
 import (
+	"bytes"
+	"encoding/gob"
 	"fmt"
 	"math/big"
 	"sync"
@@ -20,7 +22,7 @@ type BlockBuffer struct {
 
 // IBlockBuffer defines the interface for block buffer implementations
 type IBlockBuffer interface {
-	Add(blocks []common.BlockData, actualSizeBytes int64) bool
+	Add(blocks []common.BlockData) bool
 	Flush() []common.BlockData
 	ShouldFlush() bool
 	Size() (int64, int)
@@ -31,6 +33,7 @@ type IBlockBuffer interface {
 	GetMaxBlockNumber(chainId *big.Int) *big.Int
 	Clear()
 	Stats() BufferStats
+	Close() error
 }
 
 // NewBlockBuffer creates a new in-memory block buffer
@@ -49,21 +52,35 @@ func NewBlockBufferWithBadger(maxSizeMB int64, maxBlocks int) (IBlockBuffer, err
 }
 
 // Add adds blocks to the buffer and returns true if flush is needed
-func (b *BlockBuffer) Add(blocks []common.BlockData, actualSizeBytes int64) bool {
+func (b *BlockBuffer) Add(blocks []common.BlockData) bool {
 	if len(blocks) == 0 {
 		return false
 	}
 
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
+	// Calculate actual size by marshaling the entire batch once
+	// This gives us accurate size with minimal overhead since we marshal once per Add call
+	var actualSize int64
+	var buf bytes.Buffer
+	enc := gob.NewEncoder(&buf)
+
+	// Marshal all blocks to get actual serialized size
+	if err := enc.Encode(blocks); err != nil {
+		// If encoding fails, use estimation as fallback
+		log.Warn().Err(err).Msg("Failed to marshal blocks for size calculation, buffer size is not reported correctly")
+	} else {
+		actualSize = int64(buf.Len())
+	}
+
 	// Add to buffer
 	b.data = append(b.data, blocks...)
-	b.sizeBytes += actualSizeBytes
+	b.sizeBytes += actualSize
 
 	log.Debug().
 		Int("block_count", len(blocks)).
-		Int64("size_bytes", actualSizeBytes).
+		Int64("actual_size_bytes", actualSize).
 		Int64("total_size_bytes", b.sizeBytes).
 		Int("total_blocks", len(b.data)).
 		Msg("Added blocks to buffer")
@@ -248,3 +265,18 @@ func (s BufferStats) String() string {
 	return fmt.Sprintf("BufferStats{blocks=%d, size=%dMB, chains=%d}",
 		s.BlockCount, s.SizeBytes/(1024*1024), s.ChainCount)
 }
+
+// Close closes the buffer (no-op for in-memory buffer)
+func (b *BlockBuffer) Close() error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	// Clear the buffer to free memory
+	b.data = nil
+	b.sizeBytes = 0
+
+	return nil
+}
+
+// Ensure BlockBuffer implements IBlockBuffer interface
+var _ IBlockBuffer = (*BlockBuffer)(nil)

internal/storage/block_buffer_badger.go

@@ -20,7 +20,6 @@ type BadgerBlockBuffer struct {
 	mu           sync.RWMutex
 	db           *badger.DB
 	tempDir      string
-	sizeBytes    int64
 	maxSizeBytes int64
 	maxBlocks    int
 	blockCount   int
@@ -97,7 +96,7 @@ func NewBadgerBlockBuffer(maxSizeMB int64, maxBlocks int) (*BadgerBlockBuffer, e
 }
 
 // Add adds blocks to the buffer and returns true if flush is needed
-func (b *BadgerBlockBuffer) Add(blocks []common.BlockData, actualSizeBytes int64) bool {
+func (b *BadgerBlockBuffer) Add(blocks []common.BlockData) bool {
 	if len(blocks) == 0 {
 		return false
 	}
@@ -128,7 +127,6 @@ func (b *BadgerBlockBuffer) Add(blocks []common.BlockData, actualSizeBytes int64
 
 	// Update counters
 	b.blockCount += len(blocks)
-	b.sizeBytes += actualSizeBytes
 
 	// Update chain metadata for O(1) lookups
 	for _, block := range blocks {
@@ -154,8 +152,6 @@ func (b *BadgerBlockBuffer) Add(blocks []common.BlockData, actualSizeBytes int64
 
 	log.Debug().
 		Int("block_count", len(blocks)).
-		Int64("size_bytes", actualSizeBytes).
-		Int64("total_size_bytes", b.sizeBytes).
 		Int("total_blocks", b.blockCount).
 		Msg("Added blocks to badger buffer")
 
@@ -212,7 +208,6 @@ func (b *BadgerBlockBuffer) Flush() []common.BlockData {
 	// Reset counters and metadata
 	oldCount := b.blockCount
 	b.blockCount = 0
-	b.sizeBytes = 0
 	b.chainMetadata = make(map[uint64]*ChainMetadata)
 
 	log.Info().
@@ -233,7 +228,10 @@ func (b *BadgerBlockBuffer) ShouldFlush() bool {
 func (b *BadgerBlockBuffer) Size() (int64, int) {
 	b.mu.RLock()
 	defer b.mu.RUnlock()
-	return b.sizeBytes, b.blockCount
+	
+	// Get actual size from Badger's LSM tree
+	lsm, _ := b.db.Size()
+	return lsm, b.blockCount
 }
 
 // IsEmpty returns true if the buffer is empty
@@ -382,7 +380,6 @@ func (b *BadgerBlockBuffer) Clear() {
 	}
 
 	b.blockCount = 0
-	b.sizeBytes = 0
 	b.chainMetadata = make(map[uint64]*ChainMetadata)
 }
 
@@ -391,9 +388,12 @@ func (b *BadgerBlockBuffer) Stats() BufferStats {
 	b.mu.RLock()
 	defer b.mu.RUnlock()
 
+	// Get actual size from Badger
+	lsm, _ := b.db.Size()
+	
 	stats := BufferStats{
 		BlockCount: b.blockCount,
-		SizeBytes:  b.sizeBytes,
+		SizeBytes:  lsm,
 		ChainCount: len(b.chainMetadata),
 		ChainStats: make(map[uint64]ChainStats),
 	}
@@ -439,9 +439,12 @@ func (b *BadgerBlockBuffer) Close() error {
 // Private methods
 
 func (b *BadgerBlockBuffer) shouldFlushLocked() bool {
-	// Check size limit
-	if b.maxSizeBytes > 0 && b.sizeBytes >= b.maxSizeBytes {
-		return true
+	// Check size limit using Badger's actual size
+	if b.maxSizeBytes > 0 {
+		lsm, _ := b.db.Size()
+		if lsm >= b.maxSizeBytes {
+			return true
+		}
 	}
 
 	// Check block count limit
@@ -474,3 +477,6 @@ func (b *BadgerBlockBuffer) runGC() {
 		}
 	}
 }
+
+// Ensure BadgerBlockBuffer implements IBlockBuffer interface
+var _ IBlockBuffer = (*BadgerBlockBuffer)(nil)

internal/storage/block_buffer_badger_test.go

@@ -43,7 +43,7 @@ func TestBadgerBlockBufferMetadataOptimization(t *testing.T) {
 		},
 	}
 
-	buffer.Add(blocks, 1024)
+	buffer.Add(blocks)
 
 	// Test O(1) GetMaxBlockNumber
 	start := time.Now()
@@ -82,7 +82,7 @@ func TestBadgerBlockBufferMetadataOptimization(t *testing.T) {
 			},
 		},
 	}
-	buffer.Add(newBlocks, 512)
+	buffer.Add(newBlocks)
 
 	maxBlock = buffer.GetMaxBlockNumber(chainId)
 	assert.NotNil(t, maxBlock)
@@ -107,7 +107,7 @@ func BenchmarkBadgerBlockBufferGetMaxBlockNumber(b *testing.B) {
 				},
 			},
 		}
-		buffer.Add(blocks, 1024)
+		buffer.Add(blocks)
 	}
 
 	b.ResetTimer()
@@ -133,7 +133,7 @@ func BenchmarkBadgerBlockBufferStats(b *testing.B) {
 					},
 				},
 			}
-			buffer.Add(blocks, 1024)
+			buffer.Add(blocks)
 		}
 	}
 

internal/storage/s3.go

@@ -115,10 +115,11 @@ func NewS3Connector(cfg *config.S3StorageConfig) (*S3Connector, error) {
 	}
 
 	// Create buffer with configured settings
-	buffer, err := NewBlockBufferWithBadger(cfg.BufferSize, cfg.MaxBlocksPerFile)
+	var buffer IBlockBuffer
+	buffer, err = NewBadgerBlockBuffer(cfg.BufferSize, cfg.MaxBlocksPerFile)
 	if err != nil {
-		// Fall back to in-memory buffer if Badger fails
-		log.Warn().Err(err).Msg("Failed to create Badger buffer, falling back to in-memory buffer")
+		// fallback
+		log.Error().Err(err).Msg("Failed to create Badger buffer, falling back to in-memory buffer")
 		buffer = NewBlockBuffer(cfg.BufferSize, cfg.MaxBlocksPerFile)
 	}
 
@@ -144,27 +145,14 @@ func (s *S3Connector) InsertBlockData(data []common.BlockData) error {
 		return nil
 	}
 
-	// Calculate actual serialized size for accurate memory tracking
-	formattedData, err := s.formatter.FormatBlockData(data)
-	if err != nil {
-		return fmt.Errorf("failed to format block data for size calculation: %w", err)
-	}
-
-	// Use actual serialized size for accurate memory tracking
-	actualSize := int64(len(formattedData))
-	log.Debug().
-		Int("block_count", len(data)).
-		Int64("size_bytes", actualSize).
-		Int64("avg_bytes_per_block", actualSize/int64(len(data))).
-		Msg("Calculated actual block data size")
-
 	// Add to buffer and check if flush is needed
-	shouldFlush := s.buffer.Add(data, actualSize)
+	shouldFlush := s.buffer.Add(data)
 
 	// Start or reset timer when first data is added
 	s.timerMu.Lock()
-	sizeBytes, blockCount := s.buffer.Size()
-	if sizeBytes == actualSize && blockCount == len(data) && s.config.BufferTimeout > 0 {
+	_, blockCount := s.buffer.Size()
+	// Check if this is the first batch added (buffer was empty before)
+	if blockCount == len(data) && s.config.BufferTimeout > 0 {
 		// First data added to buffer, track time and start timer
 		s.lastAddTime = time.Now()
 		if s.flushTimer != nil {
@@ -357,13 +345,11 @@ func (s *S3Connector) Close() error {
 		// Wait for worker to finish
 		s.wg.Wait()
 
-		// Clean up buffer resources (especially important for BadgerBlockBuffer)
-		if badgerBuffer, ok := s.buffer.(*BadgerBlockBuffer); ok {
-			if err := badgerBuffer.Close(); err != nil {
-				log.Error().Err(err).Msg("Error closing badger buffer")
-				if closeErr == nil {
-					closeErr = err
-				}
+		// Clean up buffer resources
+		if err := s.buffer.Close(); err != nil {
+			log.Error().Err(err).Msg("Error closing buffer")
+			if closeErr == nil {
+				closeErr = err
 			}
 		}
 	})