1 files changed, 68 insertions, 44 deletions
diff --git a/swarm/storage/hasherstore.go b/swarm/storage/hasherstore.go
index 766207eae..879622b9a 100644
--- a/swarm/storage/hasherstore.go
+++ b/swarm/storage/hasherstore.go
@@ -19,10 +19,10 @@ package storage
 import (
 	"context"
 	"fmt"
-	"sync"
+	"sync/atomic"
 
 	"github.com/ethereum/go-ethereum/crypto/sha3"
-	"github.com/ethereum/go-ethereum/swarm/chunk"
+	ch "github.com/ethereum/go-ethereum/swarm/chunk"
 	"github.com/ethereum/go-ethereum/swarm/storage/encryption"
 )
 
@@ -30,31 +30,36 @@ type hasherStore struct {
 	store     ChunkStore
 	toEncrypt bool
 	hashFunc  SwarmHasher
-	hashSize  int   // content hash size
-	refSize   int64 // reference size (content hash + possibly encryption key)
-	wg        *sync.WaitGroup
-	closed    chan struct{}
+	hashSize  int           // content hash size
+	refSize   int64         // reference size (content hash + possibly encryption key)
+	nrChunks  uint64        // number of chunks to store
+	errC      chan error    // global error channel
+	doneC     chan struct{} // closed by Close() call to indicate that count is the final number of chunks
+	quitC     chan struct{} // closed to quit unterminated routines
 }
 
 // NewHasherStore creates a hasherStore object, which implements Putter and Getter interfaces.
 // With the HasherStore you can put and get chunk data (which is just []byte) into a ChunkStore
 // and the hasherStore will take core of encryption/decryption of data if necessary
