Merge pull request #4212 from onflow/leo/fix-read-leaf-nodes

Fix read leaf nodes

Author: zhangchiqing

ledger/complete/wal/checkpoint_v6_leaf_reader.go

@@ -18,11 +18,6 @@ type LeafNode struct {
 	Payload *ledger.Payload
 }
 
-type LeafNodeResult struct {
-	LeafNode *LeafNode
-	Err      error
-}
-
 func nodeToLeaf(leaf *node.Node) *LeafNode {
 	return &LeafNode{
 		Hash:    leaf.Hash(),
@@ -31,48 +26,50 @@ func nodeToLeaf(leaf *node.Node) *LeafNode {
 	}
 }
 
-func OpenAndReadLeafNodesFromCheckpointV6(dir string, fileName string, logger *zerolog.Logger) (
-	allLeafNodesCh <-chan LeafNodeResult, errToReturn error) {
+// OpenAndReadLeafNodesFromCheckpointV6 takes a channel for pushing the leaf nodes that are read from
+// the given checkpoint file specified by dir and fileName.
+// It returns when finish reading the checkpoint file and the input channel can be closed.
+func OpenAndReadLeafNodesFromCheckpointV6(allLeafNodesCh chan<- *LeafNode, dir string, fileName string, logger *zerolog.Logger) (errToReturn error) {
+	// we are the only sender of the channel, closing it after done
+	defer func() {
+		close(allLeafNodesCh)
+	}()
 
 	filepath := filePathCheckpointHeader(dir, fileName)
 
 	f, err := os.Open(filepath)
 	if err != nil {
-		return nil, fmt.Errorf("could not open file %v: %w", filepath, err)
+		return fmt.Errorf("could not open file %v: %w", filepath, err)
 	}
 	defer func(file *os.File) {
 		errToReturn = closeAndMergeError(file, errToReturn)
 	}(f)
 
 	subtrieChecksums, _, err := readCheckpointHeader(filepath, logger)
 	if err != nil {
-		return nil, fmt.Errorf("could not read header: %w", err)
+		return fmt.Errorf("could not read header: %w", err)
 	}
 
 	// ensure all checkpoint part file exists, might return os.ErrNotExist error
 	// if a file is missing
 	err = allPartFileExist(dir, fileName, len(subtrieChecksums))
 	if err != nil {
-		return nil, fmt.Errorf("fail to check all checkpoint part file exist: %w", err)
+		return fmt.Errorf("fail to check all checkpoint part file exist: %w", err)
 	}
 
-	bufSize := 1000
-	leafNodesCh := make(chan LeafNodeResult, bufSize)
-	allLeafNodesCh = leafNodesCh
-	defer func() {
-		close(leafNodesCh)
-	}()
-
 	// push leaf nodes to allLeafNodesCh
 	for i, checksum := range subtrieChecksums {
-		readCheckpointSubTrieLeafNodes(leafNodesCh, dir, fileName, i, checksum, logger)
+		err := readCheckpointSubTrieLeafNodes(allLeafNodesCh, dir, fileName, i, checksum, logger)
+		if err != nil {
+			return fmt.Errorf("fail to read checkpoint leaf nodes from %v-th subtrie file: %w", i, err)
+		}
 	}
 
-	return allLeafNodesCh, nil
+	return nil
 }
 
-func readCheckpointSubTrieLeafNodes(leafNodesCh chan<- LeafNodeResult, dir string, fileName string, index int, checksum uint32, logger *zerolog.Logger) {
-	err := processCheckpointSubTrie(dir, fileName, index, checksum, logger,
+func readCheckpointSubTrieLeafNodes(leafNodesCh chan<- *LeafNode, dir string, fileName string, index int, checksum uint32, logger *zerolog.Logger) error {
+	return processCheckpointSubTrie(dir, fileName, index, checksum, logger,
 		func(reader *Crc32Reader, nodesCount uint64) error {
 			scratch := make([]byte, 1024*4) // must not be less than 1024
 
@@ -89,21 +86,11 @@ func readCheckpointSubTrieLeafNodes(leafNodesCh chan<- LeafNodeResult, dir strin
 					return fmt.Errorf("cannot read node %d: %w", i, err)
 				}
 				if node.IsLeaf() {
-					leafNodesCh <- LeafNodeResult{
-						LeafNode: nodeToLeaf(node),
-						Err:      nil,
-					}
+					leafNodesCh <- nodeToLeaf(node)
 				}
 
 				logging(i)
 			}
 			return nil
 		})
-
-	if err != nil {
-		leafNodesCh <- LeafNodeResult{
-			LeafNode: nil,
-			Err:      err,
-		}
-	}
 }

ledger/complete/wal/checkpoint_v6_test.go

@@ -140,7 +140,7 @@ func createMultipleRandomTriesMini(t *testing.T) []*trie.MTrie {
 	var err error
 	// add tries with no shared paths
 	for i := 0; i < 5; i++ {
-		paths, payloads := randNPathPayloads(10)
+		paths, payloads := randNPathPayloads(20)
 		activeTrie, _, err = trie.NewTrieWithUpdatedRegisters(activeTrie, paths, payloads, false)
 		require.NoError(t, err, "update registers")
 		tries = append(tries, activeTrie)
@@ -318,9 +318,14 @@ func TestWriteAndReadCheckpointV6LeafEmptyTrie(t *testing.T) {
 		fileName := "checkpoint-empty-trie"
 		logger := unittest.Logger()
 		require.NoErrorf(t, StoreCheckpointV6Concurrently(tries, dir, fileName, &logger), "fail to store checkpoint")
-		resultChan, err := OpenAndReadLeafNodesFromCheckpointV6(dir, fileName, &logger)
-		require.NoErrorf(t, err, "fail to read checkpoint %v/%v", dir, fileName)
-		for range resultChan {
+
+		bufSize := 10
+		leafNodesCh := make(chan *LeafNode, bufSize)
+		go func() {
+			err := OpenAndReadLeafNodesFromCheckpointV6(leafNodesCh, dir, fileName, &logger)
+			require.NoErrorf(t, err, "fail to read checkpoint %v/%v", dir, fileName)
+		}()
+		for range leafNodesCh {
 			require.Fail(t, "should not return any nodes")
 		}
 	})
@@ -332,14 +337,17 @@ func TestWriteAndReadCheckpointV6LeafSimpleTrie(t *testing.T) {
 		fileName := "checkpoint"
 		logger := unittest.Logger()
 		require.NoErrorf(t, StoreCheckpointV6Concurrently(tries, dir, fileName, &logger), "fail to store checkpoint")
-		resultChan, err := OpenAndReadLeafNodesFromCheckpointV6(dir, fileName, &logger)
-		require.NoErrorf(t, err, "fail to read checkpoint %v/%v", dir, fileName)
+		bufSize := 1
+		leafNodesCh := make(chan *LeafNode, bufSize)
+		go func() {
+			err := OpenAndReadLeafNodesFromCheckpointV6(leafNodesCh, dir, fileName, &logger)
+			require.NoErrorf(t, err, "fail to read checkpoint %v/%v", dir, fileName)
+		}()
 		resultPayloads := make([]ledger.Payload, 0)
-		for readResult := range resultChan {
-			require.NoError(t, readResult.Err, "no errors in read results")
+		for leafNode := range leafNodesCh {
 			// avoid dummy payload from empty trie
-			if readResult.LeafNode.Payload != nil {
-				resultPayloads = append(resultPayloads, *readResult.LeafNode.Payload)
+			if leafNode.Payload != nil {
+				resultPayloads = append(resultPayloads, *leafNode.Payload)
 			}
 		}
 		require.EqualValues(t, tries[1].AllPayloads(), resultPayloads)
@@ -352,12 +360,15 @@ func TestWriteAndReadCheckpointV6LeafMultipleTries(t *testing.T) {
 		tries := createMultipleRandomTriesMini(t)
 		logger := unittest.Logger()
 		require.NoErrorf(t, StoreCheckpointV6Concurrently(tries, dir, fileName, &logger), "fail to store checkpoint")
-		resultChan, err := OpenAndReadLeafNodesFromCheckpointV6(dir, fileName, &logger)
-		require.NoErrorf(t, err, "fail to read checkpoint %v/%v", dir, fileName)
+		bufSize := 5
+		leafNodesCh := make(chan *LeafNode, bufSize)
+		go func() {
+			err := OpenAndReadLeafNodesFromCheckpointV6(leafNodesCh, dir, fileName, &logger)
+			require.NoErrorf(t, err, "fail to read checkpoint %v/%v", dir, fileName)
+		}()
 		resultPayloads := make([]ledger.Payload, 0)
-		for readResult := range resultChan {
-			require.NoError(t, readResult.Err, "no errors in read results")
-			resultPayloads = append(resultPayloads, *readResult.LeafNode.Payload)
+		for leafNode := range leafNodesCh {
+			resultPayloads = append(resultPayloads, *leafNode.Payload)
 		}
 		require.NotEmpty(t, resultPayloads)
 	})
@@ -528,7 +539,9 @@ func TestAllPartFileExistLeafReader(t *testing.T) {
 			err = os.Remove(fileToDelete)
 			require.NoError(t, err, "fail to remove part file")
 
-			_, err = OpenAndReadLeafNodesFromCheckpointV6(dir, fileName, &logger)
+			bufSize := 10
+			leafNodesCh := make(chan *LeafNode, bufSize)
+			err = OpenAndReadLeafNodesFromCheckpointV6(leafNodesCh, dir, fileName, &logger)
 			require.ErrorIs(t, err, os.ErrNotExist, "wrong error type returned")
 		}
 	})