projecteru2 · CMGS · Sep 7, 2023 · Jan 7, 2022 · Jan 16, 2022 · Jan 18, 2022
diff --git a/3rdmocks/ServerStream.go b/3rdmocks/ServerStream.go
diff --git a/cluster/calcium/replace_test.go b/cluster/calcium/replace_test.go
@@ -184,7 +184,7 @@ func TestReplaceWorkload(t *testing.T) {
 
 	engine.On("VirtualizationCreate", mock.Anything, mock.Anything).Return(&enginetypes.VirtualizationCreated{ID: "new"}, nil)
 	engine.On("VirtualizationStart", mock.Anything, mock.Anything).Return(nil)
-	engine.On("VirtualizationCopyTo", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	engine.On("VirtualizationCopyChunkTo", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	engine.On("VirtualizationInspect", mock.Anything, mock.Anything).Return(&enginetypes.VirtualizationInfo{User: "test"}, nil)
 	store.On("AddWorkload", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	// failed by remove workload

diff --git a/cluster/calcium/send.go b/cluster/calcium/send.go
@@ -1,6 +1,7 @@
 package calcium
 
 import (
+	"bytes"
 	"context"
 	"sync"
 
@@ -48,5 +49,5 @@ func (c *Calcium) Send(ctx context.Context, opts *types.SendOptions) (chan *type
 
 func (c *Calcium) doSendFileToWorkload(ctx context.Context, engine engine.API, ID string, file types.LinuxFile) error {
 	log.WithFunc("calcium.doSendFileToWorkload").Infof(ctx, "Send file to %s:%s", ID, file.Filename)
-	return engine.VirtualizationCopyTo(ctx, ID, file.Filename, file.Clone().Content, file.UID, file.GID, file.Mode)
+	return engine.VirtualizationCopyChunkTo(ctx, ID, file.Filename, int64(len(file.Content)), bytes.NewReader(file.Clone().Content), file.UID, file.GID, file.Mode)
 }
diff --git a/cluster/calcium/send_test.go b/cluster/calcium/send_test.go
@@ -19,6 +19,7 @@ func TestSend(t *testing.T) {
 	c := NewTestCluster()
 	ctx := context.Background()
 
+	// 这部分是在测试参数合法性
 	// failed by validating
 	_, err := c.Send(ctx, &types.SendOptions{IDs: []string{}, Files: []types.LinuxFile{{Content: []byte("xxx")}}})
 	assert.Error(t, err)
@@ -57,21 +58,21 @@ func TestSend(t *testing.T) {
 	// failed by engine
 	content, _ := io.ReadAll(tmpfile)
 	opts.Files[0].Content = content
-	engine.On("VirtualizationCopyTo",
+	engine.On("VirtualizationCopyChunkTo",
 		mock.Anything, mock.Anything, mock.Anything,
 		mock.Anything, mock.Anything, mock.Anything,
-		mock.Anything,
+		mock.Anything, mock.Anything,
 	).Return(types.ErrMockError).Once()
 	ch, err = c.Send(ctx, opts)
 	assert.NoError(t, err)
 	for r := range ch {
 		assert.Error(t, r.Error)
 	}
 	// success
-	engine.On("VirtualizationCopyTo",
+	engine.On("VirtualizationCopyChunkTo",
 		mock.Anything, mock.Anything, mock.Anything,
 		mock.Anything, mock.Anything, mock.Anything,
-		mock.Anything,
+		mock.Anything, mock.Anything,
 	).Return(nil)
 	ch, err = c.Send(ctx, opts)
 	assert.NoError(t, err)

diff --git a/cluster/calcium/sendlarge.go b/cluster/calcium/sendlarge.go
@@ -0,0 +1,104 @@
+package calcium
+
+import (
+	"context"
+	"io"
+	"sync"
+
+	"github.com/pkg/errors"
+	"github.com/projecteru2/core/log"
+	"github.com/projecteru2/core/types"
+	"github.com/projecteru2/core/utils"
+)
+
+// SendLargeFile send large files by stream to workload
+func (c *Calcium) SendLargeFile(ctx context.Context, inputChan chan *types.SendLargeFileOptions) chan *types.SendMessage {
+	resp := make(chan *types.SendMessage)
+	wg := &sync.WaitGroup{}
+	utils.SentryGo(func() {
+		defer close(resp)
+		senders := make(map[string]*workloadSender)
+		// for each file
+		for data := range inputChan {
+			for _, id := range data.Ids {
+				if _, ok := senders[id]; !ok {
+					log.Debugf(ctx, "[SendLargeFile] create sender for %s", id)
+					// for each container, let's create a new sender to send identical file chunk, each chunk will include the metadata of this file
+					// we need to add `waitGroup out of the `newWorkloadSender` because we need to avoid `wg.Wait()` be executing before `wg.Add()`,
+					// which will cause the goroutine in `c.newWorkloadSender` to be killed.
+					wg.Add(1)
+					sender := c.newWorkloadSender(ctx, id, resp, wg)
+					senders[id] = sender
+				}
+				senders[id].send(data)
+			}
+		}
+		for _, sender := range senders {
+			sender.close()
+		}
+		wg.Wait()
+	})
+	return resp
+}
+
+type workloadSender struct {
+	calcium *Calcium
+	id      string
+	wg      *sync.WaitGroup
+	buffer  chan *types.SendLargeFileOptions
+	resp    chan *types.SendMessage
+}
+
+func (c *Calcium) newWorkloadSender(ctx context.Context, ID string, resp chan *types.SendMessage, wg *sync.WaitGroup) *workloadSender {
+	sender := &workloadSender{
+		calcium: c,
+		id:      ID,
+		wg:      wg,
+		buffer:  make(chan *types.SendLargeFileOptions, 10),
+		resp:    resp,
+	}
+	utils.SentryGo(func() {
+		var writer *io.PipeWriter
+		curFile := ""
+		for data := range sender.buffer {
+			if curFile != "" && curFile != data.Dst {
+				log.Warnf(ctx, "[newWorkloadExecutor] receive different files %s, %s", curFile, data.Dst)
+				break
+			}
+			// ready to send
+			if curFile == "" {
+				log.Debugf(ctx, "[newWorkloadExecutor]Receive new file %s to %s", curFile, sender.id)
+				curFile = data.Dst
+				pr, pw := io.Pipe()
+				writer = pw
+				utils.SentryGo(func(ID, name string, size int64, content io.Reader, uid, gid int, mode int64) func() {
+					return func() {
+						defer wg.Done()
+						if err := sender.calcium.withWorkloadLocked(ctx, ID, func(ctx context.Context, workload *types.Workload) error {
+							err := errors.WithStack(workload.Engine.VirtualizationCopyChunkTo(ctx, ID, name, size, content, uid, gid, mode))
+							resp <- &types.SendMessage{ID: ID, Path: name, Error: err}
+							return nil
+						}); err != nil {
+							resp <- &types.SendMessage{ID: ID, Error: err}
+						}
+					}
+				}(ID, curFile, data.Size, pr, data.UID, data.GID, data.Mode))
+			}
+			n, err := writer.Write(data.Chunk)
+			if err != nil || n != len(data.Chunk) {
+				log.Errorf(ctx, err, "[newWorkloadExecutor] send file to engine err, file = %s", curFile)
+				break
+			}
+		}
+		writer.Close()
+	})
+	return sender
+}
+
+func (s *workloadSender) send(chunk *types.SendLargeFileOptions) {
+	s.buffer <- chunk
+}
+
+func (s *workloadSender) close() {
+	close(s.buffer)
+}
diff --git a/cluster/calcium/sendlarge_test.go b/cluster/calcium/sendlarge_test.go
@@ -0,0 +1,86 @@
+package calcium
+
+import (
+	"context"
+	"io/ioutil"
+	"os"
+	"testing"
+
+	enginemocks "github.com/projecteru2/core/engine/mocks"
+	lockmocks "github.com/projecteru2/core/lock/mocks"
+	storemocks "github.com/projecteru2/core/store/mocks"
+	"github.com/projecteru2/core/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestSendLarge(t *testing.T) {
+	c := NewTestCluster()
+	ctx := context.Background()
+
+	tmpfile, err := os.CreateTemp("", "example")
+	assert.NoError(t, err)
+	defer os.RemoveAll(tmpfile.Name())
+	defer tmpfile.Close()
+	opts := &types.SendLargeFileOptions{
+		Ids:   []string{"cid"},
+		Size:  1,
+		Dst:   "/tmp/1",
+		Chunk: []byte{},
+	}
+	optsChan := make(chan *types.SendLargeFileOptions)
+	store := &storemocks.Store{}
+	c.store = store
+	lock := &lockmocks.DistributedLock{}
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
+	lock.On("Unlock", mock.Anything).Return(nil)
+	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
+	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
+	ch := c.SendLargeFile(ctx, optsChan)
+	go func() {
+		optsChan <- opts
+		close(optsChan)
+	}()
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	engine := &enginemocks.API{}
+	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(
+		[]*types.Workload{{ID: "cid", Engine: engine}}, nil,
+	)
+	// failed by engine
+	content, _ := ioutil.ReadAll(tmpfile)
+	opts.Chunk = content
+	engine.On("VirtualizationCopyChunkTo",
+		mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything, mock.Anything,
+	).Return(types.ErrMockError).Once()
+	optsChan = make(chan *types.SendLargeFileOptions)
+	ch = c.SendLargeFile(ctx, optsChan)
+	go func() {
+		optsChan <- opts
+		close(optsChan)
+	}()
+	for r := range ch {
+		t.Log(r.Error)
+		assert.Error(t, r.Error)
+	}
+	// success
+	engine.On("VirtualizationCopyChunkTo",
+		mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything, mock.Anything,
+	).Return(nil)
+	optsChan = make(chan *types.SendLargeFileOptions)
+	ch = c.SendLargeFile(ctx, optsChan)
+	go func() {
+		optsChan <- opts
+		close(optsChan)
+	}()
+	for r := range ch {
+		assert.Equal(t, r.ID, "cid")
+		assert.Equal(t, r.Path, "/tmp/1")
+		assert.NoError(t, r.Error)
+	}
+}
diff --git a/cluster/cluster.go b/cluster/cluster.go
@@ -80,6 +80,7 @@ type Cluster interface {
 	// file methods
 	Copy(ctx context.Context, opts *types.CopyOptions) (chan *types.CopyMessage, error)
 	Send(ctx context.Context, opts *types.SendOptions) (chan *types.SendMessage, error)
+	SendLargeFile(ctx context.Context, opts chan *types.SendLargeFileOptions) chan *types.SendMessage
 	// image methods
 	BuildImage(ctx context.Context, opts *types.BuildOptions) (chan *types.BuildImageMessage, error)
 	CacheImage(ctx context.Context, opts *types.ImageOptions) (chan *types.CacheImageMessage, error)