From 353d0ce688a6445c57e70500bcac3d47c6d36103 Mon Sep 17 00:00:00 2001
From: Roy Crihfield <roy@manteia.ltd>
Date: Fri, 6 Oct 2023 14:45:42 +0800
Subject: [PATCH] pass context to WriteStateSnapshot

this allows canceling the workers
---
 builder.go              | 3 ++-
 test_helpers/builder.go | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/builder.go b/builder.go
index 16a02ac..c66164c 100644
--- a/builder.go
+++ b/builder.go
@@ -168,6 +168,7 @@ func (sdb *builder) WriteStateDiff(
 
 // WriteStateDiff writes a statediff object to output sinks
 func (sdb *builder) WriteStateSnapshot(
+	ctx context.Context,
 	stateRoot common.Hash, params Params,
 	nodeSink sdtypes.StateNodeSink,
 	ipldSink sdtypes.IPLDSink,
@@ -200,7 +201,7 @@ func (sdb *builder) WriteStateSnapshot(
 		}
 	}
 	// errgroup will cancel if any group fails
-	g, ctx := errgroup.WithContext(context.Background())
+	g, ctx := errgroup.WithContext(ctx)
 	for i := range subiters {
 		func(subdiv uint) {
 			g.Go(func() error {
diff --git a/test_helpers/builder.go b/test_helpers/builder.go
index 44f46bf..18a9363 100644
--- a/test_helpers/builder.go
+++ b/test_helpers/builder.go
@@ -2,6 +2,7 @@ package test_helpers
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"math/big"
 	"math/rand"
@@ -93,6 +94,7 @@ func RunStateSnapshot(
 				tr := tracker.New(recoveryFile, subtries)
 				defer tr.CloseAndSave()
 				return builder.WriteStateSnapshot(
+					context.Background(),
 					test.StateRoot, params, stateAppender, ipldAppender, tr,
 				)
 			}