introduce ReaderManager to address race condition in smaller critical

Author: jschwinger233

cluster/calcium/create.go

@@ -3,7 +3,6 @@ package calcium
 import (
 	"context"
 	"fmt"
-	"io"
 	"sync"
 
 	"github.com/projecteru2/core/cluster"
@@ -224,16 +223,12 @@ func (c *Calcium) doCreateAndStartContainer(
 
 			// Copy data to container
 			if len(opts.Data) > 0 {
-				for dst, src := range opts.Data {
-					mutexSend := func() error {
-						opts.Mux.Lock()
-						defer opts.Mux.Unlock()
-						if _, err = src.Seek(0, io.SeekStart); err != nil {
-							return err
-						}
-						return c.doSendFileToContainer(ctx, node.Engine, container.ID, dst, src, true, true)
+				for dst, readerManager := range opts.Data {
+					reader, err := readerManager.GetReader()
+					if err != nil {
+						return err
 					}
-					if err = mutexSend(); err != nil {
+					if err = c.doSendFileToContainer(ctx, node.Engine, container.ID, dst, reader, true, true); err != nil {
 						return err
 					}
 				}

cluster/calcium/replace.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"sync"
 
 	"github.com/projecteru2/core/store"
@@ -121,11 +120,9 @@ func (c *Calcium) doReplaceContainer(
 		if err != nil {
 			return nil, removeMessage, err
 		}
-		bs, err := ioutil.ReadAll(stream)
-		if err != nil {
+		if opts.DeployOptions.Data[dst], err = types.NewReaderManager(stream); err != nil {
 			return nil, removeMessage, err
 		}
-		opts.DeployOptions.Data[dst] = bytes.NewReader(bs)
 	}
 
 	createMessage := &types.CreateContainerMessage{}

rpc/transform.go

@@ -262,9 +262,11 @@ func toCoreDeployOptions(d *pb.DeployOptions) (*types.DeployOptions, error) {
 		return nil, err
 	}
 
-	data := map[string]*bytes.Reader{}
+	data := map[string]types.ReaderManager{}
 	for filename, bs := range d.Data {
-		data[filename] = bytes.NewReader(bs)
+		if data[filename], err = types.NewReaderManager(bytes.NewBuffer(bs)); err != nil {
+			return nil, err
+		}
 	}
 
 	return &types.DeployOptions{

types/options.go

@@ -2,6 +2,8 @@ package types
 
 import (
 	"bytes"
+	"io"
+	"io/ioutil"
 	"sync"
 )
 
@@ -30,7 +32,7 @@ type DeployOptions struct {
 	Labels       map[string]string        // Labels for containers
 	NodeLabels   map[string]string        // NodeLabels for filter node
 	DeployMethod string                   // Deploy method
-	Data         map[string]*bytes.Reader // For additional file data
+	Data         map[string]ReaderManager // For additional file data
 	SoftLimit    bool                     // Soft limit memory
 	NodesLimit   int                      // Limit nodes count
 	ProcessIdent string                   // ProcessIdent ident this deploy
@@ -42,6 +44,32 @@ type DeployOptions struct {
 	Mux sync.Mutex // used for concurrent send during creation
 }
 
+// ReaderManager return Reader under concurrency
+type ReaderManager interface {
+	GetReader() (io.Reader, error)
+}
+
+type readerManager struct {
+	mux sync.Mutex
+	r   io.ReadSeeker
+}
+
+func (rm *readerManager) GetReader() (_ io.Reader, err error) {
+	rm.mux.Lock()
+	defer rm.mux.Unlock()
+	buf := &bytes.Buffer{}
+	_, err = io.Copy(buf, rm.r)
+	_, err = rm.r.Seek(0, io.SeekStart)
+	return buf, nil
+}
+
+func NewReaderManager(r io.Reader) (ReaderManager, error) {
+	bs, err := ioutil.ReadAll(r)
+	return &readerManager{
+		r: bytes.NewReader(bs),
+	}, err
+}
+
 // Normalize keeps deploy options consistant
 func (o *DeployOptions) Normalize() {
 	o.Storage += o.Volumes.TotalSize()