From e87084aa50cb4219525abfe52f7fe42d7b800e80 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Thu, 11 Jul 2024 11:46:50 -0700
Subject: [PATCH] small refactor

---
 storage/badger/commits.go                     |  4 ----
 ...utaton_result.go => computation_result.go} |  0
 ...ult_test.go => computation_result_test.go} |  0
 storage/pebble/headers.go                     | 19 ++++++-------------
 storage/pebble/operation/batch.go             |  1 +
 storage/pebble/operation/common.go            |  2 +-
 storage/pebble/qcs.go                         |  2 +-
 7 files changed, 9 insertions(+), 19 deletions(-)
 rename storage/pebble/{computaton_result.go => computation_result.go} (100%)
 rename storage/pebble/{computaton_result_test.go => computation_result_test.go} (100%)

diff --git a/storage/badger/commits.go b/storage/badger/commits.go
index 9dbf6601699..11a4e4aa8e2 100644
--- a/storage/badger/commits.go
+++ b/storage/badger/commits.go
@@ -70,10 +70,6 @@ func (c *Commits) BatchStore(blockID flow.Identifier, commit flow.StateCommitmen
 	return operation.BatchIndexStateCommitment(blockID, commit)(writeBatch)
 }
 
-func (c *Commits) BatchStore2(blockID flow.Identifier, commit flow.StateCommitment, tx storage.BatchWriter) error {
-	return operation.BatchIndexStateCommitment(blockID, commit)(tx)
-}
-
 func (c *Commits) ByBlockID(blockID flow.Identifier) (flow.StateCommitment, error) {
 	tx := c.db.NewTransaction(false)
 	defer tx.Discard()
diff --git a/storage/pebble/computaton_result.go b/storage/pebble/computation_result.go
similarity index 100%
rename from storage/pebble/computaton_result.go
rename to storage/pebble/computation_result.go
diff --git a/storage/pebble/computaton_result_test.go b/storage/pebble/computation_result_test.go
similarity index 100%
rename from storage/pebble/computaton_result_test.go
rename to storage/pebble/computation_result_test.go
diff --git a/storage/pebble/headers.go b/storage/pebble/headers.go
index be0c4848243..c42bcb9bf53 100644
--- a/storage/pebble/headers.go
+++ b/storage/pebble/headers.go
@@ -21,6 +21,10 @@ type Headers struct {
 
 func NewHeaders(collector module.CacheMetrics, db *pebble.DB) *Headers {
 
+	store := func(blockID flow.Identifier, header *flow.Header) func(storage.PebbleReaderBatchWriter) error {
+		return storage.OnlyWriter(operation.InsertHeader(blockID, header))
+	}
+
 	// CAUTION: should only be used to index FINALIZED blocks by their
 	// respective height
 	storeHeight := func(height uint64, id flow.Identifier) func(storage.PebbleReaderBatchWriter) error {
@@ -47,6 +51,7 @@ func NewHeaders(collector module.CacheMetrics, db *pebble.DB) *Headers {
 		db: db,
 		cache: newCache(collector, metrics.ResourceHeader,
 			withLimit[flow.Identifier, *flow.Header](4*flow.DefaultTransactionExpiry),
+			withStore(store),
 			withRetrieve(retrieve)),
 
 		heightCache: newCache(collector, metrics.ResourceFinalizedHeight,
@@ -59,19 +64,7 @@ func NewHeaders(collector module.CacheMetrics, db *pebble.DB) *Headers {
 }
 
 func (h *Headers) storePebble(blockID flow.Identifier, header *flow.Header) func(storage.PebbleReaderBatchWriter) error {
-	return func(rw storage.PebbleReaderBatchWriter) error {
-		rw.AddCallback(func() {
-			h.cache.Insert(blockID, header)
-		})
-
-		_, tx := rw.ReaderWriter()
-		err := operation.InsertHeader(blockID, header)(tx)
-		if err != nil {
-			return fmt.Errorf("could not store header %v: %w", blockID, err)
-		}
-
-		return nil
-	}
+	return h.cache.PutPebble(blockID, header)
 }
 
 func (h *Headers) retrieveTx(blockID flow.Identifier) func(pebble.Reader) (*flow.Header, error) {
diff --git a/storage/pebble/operation/batch.go b/storage/pebble/operation/batch.go
index bd57aebc5f7..9f92edc4dc0 100644
--- a/storage/pebble/operation/batch.go
+++ b/storage/pebble/operation/batch.go
@@ -9,6 +9,7 @@ import (
 	"github.com/onflow/flow-go/storage"
 )
 
+// batchWriter wraps the storage.BatchWriter to make it compatible with pebble.Writer
 type batchWriter struct {
 	batch storage.BatchWriter
 }
diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index 2fcb8597f20..221a438b21e 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -379,7 +379,7 @@ func findHighestAtOrBelow(
 
 		key := append(prefix, b(height)...)
 		it, err := r.NewIter(&pebble.IterOptions{
-			UpperBound: append(key, 0xFF),
+			UpperBound: append(key, ffBytes...),
 		})
 		if err != nil {
 			return fmt.Errorf("can not create iterator: %w", err)
diff --git a/storage/pebble/qcs.go b/storage/pebble/qcs.go
index 03ad32a4273..b8e651aa7be 100644
--- a/storage/pebble/qcs.go
+++ b/storage/pebble/qcs.go
@@ -15,7 +15,7 @@ import (
 
 // QuorumCertificates implements persistent storage for quorum certificates.
 type QuorumCertificates struct {
-	storing *sync.Mutex
+	storing sync.Mutex
 	db      *pebble.DB
 	cache   *Cache[flow.Identifier, *flow.QuorumCertificate]
 }