swarm: Adding context to more functions

swarm/api/api.go

@@ -239,7 +239,7 @@ func (a *API) Retrieve(ctx context.Context, addr storage.Address) (reader storag
 }
 
 // Store wraps the Store API call of the embedded FileStore
-func (a *API) Store(ctx context.Context, data io.Reader, size int64, toEncrypt bool) (addr storage.Address, wait func(), err error) {
+func (a *API) Store(ctx context.Context, data io.Reader, size int64, toEncrypt bool) (addr storage.Address, wait func(ctx context.Context) error, err error) {
 	log.Debug("api.store", "size", size)
 	return a.fileStore.Store(ctx, data, size, toEncrypt)
 }
@@ -286,7 +286,7 @@ func (a *API) Resolve(ctx context.Context, uri *URI) (storage.Address, error) {
 }
 
 // Put provides singleton manifest creation on top of FileStore store
-func (a *API) Put(ctx context.Context, content string, contentType string, toEncrypt bool) (k storage.Address, wait func(), err error) {
+func (a *API) Put(ctx context.Context, content string, contentType string, toEncrypt bool) (k storage.Address, wait func(context.Context) error, err error) {
 	apiPutCount.Inc(1)
 	r := strings.NewReader(content)
 	key, waitContent, err := a.fileStore.Store(ctx, r, int64(len(content)), toEncrypt)
@@ -301,9 +301,12 @@ func (a *API) Put(ctx context.Context, content string, contentType string, toEnc
 		apiPutFail.Inc(1)
 		return nil, nil, err
 	}
-	return key, func() {
+	return key, func(ctx context.Context) error {
-		waitContent()
+		err := waitContent(ctx)
-		waitManifest()
+		if err != nil {
+			return err
+		}
+		return waitManifest(ctx)
 	}, nil
 }
 

swarm/api/api_test.go

@@ -109,11 +109,15 @@ func TestApiPut(t *testing.T) {
 	testAPI(t, func(api *API, toEncrypt bool) {
 		content := "hello"
 		exp := expResponse(content, "text/plain", 0)
-		addr, wait, err := api.Put(context.TODO(), content, exp.MimeType, toEncrypt)
+		ctx := context.TODO()
+		addr, wait, err := api.Put(ctx, content, exp.MimeType, toEncrypt)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
-		wait()
 		resp := testGet(t, api, addr.Hex(), "")
 		checkResponse(t, resp, exp)
 	})

swarm/api/filesystem.go

@@ -114,12 +114,13 @@ func (fs *FileSystem) Upload(lpath, index string, toEncrypt bool) (string, error
 			if err == nil {
 				stat, _ := f.Stat()
 				var hash storage.Address
-				var wait func()
+				var wait func(context.Context) error
-				hash, wait, err = fs.api.fileStore.Store(context.TODO(), f, stat.Size(), toEncrypt)
+				ctx := context.TODO()
+				hash, wait, err = fs.api.fileStore.Store(ctx, f, stat.Size(), toEncrypt)
 				if hash != nil {
 					list[i].Hash = hash.Hex()
 				}
-				wait()
+				err = wait(ctx)
 				awg.Done()
 				if err == nil {
 					first512 := make([]byte, 512)

swarm/api/filesystem_test.go

@@ -106,8 +106,13 @@ func TestApiDirUploadModify(t *testing.T) {
 			t.Errorf("unexpected error: %v", err)
 			return
 		}
-		hash, wait, err := api.Store(context.TODO(), bytes.NewReader(index), int64(len(index)), toEncrypt)
+		ctx := context.TODO()
-		wait()
+		hash, wait, err := api.Store(ctx, bytes.NewReader(index), int64(len(index)), toEncrypt)
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+			return
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Errorf("unexpected error: %v", err)
 			return

swarm/api/http/server_test.go

@@ -383,15 +383,19 @@ func testBzzGetPath(encrypted bool, t *testing.T) {
 
 	for i, mf := range testmanifest {
 		reader[i] = bytes.NewReader([]byte(mf))
-		var wait func()
+		var wait func(context.Context) error
-		addr[i], wait, err = srv.FileStore.Store(context.TODO(), reader[i], int64(len(mf)), encrypted)
+		ctx := context.TODO()
+		addr[i], wait, err = srv.FileStore.Store(ctx, reader[i], int64(len(mf)), encrypted)
 		for j := i + 1; j < len(testmanifest); j++ {
 			testmanifest[j] = strings.Replace(testmanifest[j], fmt.Sprintf("<key%v>", i), addr[i].Hex(), -1)
 		}
 		if err != nil {
 			t.Fatal(err)
 		}
-		wait()
+		err = wait(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
 	}
 
 	rootRef := addr[2].Hex()

swarm/api/manifest.go

@@ -69,7 +69,7 @@ func (a *API) NewManifest(ctx context.Context, toEncrypt bool) (storage.Address,
 		return nil, err
 	}
 	key, wait, err := a.Store(ctx, bytes.NewReader(data), int64(len(data)), toEncrypt)
-	wait()
+	wait(ctx)
 	return key, err
 }
 
@@ -382,8 +382,12 @@ func (mt *manifestTrie) recalcAndStore() error {
 	}
 
 	sr := bytes.NewReader(manifest)
-	key, wait, err2 := mt.fileStore.Store(context.TODO(), sr, int64(len(manifest)), mt.encrypted)
+	ctx := context.TODO()
-	wait()
+	key, wait, err2 := mt.fileStore.Store(ctx, sr, int64(len(manifest)), mt.encrypted)
+	if err2 != nil {
+		return err2
+	}
+	err2 = wait(ctx)
 	mt.ref = key
 	return err2
 }

swarm/api/storage.go

@@ -46,7 +46,7 @@ func NewStorage(api *API) *Storage {
 // its content type
 //
 // DEPRECATED: Use the HTTP API instead
-func (s *Storage) Put(ctx context.Context, content string, contentType string, toEncrypt bool) (storage.Address, func(), error) {
+func (s *Storage) Put(ctx context.Context, content string, contentType string, toEncrypt bool) (storage.Address, func(context.Context) error, error) {
 	return s.api.Put(ctx, content, contentType, toEncrypt)
 }
 

swarm/api/storage_test.go

@@ -32,11 +32,15 @@ func TestStoragePutGet(t *testing.T) {
 		content := "hello"
 		exp := expResponse(content, "text/plain", 0)
 		// exp := expResponse([]byte(content), "text/plain", 0)
-		bzzkey, wait, err := api.Put(context.TODO(), content, exp.MimeType, toEncrypt)
+		ctx := context.TODO()
+		bzzkey, wait, err := api.Put(ctx, content, exp.MimeType, toEncrypt)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
-		wait()
 		bzzhash := bzzkey.Hex()
 		// to check put against the API#Get
 		resp0 := testGet(t, api.api, bzzhash, "")

swarm/network/stream/delivery_test.go

@@ -345,9 +345,13 @@ func testDeliveryFromNodes(t *testing.T, nodes, conns, chunkCount int, skipCheck
 	// here we distribute chunks of a random file into Stores of nodes 1 to nodes
 	rrFileStore := storage.NewFileStore(newRoundRobinStore(sim.Stores[1:]...), storage.NewFileStoreParams())
 	size := chunkCount * chunkSize
-	fileHash, wait, err := rrFileStore.Store(context.TODO(), io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+	ctx := context.TODO()
+	fileHash, wait, err := rrFileStore.Store(ctx, io.LimitReader(crand.Reader, int64(size)), int64(size), false)
 	// wait until all chunks stored
-	wait()
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	err = wait(ctx)
 	if err != nil {
 		t.Fatal(err.Error())
 	}
@@ -627,9 +631,13 @@ Loop:
 		hashes := make([]storage.Address, chunkCount)
 		for i := 0; i < chunkCount; i++ {
 			// create actual size real chunks
-			hash, wait, err := remoteFileStore.Store(context.TODO(), io.LimitReader(crand.Reader, int64(chunkSize)), int64(chunkSize), false)
+			ctx := context.TODO()
+			hash, wait, err := remoteFileStore.Store(ctx, io.LimitReader(crand.Reader, int64(chunkSize)), int64(chunkSize), false)
+			if err != nil {
+				b.Fatalf("expected no error. got %v", err)
+			}
 			// wait until all chunks stored
-			wait()
+			err = wait(ctx)
 			if err != nil {
 				b.Fatalf("expected no error. got %v", err)
 			}

swarm/network/stream/intervals_test.go

@@ -117,8 +117,12 @@ func testIntervals(t *testing.T, live bool, history *Range, skipCheck bool) {
 
 	fileStore := storage.NewFileStore(sim.Stores[0], storage.NewFileStoreParams())
 	size := chunkCount * chunkSize
-	_, wait, err := fileStore.Store(context.TODO(), io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+	ctx := context.TODO()
-	wait()
+	_, wait, err := fileStore.Store(ctx, io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = wait(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}

swarm/network/stream/snapshot_retrieval_test.go

@@ -765,9 +765,13 @@ func uploadFilesToNodes(nodes []*simulations.Node) ([]storage.Address, []string,
 			return nil, nil, err
 		}
 		//store it (upload it) on the FileStore
-		rk, wait, err := fileStore.Store(context.TODO(), strings.NewReader(rfiles[i]), int64(len(rfiles[i])), false)
+		ctx := context.TODO()
+		rk, wait, err := fileStore.Store(ctx, strings.NewReader(rfiles[i]), int64(len(rfiles[i])), false)
 		log.Debug("Uploaded random string file to node")
-		wait()
+		if err != nil {
+			return nil, nil, err
+		}
+		err = wait(ctx)
 		if err != nil {
 			return nil, nil, err
 		}

swarm/network/stream/snapshot_sync_test.go

@@ -581,8 +581,12 @@ func uploadFileToSingleNodeStore(id discover.NodeID, chunkCount int) ([]storage.
 	fileStore := storage.NewFileStore(lstore, storage.NewFileStoreParams())
 	var rootAddrs []storage.Address
 	for i := 0; i < chunkCount; i++ {
-		rk, wait, err := fileStore.Store(context.TODO(), io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+		ctx := context.TODO()
-		wait()
+		rk, wait, err := fileStore.Store(ctx, io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+		if err != nil {
+			return nil, err
+		}
+		err = wait(ctx)
 		if err != nil {
 			return nil, err
 		}

swarm/network/stream/syncer_test.go

@@ -202,9 +202,12 @@ func testSyncBetweenNodes(t *testing.T, nodes, conns, chunkCount int, skipCheck
 		// here we distribute chunks of a random file into stores 1...nodes
 		rrFileStore := storage.NewFileStore(newRoundRobinStore(sim.Stores[1:]...), storage.NewFileStoreParams())
 		size := chunkCount * chunkSize
-		_, wait, err := rrFileStore.Store(context.TODO(), io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+		_, wait, err := rrFileStore.Store(ctx, io.LimitReader(crand.Reader, int64(size)), int64(size), false)
+		if err != nil {
+			t.Fatal(err.Error())
+		}
 		// need to wait cos we then immediately collect the relevant bin content
-		wait()
+		wait(ctx)
 		if err != nil {
 			t.Fatal(err.Error())
 		}

swarm/network_test.go

@@ -508,14 +508,15 @@ func uploadFile(swarm *Swarm) (storage.Address, string, error) {
 	// File data is very short, but it is ensured that its
 	// uniqueness is very certain.
 	data := fmt.Sprintf("test content %s %x", time.Now().Round(0), b)
-	k, wait, err := swarm.api.Put(context.TODO(), data, "text/plain", false)
+	ctx := context.TODO()
+	k, wait, err := swarm.api.Put(ctx, data, "text/plain", false)
 	if err != nil {
 		return nil, "", err
 	}
 	if wait != nil {
-		wait()
+		err = wait(ctx)
 	}
-	return k, data, nil
+	return k, data, err
 }
 
 // retrieve is the function that is used for checking the availability of

swarm/storage/chunker.go

@@ -16,6 +16,7 @@
 package storage
 
 import (
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -126,7 +127,7 @@ type TreeChunker struct {
 	The chunks are not meant to be validated by the chunker when joining. This
 	is because it is left to the DPA to decide which sources are trusted.
 */
-func TreeJoin(addr Address, getter Getter, depth int) *LazyChunkReader {
+func TreeJoin(ctx context.Context, addr Address, getter Getter, depth int) *LazyChunkReader {
 	jp := &JoinerParams{
 		ChunkerParams: ChunkerParams{
 			chunkSize: DefaultChunkSize,
@@ -137,14 +138,14 @@ func TreeJoin(addr Address, getter Getter, depth int) *LazyChunkReader {
 		depth:  depth,
 	}
 
-	return NewTreeJoiner(jp).Join()
+	return NewTreeJoiner(jp).Join(ctx)
 }
 
 /*
 	When splitting, data is given as a SectionReader, and the key is a hashSize long byte slice (Key), the root hash of the entire content will fill this once processing finishes.
 	New chunks to store are store using the putter which the caller provides.
 */
-func TreeSplit(data io.Reader, size int64, putter Putter) (k Address, wait func(), err error) {
+func TreeSplit(ctx context.Context, data io.Reader, size int64, putter Putter) (k Address, wait func(context.Context) error, err error) {
 	tsp := &TreeSplitterParams{
 		SplitterParams: SplitterParams{
 			ChunkerParams: ChunkerParams{
@@ -156,7 +157,7 @@ func TreeSplit(data io.Reader, size int64, putter Putter) (k Address, wait func(
 		},
 		size: size,
 	}
-	return NewTreeSplitter(tsp).Split()
+	return NewTreeSplitter(tsp).Split(ctx)
 }
 
 func NewTreeJoiner(params *JoinerParams) *TreeChunker {
@@ -224,7 +225,7 @@ func (tc *TreeChunker) decrementWorkerCount() {
 	tc.workerCount -= 1
 }
 
-func (tc *TreeChunker) Split() (k Address, wait func(), err error) {
+func (tc *TreeChunker) Split(ctx context.Context) (k Address, wait func(context.Context) error, err error) {
 	if tc.chunkSize <= 0 {
 		panic("chunker must be initialised")
 	}
@@ -380,7 +381,7 @@ type LazyChunkReader struct {
 	getter    Getter
 }
 
-func (tc *TreeChunker) Join() *LazyChunkReader {
+func (tc *TreeChunker) Join(ctx context.Context) *LazyChunkReader {
 	return &LazyChunkReader{
 		key:       tc.addr,
 		chunkSize: tc.chunkSize,

swarm/storage/chunker_test.go

@@ -18,6 +18,7 @@ package storage
 
 import (
 	"bytes"
+	"context"
 	"crypto/rand"
 	"encoding/binary"
 	"errors"
@@ -81,7 +82,7 @@ func testRandomBrokenData(n int, tester *chunkerTester) {
 	putGetter := newTestHasherStore(NewMapChunkStore(), SHA3Hash)
 
 	expectedError := fmt.Errorf("Broken reader")
-	addr, _, err := TreeSplit(brokendata, int64(n), putGetter)
+	addr, _, err := TreeSplit(context.TODO(), brokendata, int64(n), putGetter)
 	if err == nil || err.Error() != expectedError.Error() {
 		tester.t.Fatalf("Not receiving the correct error! Expected %v, received %v", expectedError, err)
 	}
@@ -104,20 +105,24 @@ func testRandomData(usePyramid bool, hash string, n int, tester *chunkerTester)
 	putGetter := newTestHasherStore(NewMapChunkStore(), hash)
 
 	var addr Address
-	var wait func()
+	var wait func(context.Context) error
 	var err error
+	ctx := context.TODO()
 	if usePyramid {
-		addr, wait, err = PyramidSplit(data, putGetter, putGetter)
+		addr, wait, err = PyramidSplit(ctx, data, putGetter, putGetter)
 	} else {
-		addr, wait, err = TreeSplit(data, int64(n), putGetter)
+		addr, wait, err = TreeSplit(ctx, data, int64(n), putGetter)
 	}
 	if err != nil {
 		tester.t.Fatalf(err.Error())
 	}
 	tester.t.Logf(" Key = %v\n", addr)
-	wait()
+	err = wait(ctx)
+	if err != nil {
+		tester.t.Fatalf(err.Error())
+	}
 
-	reader := TreeJoin(addr, putGetter, 0)
+	reader := TreeJoin(context.TODO(), addr, putGetter, 0)
 	output := make([]byte, n)
 	r, err := reader.Read(output)
 	if r != n || err != io.EOF {
@@ -200,11 +205,15 @@ func TestDataAppend(t *testing.T) {
 		chunkStore := NewMapChunkStore()
 		putGetter := newTestHasherStore(chunkStore, SHA3Hash)
 
-		addr, wait, err := PyramidSplit(data, putGetter, putGetter)
+		ctx := context.TODO()
+		addr, wait, err := PyramidSplit(ctx, data, putGetter, putGetter)
+		if err != nil {
+			tester.t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			tester.t.Fatalf(err.Error())
 		}
-		wait()
 
 		//create a append data stream
 		appendInput, found := tester.inputs[uint64(m)]
@@ -217,13 +226,16 @@ func TestDataAppend(t *testing.T) {
 		}
 
 		putGetter = newTestHasherStore(chunkStore, SHA3Hash)
-		newAddr, wait, err := PyramidAppend(addr, appendData, putGetter, putGetter)
+		newAddr, wait, err := PyramidAppend(ctx, addr, appendData, putGetter, putGetter)
+		if err != nil {
+			tester.t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			tester.t.Fatalf(err.Error())
 		}
-		wait()
 
-		reader := TreeJoin(newAddr, putGetter, 0)
+		reader := TreeJoin(ctx, newAddr, putGetter, 0)
 		newOutput := make([]byte, n+m)
 		r, err := reader.Read(newOutput)
 		if r != (n + m) {
@@ -282,12 +294,16 @@ func benchmarkSplitJoin(n int, t *testing.B) {
 		data := testDataReader(n)
 
 		putGetter := newTestHasherStore(NewMapChunkStore(), SHA3Hash)
-		key, wait, err := PyramidSplit(data, putGetter, putGetter)
+		ctx := context.TODO()
+		key, wait, err := PyramidSplit(ctx, data, putGetter, putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
-		wait()
+		err = wait(ctx)
-		reader := TreeJoin(key, putGetter, 0)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		reader := TreeJoin(ctx, key, putGetter, 0)
 		benchReadAll(reader)
 	}
 }
@@ -298,7 +314,7 @@ func benchmarkSplitTreeSHA3(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, SHA3Hash)
 
-		_, _, err := TreeSplit(data, int64(n), putGetter)
+		_, _, err := TreeSplit(context.TODO(), data, int64(n), putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -311,7 +327,7 @@ func benchmarkSplitTreeBMT(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, BMTHash)
 
-		_, _, err := TreeSplit(data, int64(n), putGetter)
+		_, _, err := TreeSplit(context.TODO(), data, int64(n), putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -324,7 +340,7 @@ func benchmarkSplitPyramidSHA3(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, SHA3Hash)
 
-		_, _, err := PyramidSplit(data, putGetter, putGetter)
+		_, _, err := PyramidSplit(context.TODO(), data, putGetter, putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -338,7 +354,7 @@ func benchmarkSplitPyramidBMT(n int, t *testing.B) {
 		data := testDataReader(n)
 		putGetter := newTestHasherStore(&fakeChunkStore{}, BMTHash)
 
-		_, _, err := PyramidSplit(data, putGetter, putGetter)
+		_, _, err := PyramidSplit(context.TODO(), data, putGetter, putGetter)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
@@ -354,18 +370,25 @@ func benchmarkSplitAppendPyramid(n, m int, t *testing.B) {
 		chunkStore := NewMapChunkStore()
 		putGetter := newTestHasherStore(chunkStore, SHA3Hash)
 
-		key, wait, err := PyramidSplit(data, putGetter, putGetter)
+		ctx := context.TODO()
+		key, wait, err := PyramidSplit(ctx, data, putGetter, putGetter)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
-		wait()
 
 		putGetter = newTestHasherStore(chunkStore, SHA3Hash)
-		_, wait, err = PyramidAppend(key, data1, putGetter, putGetter)
+		_, wait, err = PyramidAppend(ctx, key, data1, putGetter, putGetter)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		err = wait(ctx)
 		if err != nil {
 			t.Fatalf(err.Error())
 		}
-		wait()
 	}
 }
 

swarm/storage/filestore.go

@@ -82,15 +82,15 @@ func NewFileStore(store ChunkStore, params *FileStoreParams) *FileStore {
 func (f *FileStore) Retrieve(ctx context.Context, addr Address) (reader *LazyChunkReader, isEncrypted bool) {
 	isEncrypted = len(addr) > f.hashFunc().Size()
 	getter := NewHasherStore(f.ChunkStore, f.hashFunc, isEncrypted)
-	reader = TreeJoin(addr, getter, 0)
+	reader = TreeJoin(ctx, addr, getter, 0)
 	return
 }
 
 // Public API. Main entry point for document storage directly. Used by the
 // FS-aware API and httpaccess
-func (f *FileStore) Store(ctx context.Context, data io.Reader, size int64, toEncrypt bool) (addr Address, wait func(), err error) {
+func (f *FileStore) Store(ctx context.Context, data io.Reader, size int64, toEncrypt bool) (addr Address, wait func(context.Context) error, err error) {
 	putter := NewHasherStore(f.ChunkStore, f.hashFunc, toEncrypt)
-	return PyramidSplit(data, putter, putter)
+	return PyramidSplit(ctx, data, putter, putter)
 }
 
 func (f *FileStore) HashSize() int {

swarm/storage/filestore_test.go

@@ -50,11 +50,15 @@ func testFileStoreRandom(toEncrypt bool, t *testing.T) {
 	defer os.RemoveAll("/tmp/bzz")
 
 	reader, slice := generateRandomData(testDataSize)
-	key, wait, err := fileStore.Store(context.TODO(), reader, testDataSize, toEncrypt)
+	ctx := context.TODO()
+	key, wait, err := fileStore.Store(ctx, reader, testDataSize, toEncrypt)
 	if err != nil {
 		t.Errorf("Store error: %v", err)
 	}
-	wait()
+	err = wait(ctx)
+	if err != nil {
+		t.Fatalf("Store waitt error: %v", err.Error())
+	}
 	resultReader, isEncrypted := fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)
@@ -111,11 +115,15 @@ func testFileStoreCapacity(toEncrypt bool, t *testing.T) {
 	}
 	fileStore := NewFileStore(localStore, NewFileStoreParams())
 	reader, slice := generateRandomData(testDataSize)
-	key, wait, err := fileStore.Store(context.TODO(), reader, testDataSize, toEncrypt)
+	ctx := context.TODO()
+	key, wait, err := fileStore.Store(ctx, reader, testDataSize, toEncrypt)
+	if err != nil {
+		t.Errorf("Store error: %v", err)
+	}
+	err = wait(ctx)
 	if err != nil {
 		t.Errorf("Store error: %v", err)
 	}
-	wait()
 	resultReader, isEncrypted := fileStore.Retrieve(context.TODO(), key)
 	if isEncrypted != toEncrypt {
 		t.Fatalf("isEncrypted expected %v got %v", toEncrypt, isEncrypted)

swarm/storage/hasherstore.go

@@ -17,6 +17,7 @@
 package storage
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
@@ -126,9 +127,10 @@ func (h *hasherStore) Close() {
 // Wait returns when
 //    1) the Close() function has been called and
 //    2) all the chunks which has been Put has been stored
-func (h *hasherStore) Wait() {
+func (h *hasherStore) Wait(ctx context.Context) error {
 	<-h.closed
 	h.wg.Wait()
+	return nil
 }
 
 func (h *hasherStore) createHash(chunkData ChunkData) Address {

swarm/storage/hasherstore_test.go

@@ -18,6 +18,7 @@ package storage
 
 import (
 	"bytes"
+	"context"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/swarm/storage/encryption"
@@ -60,7 +61,10 @@ func TestHasherStore(t *testing.T) {
 		hasherStore.Close()
 
 		// Wait until chunks are really stored
-		hasherStore.Wait()
+		err = hasherStore.Wait(context.TODO())
+		if err != nil {
+			t.Fatalf("Expected no error got \"%v\"", err)
+		}
 
 		// Get the first chunk
 		retrievedChunkData1, err := hasherStore.Get(key1)

swarm/storage/pyramid.go

@@ -17,6 +17,7 @@
 package storage
 
 import (
+	"context"
 	"encoding/binary"
 	"errors"
 	"io"
@@ -99,12 +100,12 @@ func NewPyramidSplitterParams(addr Address, reader io.Reader, putter Putter, get
 	When splitting, data is given as a SectionReader, and the key is a hashSize long byte slice (Key), the root hash of the entire content will fill this once processing finishes.
 	New chunks to store are store using the putter which the caller provides.
 */
-func PyramidSplit(reader io.Reader, putter Putter, getter Getter) (Address, func(), error) {
+func PyramidSplit(ctx context.Context, reader io.Reader, putter Putter, getter Getter) (Address, func(context.Context) error, error) {
-	return NewPyramidSplitter(NewPyramidSplitterParams(nil, reader, putter, getter, DefaultChunkSize)).Split()
+	return NewPyramidSplitter(NewPyramidSplitterParams(nil, reader, putter, getter, DefaultChunkSize)).Split(ctx)
 }
 
-func PyramidAppend(addr Address, reader io.Reader, putter Putter, getter Getter) (Address, func(), error) {
+func PyramidAppend(ctx context.Context, addr Address, reader io.Reader, putter Putter, getter Getter) (Address, func(context.Context) error, error) {
-	return NewPyramidSplitter(NewPyramidSplitterParams(addr, reader, putter, getter, DefaultChunkSize)).Append()
+	return NewPyramidSplitter(NewPyramidSplitterParams(addr, reader, putter, getter, DefaultChunkSize)).Append(ctx)
 }
 
 // Entry to create a tree node
@@ -203,7 +204,7 @@ func (pc *PyramidChunker) decrementWorkerCount() {
 	pc.workerCount -= 1
 }
 
-func (pc *PyramidChunker) Split() (k Address, wait func(), err error) {
+func (pc *PyramidChunker) Split(ctx context.Context) (k Address, wait func(context.Context) error, err error) {
 	log.Debug("pyramid.chunker: Split()")
 
 	pc.wg.Add(1)
@@ -235,7 +236,7 @@ func (pc *PyramidChunker) Split() (k Address, wait func(), err error) {
 
 }
 
-func (pc *PyramidChunker) Append() (k Address, wait func(), err error) {
+func (pc *PyramidChunker) Append(ctx context.Context) (k Address, wait func(context.Context) error, err error) {
 	log.Debug("pyramid.chunker: Append()")
 	// Load the right most unfinished tree chunks in every level
 	pc.loadTree()

swarm/storage/types.go

@@ -18,6 +18,7 @@ package storage
 
 import (
 	"bytes"
+	"context"
 	"crypto"
 	"crypto/rand"
 	"encoding/binary"
@@ -303,7 +304,7 @@ type Putter interface {
 	// Close is to indicate that no more chunk data will be Put on this Putter
 	Close()
 	// Wait returns if all data has been store and the Close() was called.
-	Wait()
+	Wait(context.Context) error
 }
 
 // Getter is an interface to retrieve a chunk's data by its reference

swarm/swarm_test.go

@@ -348,12 +348,16 @@ func testLocalStoreAndRetrieve(t *testing.T, swarm *Swarm, n int, randomData boo
 	}
 	dataPut := string(slice)
 
-	k, wait, err := swarm.api.Store(context.TODO(), strings.NewReader(dataPut), int64(len(dataPut)), false)
+	ctx := context.TODO()
+	k, wait, err := swarm.api.Store(ctx, strings.NewReader(dataPut), int64(len(dataPut)), false)
 	if err != nil {
 		t.Fatal(err)
 	}
 	if wait != nil {
-		wait()
+		err = wait(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
 	}
 
 	r, _ := swarm.api.Retrieve(context.TODO(), k)