Added wait group that waits until local storages finish

bzz/chunker.go

@@ -58,10 +58,11 @@ type Chunker interface {
 	/*
 	 	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 coming to caller via the chunk storage channel, which the caller provides.
+	 	wg is a Waitgroup (can be nil) that can be used to block until the local storage finishes
 	 	The caller gets returned an error channel, if an error is encountered during splitting, it is fed to errC error channel.
 	   A closed error signals process completion at which point the key can be considered final if there were no errors.
 	*/
-	Split(key Key, data SectionReader, chunkC chan *Chunk) chan error
+	Split(key Key, data SectionReader, chunkC chan *Chunk, wg *sync.WaitGroup) chan error
 	/*
 		Join reconstructs original content based on a root key.
 		When joining, the caller gets returned a Lazy SectionReader
@@ -133,7 +134,12 @@ func (self *TreeChunker) Hash(size int64, input []byte) []byte {
 	return hasher.Sum(nil)
 }
 
-func (self *TreeChunker) Split(key Key, data SectionReader, chunkC chan *Chunk) (errC chan error) {
+func (self *TreeChunker) Split(key Key, data SectionReader, chunkC chan *Chunk, swg *sync.WaitGroup) (errC chan error) {
+
+	if swg != nil {
+		swg.Add(1)
+		defer swg.Done()
+	}
 
 	if self.chunkSize <= 0 {
 		panic("chunker must be initialised")
@@ -164,7 +170,7 @@ func (self *TreeChunker) Split(key Key, data SectionReader, chunkC chan *Chunk)
 		// dpaLogger.Debugf("split request received for data (%v bytes, depth: %v)", size, depth)
 
 		//launch actual recursive function passing the workgroup
-		self.split(depth, treeSize/self.Branches, key, data, chunkC, rerrC, wg)
+		self.split(depth, treeSize/self.Branches, key, data, chunkC, rerrC, wg, swg)
 	}()
 
 	// closes internal error channel if all subprocesses in the workgroup finished
@@ -190,7 +196,7 @@ func (self *TreeChunker) Split(key Key, data SectionReader, chunkC chan *Chunk)
 	return
 }
 
-func (self *TreeChunker) split(depth int, treeSize int64, key Key, data SectionReader, chunkC chan *Chunk, errc chan error, parentWg *sync.WaitGroup) {
+func (self *TreeChunker) split(depth int, treeSize int64, key Key, data SectionReader, chunkC chan *Chunk, errc chan error, parentWg *sync.WaitGroup, swg *sync.WaitGroup) {
 
 	defer parentWg.Done()
 
@@ -238,7 +244,7 @@ func (self *TreeChunker) split(depth int, treeSize int64, key Key, data SectionR
 			subTreeKey := chunk[i*self.hashSize : (i+1)*self.hashSize]
 
 			childrenWg.Add(1)
-			go self.split(depth-1, treeSize/self.Branches, subTreeKey, subTreeData, chunkC, errc, childrenWg)
+			go self.split(depth-1, treeSize/self.Branches, subTreeKey, subTreeData, chunkC, errc, childrenWg, swg)
 
 			i++
 			pos += treeSize
@@ -255,6 +261,11 @@ func (self *TreeChunker) split(depth int, treeSize int64, key Key, data SectionR
 			Key:  hash,
 			Data: chunkData,
 			Size: size,
+			wg:   swg,
+		}
+
+		if swg != nil {
+			swg.Add(1)
 		}
 	}
 	// send off new chunk to storage

bzz/chunker_test.go

@@ -2,7 +2,7 @@ package bzz
 
 import (
 	"bytes"
-	"fmt"
+	// "fmt"
 	"io"
 	"testing"
 	"time"
@@ -37,7 +37,7 @@ func (self *chunkerTester) Split(chunker *TreeChunker, l int) (key Key, input []
 	input = slice
 	key = make([]byte, 32)
 	chunkC := make(chan *Chunk, 1000)
-	errC := chunker.Split(key, data, chunkC)
+	errC := chunker.Split(key, data, chunkC, nil)
 	quitC := make(chan bool)
 	timeout := time.After(600 * time.Second)
 
@@ -121,19 +121,17 @@ func (self *chunkerTester) Join(chunker *TreeChunker, key Key, c int) SectionRea
 
 func testRandomData(chunker *TreeChunker, tester *chunkerTester, n int, chunks int, t *testing.T) {
 	key, input := tester.Split(chunker, n)
-	fmt.Printf("split returned\n")
-	fmt.Printf("input\n%x\n", input)
 
 	tester.checkChunks(t, chunks)
-	t.Logf("chunks: %v", tester.chunks)
+	time.Sleep(100 * time.Millisecond)
-	fmt.Printf("chunks: %v", tester.chunks)
+
 	reader := tester.Join(chunker, key, 0)
 	output := make([]byte, n)
 	_, err := reader.Read(output)
 	if err != io.EOF {
 		t.Errorf("read error %v\n", err)
 	}
-	t.Logf(" IN: %x\nOUT: %x\n", input, output)
+	// t.Logf(" IN: %x\nOUT: %x\n", input, output)
 	if !bytes.Equal(output, input) {
 		t.Errorf("input and output mismatch\n IN: %x\nOUT: %x\n", input, output)
 	}
@@ -151,7 +149,7 @@ func TestRandomData(t *testing.T) {
 	testRandomData(chunker, tester, 70, 3, t)
 	testRandomData(chunker, tester, 179, 5, t)
 	testRandomData(chunker, tester, 253, 7, t)
-	t.Logf("chunks %v", tester.chunks)
+	// t.Logf("chunks %v", tester.chunks)
 }
 
 func chunkerAndTester() (chunker *TreeChunker, tester *chunkerTester) {

bzz/common_test.go

@@ -3,6 +3,7 @@ package bzz
 import (
 	"crypto/rand"
 	"io"
+	"sync"
 	"testing"
 )
 
@@ -26,7 +27,9 @@ func randomChunks(l int64, branches int64, chunkC chan *Chunk) (key Key, errC ch
 	if err != nil {
 		panic("no rand")
 	}
-	errC = chunker.Split(key, NewChunkReaderFromBytes(b), chunkC)
+	wg := &sync.WaitGroup{}
+	errC = chunker.Split(key, NewChunkReaderFromBytes(b), chunkC, wg)
+	wg.Wait()
 	return
 }
 

bzz/dpa.go

@@ -46,7 +46,7 @@ type DPA struct {
 
 	lock    sync.Mutex
 	running bool
-	wg      sync.WaitGroup
+	wg      *sync.WaitGroup
 	quitC   chan bool
 }
 
@@ -61,6 +61,7 @@ type Chunk struct {
 	Key  Key            // always
 	C    chan bool      // to signal data delivery by the dpa
 	req  *requestStatus //
+	wg   *sync.WaitGroup
 }
 
 type ChunkStore interface {
@@ -74,9 +75,9 @@ func (self *DPA) Retrieve(key Key) SectionReader {
 	// we can add subscriptions etc. or timeout here
 }
 
-func (self *DPA) Store(data SectionReader) (key Key, err error) {
+func (self *DPA) Store(data SectionReader, wg *sync.WaitGroup) (key Key, err error) {
 	key = make([]byte, self.Chunker.KeySize())
-	errC := self.Chunker.Split(key, data, self.storeC)
+	errC := self.Chunker.Split(key, data, self.storeC, wg)
 
 SPLIT:
 	for {
@@ -154,6 +155,9 @@ func (self *DPA) storeLoop() {
 		for ch := range self.storeC {
 			// go func(chunk *Chunk) {
 			self.ChunkStore.Put(ch)
+			if ch.wg != nil {
+				ch.wg.Done()
+			}
 			// self.ChunkStore.Put(chunk)
 			// }(ch)
 			select {

bzz/dpa_test.go

@@ -6,8 +6,8 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"sync"
 	"testing"
-	// "time"
 )
 
 const testDataSize = 0x1000000
@@ -33,10 +33,12 @@ func TestDPArandom(t *testing.T) {
 	}
 	dpa.Start()
 	reader, slice := testDataReader(testDataSize)
-	key, err := dpa.Store(reader)
+	wg := &sync.WaitGroup{}
+	key, err := dpa.Store(reader, wg)
 	if err != nil {
 		t.Errorf("Store error: %v", err)
 	}
+	wg.Wait()
 	resultReader := dpa.Retrieve(key)
 	resultSlice := make([]byte, len(slice))
 	n, err := resultReader.ReadAt(resultSlice, 0)
@@ -46,9 +48,9 @@ func TestDPArandom(t *testing.T) {
 	if n != len(slice) {
 		t.Errorf("Slice size error got %d, expected %d.", n, len(slice))
 	}
-	// if !bytes.Equal(slice, resultSlice) {
+	if !bytes.Equal(slice, resultSlice) {
-	// 	t.Errorf("Comparison error.")
+		t.Errorf("Comparison error.")
-	// }
+	}
 	ioutil.WriteFile("/tmp/slice.bzz.16M", slice, 0666)
 	ioutil.WriteFile("/tmp/result.bzz.16M", resultSlice, 0666)
 	localStore.memStore = newMemStore(dbStore)
@@ -89,10 +91,12 @@ func TestDPA_capacity(t *testing.T) {
 	}
 	dpa.Start()
 	reader, slice := testDataReader(testDataSize)
-	key, err := dpa.Store(reader)
+	wg := &sync.WaitGroup{}
+	key, err := dpa.Store(reader, wg)
 	if err != nil {
 		t.Errorf("Store error: %v", err)
 	}
+	wg.Wait()
 	resultReader := dpa.Retrieve(key)
 	resultSlice := make([]byte, len(slice))
 	n, err := resultReader.ReadAt(resultSlice, 0)

bzz/httpaccess.go

@@ -98,7 +98,7 @@ func handler(w http.ResponseWriter, r *http.Request, dpa *DPA) {
 			key, err := dpa.Store(io.NewSectionReader(&sequentialReader{
 				reader: r.Body,
 				ahead:  make(map[int64]chan bool),
-			}, 0, r.ContentLength))
+			}, 0, r.ContentLength), nil)
 			if err == nil {
 				fmt.Fprintf(w, "%064x", key)
 				dpaLogger.Debugf("Swarm: Object %064x stored", key)

bzz/localstore.go

@@ -10,7 +10,15 @@ type localStore struct {
 // its integrity is checked ?
 func (self *localStore) Put(chunk *Chunk) {
 	self.memStore.Put(chunk)
-	go self.dbStore.Put(chunk)
+	if chunk.wg != nil {
+		chunk.wg.Add(1)
+	}
+	go func() {
+		self.dbStore.Put(chunk)
+		if chunk.wg != nil {
+			chunk.wg.Done()
+		}
+	}()
 }
 
 // Get(chunk *Chunk) looks up a chunk in the local stores