Add BatchPreload to decode slabs in parallel and cache

storage.go

@@ -1560,3 +1560,155 @@ func (s *PersistentSlabStorage) getAllChildReferences(slab Slab) (
 
  return references, brokenReferences, nil
 }
+
+// BatchPreload decodeds and caches slabs of given ids in parallel.
+// This is useful for storage health or data validation in migration programs.
+func (s *PersistentSlabStorage) BatchPreload(ids []SlabID, numWorkers int) error {
+ if len(ids) == 0 {
+ return nil
+ }
+
+ // Use 11 for min slab count for parallel decoding because micro benchmarks showed
+ // performance regression for <= 10 slabs when decoding slabs in parallel.
+ const minCountForBatchPreload = 11
+ if len(ids) < minCountForBatchPreload {
+
+ for _, id := range ids {
+ // fetch from base storage last
+ data, ok, err := s.baseStorage.Retrieve(id)
+ if err != nil {
+ // Wrap err as external error (if needed) because err is returned by BaseStorage interface.
+ return wrapErrorfAsExternalErrorIfNeeded(err, fmt.Sprintf("failed to retrieve slab %s", id))
+ }
+ if !ok {
+ continue
+ }
+
+ slab, err := DecodeSlab(id, data, s.cborDecMode, s.DecodeStorable, s.DecodeTypeInfo)
+ if err != nil {
+ // err is already categorized by DecodeSlab().
+ return err
+ }
+
+ // save decoded slab to cache
+ s.cache[id] = slab
+ }
+
+ return nil
+ }
+
+ type slabToBeDecoded struct {
+ slabID SlabID
+ data []byte
+ }
+
+ type decodedSlab struct {
+ slabID SlabID
+ slab Slab
+ err error
+ }
+
+ // Define decoder (worker) to decode slabs in parallel
+ decoder := func(wg *sync.WaitGroup, done <-chan struct{}, jobs <-chan slabToBeDecoded, results chan<- decodedSlab) {
+ defer wg.Done()
+
+ for slabData := range jobs {
+ // Check if goroutine is signaled to stop before proceeding.
+ select {
+ case <-done:
+ return
+ default:
+ }
+
+ id := slabData.slabID
+ data := slabData.data
+
+ slab, err := DecodeSlab(id, data, s.cborDecMode, s.DecodeStorable, s.DecodeTypeInfo)
+ // err is already categorized by DecodeSlab().
+ results <- decodedSlab{
+ slabID: id,
+ slab: slab,
+ err: err,
+ }
+ }
+ }
+
+ if numWorkers > len(ids) {
+ numWorkers = len(ids)
+ }
+
+ var wg sync.WaitGroup
+
+ // Construct done signal channel
+ done := make(chan struct{})
+
+ // Construct job queue
+ jobs := make(chan slabToBeDecoded, len(ids))
+
+ // Construct result queue
+ results := make(chan decodedSlab, len(ids))
+
+ defer func() {
+ // This ensures that all goroutines are stopped before output channel is closed.
+
+ // Wait for all goroutines to finish
+ wg.Wait()
+
+ // Close output channel
+ close(results)
+ }()
+
+ // Preallocate cache map if empty
+ if len(s.cache) == 0 {
+ s.cache = make(map[SlabID]Slab, len(ids))
+ }
+
+ // Launch workers
+ wg.Add(numWorkers)
+ for i := 0; i < numWorkers; i++ {
+ go decoder(&wg, done, jobs, results)
+ }
+
+ // Send jobs
+ jobCount := 0
+ {
+ // Need to close input channel (jobs) here because
+ // if there isn't any job in jobs channel,
+ // done is never processed inside loop "for slabData := range jobs".
+ defer close(jobs)
+
+ for _, id := range ids {
+ // fetch from base storage last
+ data, ok, err := s.baseStorage.Retrieve(id)
+ if err != nil {
+ // Closing done channel signals goroutines to stop.
+ close(done)
+ // Wrap err as external error (if needed) because err is returned by BaseStorage interface.
+ return wrapErrorfAsExternalErrorIfNeeded(err, fmt.Sprintf("failed to retrieve slab %s", id))
+ }
+ if !ok {
+ continue
+ }
+
+ jobs <- slabToBeDecoded{id, data}
+ jobCount++
+ }
+ }
+
+ // Process results
+ for i := 0; i < jobCount; i++ {
+ result := <-results
+
+ if result.err != nil {
+ // Closing done channel signals goroutines to stop.
+ close(done)
+ // result.err is already categorized by DecodeSlab().
+ return result.err
+ }
+
+ // save decoded slab to cache
+ s.cache[result.slabID] = result.slab
+ }
+
+ return nil
+}

storage_bench_test.go

