make inserts and reorgs atomic

Author: iuwqyir

internal/orchestrator/committer.go

@@ -168,18 +168,17 @@ func (c *Committer) commit(blockData []common.BlockData) error {
 	}
 	log.Debug().Msgf("Committing %d blocks", len(blockNumbers))
 
+	if err := c.storage.MainStorage.InsertBlockData(blockData); err != nil {
+		log.Error().Err(err).Msgf("Failed to commit blocks: %v", blockNumbers)
+		return fmt.Errorf("error saving data to main storage: %v", err)
+	}
+
 	go func() {
 		if err := c.publisher.PublishBlockData(blockData); err != nil {
 			log.Error().Err(err).Msg("Failed to publish block data to kafka")
 		}
 	}()
 
-	// TODO if next parts (saving or deleting) fail, we'll have to do a rollback
-	if err := c.storage.MainStorage.InsertBlockData(blockData); err != nil {
-		log.Error().Err(err).Msgf("Failed to commit blocks: %v", blockNumbers)
-		return fmt.Errorf("error saving data to main storage: %v", err)
-	}
-
 	if err := c.storage.StagingStorage.DeleteStagingData(blockData); err != nil {
 		return fmt.Errorf("error deleting data from staging storage: %v", err)
 	}

internal/orchestrator/reorg_handler.go

@@ -281,13 +281,10 @@ func (rh *ReorgHandler) handleReorg(reorgedBlockNumbers []*big.Int) error {
 		})
 		blocksToDelete = append(blocksToDelete, result.BlockNumber)
 	}
-	// TODO make delete and insert atomic
-	deletedBlockData, err := rh.storage.MainStorage.DeleteBlockData(rh.rpc.GetChainID(), blocksToDelete)
+
+	deletedBlockData, err := rh.storage.MainStorage.ReplaceBlockData(data)
 	if err != nil {
-		return fmt.Errorf("error deleting data for blocks %v: %w", blocksToDelete, err)
-	}
-	if err := rh.storage.MainStorage.InsertBlockData(data); err != nil {
-		return fmt.Errorf("error saving data to main storage: %w", err)
+		return fmt.Errorf("error replacing reorged data for blocks %v: %w", blocksToDelete, err)
 	}
 	if rh.publisher != nil {
 		// Publish block data asynchronously

internal/orchestrator/reorg_handler_test.go

@@ -504,8 +504,7 @@ func TestHandleReorg(t *testing.T) {
 	})
 	mockOrchestratorStorage.EXPECT().GetLastReorgCheckedBlockNumber(big.NewInt(1)).Return(big.NewInt(3), nil)
 
-	mockMainStorage.EXPECT().DeleteBlockData(big.NewInt(1), mock.Anything).Return([]common.BlockData{}, nil)
-	mockMainStorage.EXPECT().InsertBlockData(mock.Anything).Return(nil)
+	mockMainStorage.EXPECT().ReplaceBlockData(mock.Anything).Return([]common.BlockData{}, nil)
 
 	handler := NewReorgHandler(mockRPC, mockStorage)
 	err := handler.handleReorg([]*big.Int{big.NewInt(1), big.NewInt(2), big.NewInt(3)})
@@ -611,12 +610,9 @@ func TestHandleReorgWithSingleBlockReorg(t *testing.T) {
 		{BlockNumber: big.NewInt(105), Data: common.BlockData{}},
 	})
 
-	mockMainStorage.EXPECT().DeleteBlockData(big.NewInt(1), mock.MatchedBy(func(blocks []*big.Int) bool {
+	mockMainStorage.EXPECT().ReplaceBlockData(mock.MatchedBy(func(blocks []common.BlockData) bool {
 		return len(blocks) == 1
 	})).Return([]common.BlockData{}, nil)
-	mockMainStorage.EXPECT().InsertBlockData(mock.MatchedBy(func(data []common.BlockData) bool {
-		return len(data) == 1
-	})).Return(nil)
 
 	handler := NewReorgHandler(mockRPC, mockStorage)
 	mostRecentBlockChecked, err := handler.RunFromBlock(big.NewInt(99))
@@ -679,12 +675,9 @@ func TestHandleReorgWithLatestBlockReorged(t *testing.T) {
 		{BlockNumber: big.NewInt(108), Data: common.BlockData{}},
 	})
 
-	mockMainStorage.EXPECT().DeleteBlockData(big.NewInt(1), mock.MatchedBy(func(blocks []*big.Int) bool {
-		return len(blocks) == 8
-	})).Return([]common.BlockData{}, nil)
-	mockMainStorage.EXPECT().InsertBlockData(mock.MatchedBy(func(data []common.BlockData) bool {
+	mockMainStorage.EXPECT().ReplaceBlockData(mock.MatchedBy(func(data []common.BlockData) bool {
 		return len(data) == 8
-	})).Return(nil)
+	})).Return([]common.BlockData{}, nil)
 
 	handler := NewReorgHandler(mockRPC, mockStorage)
 	mostRecentBlockChecked, err := handler.RunFromBlock(big.NewInt(99))
@@ -743,12 +736,9 @@ func TestHandleReorgWithManyBlocks(t *testing.T) {
 		{BlockNumber: big.NewInt(103), Data: common.BlockData{}},
 	})
 
-	mockMainStorage.EXPECT().DeleteBlockData(big.NewInt(1), mock.MatchedBy(func(blocks []*big.Int) bool {
-		return len(blocks) == 5
-	})).Return([]common.BlockData{}, nil)
-	mockMainStorage.EXPECT().InsertBlockData(mock.MatchedBy(func(data []common.BlockData) bool {
+	mockMainStorage.EXPECT().ReplaceBlockData(mock.MatchedBy(func(data []common.BlockData) bool {
 		return len(data) == 5
-	})).Return(nil)
+	})).Return([]common.BlockData{}, nil)
 
 	handler := NewReorgHandler(mockRPC, mockStorage)
 	mostRecentBlockChecked, err := handler.RunFromBlock(big.NewInt(99))