-func NewHasherStore(chunkStore ChunkStore, hashFunc SwarmHasher, toEncrypt bool) *hasherStore {
+func NewHasherStore(store ChunkStore, hashFunc SwarmHasher, toEncrypt bool) *hasherStore {
 	hashSize := hashFunc().Size()
 	refSize := int64(hashSize)
 	if toEncrypt {
 		refSize += encryption.KeyLength
 	}
 
-	return &hasherStore{
-		store:     chunkStore,
+	h := &hasherStore{
+		store:     store,
 		toEncrypt: toEncrypt,
 		hashFunc:  hashFunc,
 		hashSize:  hashSize,
 		refSize:   refSize,
-		wg:        &sync.WaitGroup{},
-		closed:    make(chan struct{}),
+		errC:      make(chan error),
+		doneC:     make(chan struct{}),
+		quitC:     make(chan struct{}),
 	}
+
+	return h
 }
 
 // Put stores the chunkData into the ChunkStore of the hasherStore and returns the reference.
@@ -62,7 +67,6 @@ func NewHasherStore(chunkStore ChunkStore, hashFunc SwarmHasher, toEncrypt bool)
 // Asynchronous function, the data will not necessarily be stored when it returns.
 func (h *hasherStore) Put(ctx context.Context, chunkData ChunkData) (Reference, error) {
 	c := chunkData
-	size := chunkData.Size()
 	var encryptionKey encryption.Key
 	if h.toEncrypt {
 		var err error
@@ -71,29 +75,28 @@ func (h *hasherStore) Put(ctx context.Context, chunkData ChunkData) (Reference,
 			return nil, err
 		}
 	}
-	chunk := h.createChunk(c, size)
-
+	chunk := h.createChunk(c)
 	h.storeChunk(ctx, chunk)
 
-	return Reference(append(chunk.Addr, encryptionKey...)), nil
+	return Reference(append(chunk.Address(), encryptionKey...)), nil
 }
 
 // Get returns data of the chunk with the given reference (retrieved from the ChunkStore of hasherStore).
 // If the data is encrypted and the reference contains an encryption key, it will be decrypted before
 // return.
 func (h *hasherStore) Get(ctx context.Context, ref Reference) (ChunkData, error) {
-	key, encryptionKey, err := parseReference(ref, h.hashSize)
+	addr, encryptionKey, err := parseReference(ref, h.hashSize)
 	if err != nil {
 		return nil, err
 	}
-	toDecrypt := (encryptionKey != nil)
 
-	chunk, err := h.store.Get(ctx, key)
+	chunk, err := h.store.Get(ctx, addr)
 	if err != nil {
 		return nil, err
 	}
 
-	chunkData := chunk.SData
+	chunkData := ChunkData(chunk.Data())
+	toDecrypt := (encryptionKey != nil)
 	if toDecrypt {
 		var err error
 		chunkData, err = h.decryptChunkData(chunkData, encryptionKey)
@@ -107,16 +110,40 @@ func (h *hasherStore) Get(ctx context.Context, ref Reference) (ChunkData, error)
 // Close indicates that no more chunks will be put with the hasherStore, so the Wait
 // function can return when all the previously put chunks has been stored.
 func (h *hasherStore) Close() {
-	close(h.closed)
+	close(h.doneC)
 }
 
 // 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(ctx context.Context) error {
-	<-h.closed
-	h.wg.Wait()
-	return nil
+	defer close(h.quitC)
+	var nrStoredChunks uint64 // number of stored chunks
+	var done bool
+	doneC := h.doneC
+	for {
+		select {
+		// if context is done earlier, just return with the error
+		case <-ctx.Done():
+			return ctx.Err()
+		// doneC is closed if all chunks have been submitted, from then we just wait until all of them are also stored
+		case <-doneC:
+			done = true
+			doneC = nil
+		// a chunk has been stored, if err is nil, then successfully, so increase the stored chunk counter
+		case err := <-h.errC:
+			if err != nil {
+				return err
+			}
+			nrStoredChunks++
+		}
+		// if all the chunks have been submitted and all of them are stored, then we can return
+		if done {
+			if nrStoredChunks >= atomic.LoadUint64(&h.nrChunks) {
+				return nil
+			}
+		}
+	}
 }
 
 func (h *hasherStore) createHash(chunkData ChunkData) Address {
@@ -126,12 +153,9 @@ func (h *hasherStore) createHash(chunkData ChunkData) Address {
 	return hasher.Sum(nil)
 }
 
-func (h *hasherStore) createChunk(chunkData ChunkData, chunkSize int64) *Chunk {
+func (h *hasherStore) createChunk(chunkData ChunkData) *chunk {
 	hash := h.createHash(chunkData)
-	chunk := NewChunk(hash, nil)
-	chunk.SData = chunkData
-	chunk.Size = chunkSize
-
+	chunk := NewChunk(hash, chunkData)
 	return chunk
 }
 
@@ -162,10 +186,10 @@ func (h *hasherStore) decryptChunkData(chunkData ChunkData, encryptionKey encryp
 
 	// removing extra bytes which were just added for padding
 	length := ChunkData(decryptedSpan).Size()
-	for length > chunk.DefaultSize {
-		length = length + (chunk.DefaultSize - 1)
-		length = length / chunk.DefaultSize
-		length *= h.refSize
+	for length > ch.DefaultSize {
+		length = length + (ch.DefaultSize - 1)
+		length = length / ch.DefaultSize
+		length *= uint64(h.refSize)
 	}
 
 	c := make(ChunkData, length+8)
@@ -205,32 +229,32 @@ func (h *hasherStore) decrypt(chunkData ChunkData, key encryption.Key) ([]byte,
 }
 
 func (h *hasherStore) newSpanEncryption(key encryption.Key) encryption.Encryption {
-	return encryption.New(key, 0, uint32(chunk.DefaultSize/h.refSize), sha3.NewKeccak256)
+	return encryption.New(key, 0, uint32(ch.DefaultSize/h.refSize), sha3.NewKeccak256)
 }
 
 func (h *hasherStore) newDataEncryption(key encryption.Key) encryption.Encryption {
-	return encryption.New(key, int(chunk.DefaultSize), 0, sha3.NewKeccak256)
+	return encryption.New(key, int(ch.DefaultSize), 0, sha3.NewKeccak256)
 }
 
-func (h *hasherStore) storeChunk(ctx context.Context, chunk *Chunk) {
-	h.wg.Add(1)
+func (h *hasherStore) storeChunk(ctx context.Context, chunk *chunk) {
+	atomic.AddUint64(&h.nrChunks, 1)
 	go func() {
-		<-chunk.dbStoredC
-		h.wg.Done()
+		select {
+		case h.errC <- h.store.Put(ctx, chunk):
+		case <-h.quitC:
+		}
 	}()
-	h.store.Put(ctx, chunk)
 }
 
 func parseReference(ref Reference, hashSize int) (Address, encryption.Key, error) {
-	encryptedKeyLength := hashSize + encryption.KeyLength
+	encryptedRefLength := hashSize + encryption.KeyLength
 	switch len(ref) {
-	case KeyLength:
+	case AddressLength:
 		return Address(ref), nil, nil
-	case encryptedKeyLength:
+	case encryptedRefLength:
 		encKeyIdx := len(ref) - encryption.KeyLength
 		return Address(ref[:encKeyIdx]), encryption.Key(ref[encKeyIdx:]), nil
 	default:
-		return nil, nil, fmt.Errorf("Invalid reference length, expected %v or %v got %v", hashSize, encryptedKeyLength, len(ref))
+		return nil, nil, fmt.Errorf("Invalid reference length, expected %v or %v got %v", hashSize, encryptedRefLength, len(ref))
 	}
-
 }