@@ -132,3 +132,122 @@ func BenchmarkStorageNondeterministicFastCommit(b *testing.B) {
  benchmarkNondeterministicFastCommit(b, fixedSeed, 100_000)
  benchmarkNondeterministicFastCommit(b, fixedSeed, 1_000_000)
 }
+
+func benchmarkRetrieve(b *testing.B, seed int64, numberOfSlabs int) {
+
+ r := rand.New(rand.NewSource(seed))
+
+ encMode, err := cbor.EncOptions{}.EncMode()
+ require.NoError(b, err)
+
+ decMode, err := cbor.DecOptions{}.DecMode()
+ require.NoError(b, err)
+
+ encodedSlabs := make(map[SlabID][]byte)
+ ids := make([]SlabID, 0, numberOfSlabs)
+ for i := 0; i < numberOfSlabs; i++ {
+ addr := generateRandomAddress(r)
+
+ var index SlabIndex
+ binary.BigEndian.PutUint64(index[:], uint64(i))
+
+ id := SlabID{addr, index}
+
+ slab := generateLargeSlab(id)
+
+ data, err := EncodeSlab(slab, encMode)
+ require.NoError(b, err)
+
+ encodedSlabs[id] = data
+ ids = append(ids, id)
+ }
+
+ b.Run(strconv.Itoa(numberOfSlabs), func(b *testing.B) {
+ for i := 0; i < b.N; i++ {
+ b.StopTimer()
+
+ baseStorage := NewInMemBaseStorageFromMap(encodedSlabs)
+ storage := NewPersistentSlabStorage(baseStorage, encMode, decMode, decodeStorable, decodeTypeInfo)
+
+ b.StartTimer()
+
+ for _, id := range ids {
+ _, found, err := storage.Retrieve(id)
+ require.True(b, found)
+ require.NoError(b, err)
+ }
+ }
+ })
+}
+
+func benchmarkBatchPreload(b *testing.B, seed int64, numberOfSlabs int) {
+
+ r := rand.New(rand.NewSource(seed))
+
+ encMode, err := cbor.EncOptions{}.EncMode()
+ require.NoError(b, err)
+
+ decMode, err := cbor.DecOptions{}.DecMode()
+ require.NoError(b, err)
+
+ encodedSlabs := make(map[SlabID][]byte)
+ ids := make([]SlabID, 0, numberOfSlabs)
+ for i := 0; i < numberOfSlabs; i++ {
+ addr := generateRandomAddress(r)
+
+ var index SlabIndex
+ binary.BigEndian.PutUint64(index[:], uint64(i))
+
+ id := SlabID{addr, index}
+
+ slab := generateLargeSlab(id)
+
+ data, err := EncodeSlab(slab, encMode)
+ require.NoError(b, err)
+
+ encodedSlabs[id] = data
+ ids = append(ids, id)
+ }
+
+ b.Run(strconv.Itoa(numberOfSlabs), func(b *testing.B) {
+ for i := 0; i < b.N; i++ {
+ b.StopTimer()
+
+ baseStorage := NewInMemBaseStorageFromMap(encodedSlabs)
+ storage := NewPersistentSlabStorage(baseStorage, encMode, decMode, decodeStorable, decodeTypeInfo)
+
+ b.StartTimer()
+
+ err = storage.BatchPreload(ids, runtime.NumCPU())
+ require.NoError(b, err)
+
+ for _, id := range ids {
+ _, found, err := storage.Retrieve(id)
+ require.True(b, found)
+ require.NoError(b, err)
+ }
+ }
+ })
+}
+
+func BenchmarkStorageRetrieve(b *testing.B) {
+ fixedSeed := int64(1234567) // intentionally use fixed constant rather than time, etc.
+
+ benchmarkRetrieve(b, fixedSeed, 10)
+ benchmarkRetrieve(b, fixedSeed, 100)
+ benchmarkRetrieve(b, fixedSeed, 1_000)
+ benchmarkRetrieve(b, fixedSeed, 10_000)
+ benchmarkRetrieve(b, fixedSeed, 100_000)
+ benchmarkRetrieve(b, fixedSeed, 1_000_000)
+}
+
+func BenchmarkStorageBatchPreload(b *testing.B) {
+ fixedSeed := int64(1234567) // intentionally use fixed constant rather than time, etc.
+
+ benchmarkBatchPreload(b, fixedSeed, 10)
+ benchmarkBatchPreload(b, fixedSeed, 100)
+ benchmarkBatchPreload(b, fixedSeed, 1_000)
+ benchmarkBatchPreload(b, fixedSeed, 10_000)
+ benchmarkBatchPreload(b, fixedSeed, 100_000)
+ benchmarkBatchPreload(b, fixedSeed, 1_000_000)
+}