Merge pull request #228 from thirdweb-dev/06-16-parallelize_staging_cleanup

Make staging cleanup async and conditional

Author: iuwqyir

internal/orchestrator/committer.go

@@ -259,12 +259,16 @@ func (c *Committer) commit(ctx context.Context, blockData []common.BlockData) er
 		}
 	}()
 
-	stagingDeleteStart := time.Now()
-	if err := c.storage.StagingStorage.DeleteStagingData(blockData); err != nil {
-		return fmt.Errorf("error deleting data from staging storage: %v", err)
+	if c.workMode == WorkModeBackfill {
+		go func() {
+			stagingDeleteStart := time.Now()
+			if err := c.storage.StagingStorage.DeleteStagingData(blockData); err != nil {
+				log.Error().Err(err).Msg("Failed to delete staging data")
+			}
+			log.Debug().Str("metric", "staging_delete_duration").Msgf("StagingStorage.DeleteStagingData duration: %f", time.Since(stagingDeleteStart).Seconds())
+			metrics.StagingDeleteDuration.Observe(time.Since(stagingDeleteStart).Seconds())
+		}()
 	}
-	log.Debug().Str("metric", "staging_delete_duration").Msgf("StagingStorage.DeleteStagingData duration: %f", time.Since(stagingDeleteStart).Seconds())
-	metrics.StagingDeleteDuration.Observe(time.Since(stagingDeleteStart).Seconds())
 
 	// Find highest block number from committed blocks
 	highestBlock := blockData[0].Block

internal/orchestrator/committer_test.go

@@ -329,12 +329,25 @@ func TestCommit(t *testing.T) {
 		{Block: common.Block{Number: big.NewInt(102)}},
 	}
 
+	// Create a channel to signal when DeleteStagingData is called
+	deleteDone := make(chan struct{})
+
 	mockMainStorage.EXPECT().InsertBlockData(blockData).Return(nil)
-	mockStagingStorage.EXPECT().DeleteStagingData(blockData).Return(nil)
+	mockStagingStorage.EXPECT().DeleteStagingData(blockData).RunAndReturn(func(data []common.BlockData) error {
+		close(deleteDone)
+		return nil
+	})
 
 	err := committer.commit(context.Background(), blockData)
-
 	assert.NoError(t, err)
+
+	// Wait for DeleteStagingData to be called with a timeout
+	select {
+	case <-deleteDone:
+		// Success - DeleteStagingData was called
+	case <-time.After(2 * time.Second):
+		t.Fatal("DeleteStagingData was not called within timeout period")
+	}
 }
 
 func TestHandleGap(t *testing.T) {