monorepo

Co-authored-by: Malte Poll <mp@edgeless.systems>
Co-authored-by: katexochen <katexochen@users.noreply.github.com>
Co-authored-by: Daniel Weiße <dw@edgeless.systems>
Co-authored-by: Thomas Tendyck <tt@edgeless.systems>
Co-authored-by: Benedict Schlueter <bs@edgeless.systems>
Co-authored-by: leongross <leon.gross@rub.de>
Co-authored-by: Moritz Eckert <m1gh7ym0@gmail.com>

debugd/coordinator/streamer.go

@@ -0,0 +1,124 @@
+package coordinator
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"sync"
+
+	pb "github.com/edgelesssys/constellation/debugd/service"
+	"github.com/schollz/progressbar/v3"
+	"github.com/spf13/afero"
+)
+
+// FileStreamer handles reading and writing of a file using a stream of chunks.
+type FileStreamer struct {
+	fs  afero.Fs
+	mux sync.RWMutex
+}
+
+// ReadChunkStream is abstraction over a gRPC stream that allows us to receive chunks via gRPC.
+type ReadChunkStream interface {
+	Recv() (*pb.Chunk, error)
+}
+
+// WriteChunkStream is abstraction over a gRPC stream that allows us to send chunks via gRPC.
+type WriteChunkStream interface {
+	Send(chunk *pb.Chunk) error
+}
+
+// NewFileStreamer creates a new FileStreamer.
+func NewFileStreamer(fs afero.Fs) *FileStreamer {
+	return &FileStreamer{
+		fs:  fs,
+		mux: sync.RWMutex{},
+	}
+}
+
+// WriteStream opens a file to write to and streams chunks from a gRPC stream into the file.
+func (f *FileStreamer) WriteStream(filename string, stream ReadChunkStream, showProgress bool) error {
+	f.mux.Lock()
+	defer f.mux.Unlock()
+	file, err := f.fs.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o755)
+	if err != nil {
+		return fmt.Errorf("open %v for writing failed: %w", filename, err)
+	}
+	defer file.Close()
+
+	var bar *progressbar.ProgressBar
+	if showProgress {
+		bar = progressbar.NewOptions64(
+			-1,
+			progressbar.OptionSetDescription("receiving coordinator"),
+			progressbar.OptionShowBytes(true),
+			progressbar.OptionClearOnFinish(),
+		)
+		defer bar.Close()
+	}
+
+	for {
+		chunk, err := stream.Recv()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return fmt.Errorf("reading stream failed: %w", err)
+		}
+		if _, err := file.Write(chunk.Content); err != nil {
+			return fmt.Errorf("writing chunk to disk failed: %w", err)
+		}
+		if showProgress {
+			bar.Add(len(chunk.Content))
+		}
+	}
+
+	return nil
+}
+
+// ReadStream opens a file to read from and streams its contents chunkwise over gRPC.
+func (f *FileStreamer) ReadStream(filename string, stream WriteChunkStream, chunksize uint, showProgress bool) error {
+	f.mux.RLock()
+	defer f.mux.RUnlock()
+	if chunksize == 0 {
+		return errors.New("invalid chunksize")
+	}
+	file, err := f.fs.OpenFile(filename, os.O_RDONLY, 0o755)
+	if err != nil {
+		return fmt.Errorf("open %v for reading failed: %w", filename, err)
+	}
+	defer file.Close()
+
+	var bar *progressbar.ProgressBar
+	if showProgress {
+		stat, err := file.Stat()
+		if err != nil {
+			return fmt.Errorf("performing stat on %v to get the file size failed: %w", filename, err)
+		}
+		bar = progressbar.NewOptions64(
+			stat.Size(),
+			progressbar.OptionSetDescription("uploading coordinator"),
+			progressbar.OptionShowBytes(true),
+			progressbar.OptionClearOnFinish(),
+		)
+		defer bar.Close()
+	}
+
+	buf := make([]byte, chunksize)
+	for {
+		n, err := file.Read(buf)
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return nil
+			}
+			return fmt.Errorf("reading file chunk failed: %w", err)
+		}
+
+		if err = stream.Send(&pb.Chunk{Content: buf[:n]}); err != nil {
+			return fmt.Errorf("sending chunk failed: %w", err)
+		}
+		if showProgress {
+			bar.Add(n)
+		}
+	}
+}