Repository: googleforgames/open-saves
Pull request #268 (2021-07-08 23:09:37)

Feat: Chunk blob support
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://github.com/googleforgames/open-saves/blob/main/docs/contributing.md and developer guide https://github.com/googleforgames/open-saves/blob/main/docs/development.md
2. Please label this pull request according to what type of issue you are addressing.
3. Ensure you have added or ran the appropriate tests for your PR.
-->

**What type of PR is this?**
/kind feat

**What this PR does / Why we need it**:
Add chunk blob support to the server code. This PR adds implementations to the following methods:
- CreateChunkedBlob
- UploadChunk
- CommitChunkedUpload
- AbortChunkedUpload
- GetBlobChunk

**Which issue(s) this PR fixes**:
<!--
*Automatically closes linked issue when PR is merged.
Usage: `Closes #<issue number>`, or `Closes (paste link of issue)`.
-->
Closes #

**Special notes for your reviewer**:
I still need to implement garbage collector support.


Patch:
diff --git a/api/open_saves.pb.go b/api/open_saves.pb.go
index 6543bb3..444078b 100644
--- a/api/open_saves.pb.go
+++ b/api/open_saves.pb.go
@@ -15,7 +15,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.25.0
-// 	protoc        v3.15.8
+// 	protoc        v3.17.3
 // source: api/open_saves.proto
 
 package opensaves
diff --git a/internal/app/server/open_saves.go b/internal/app/server/open_saves.go
index 6d7d612..c429848 100644
--- a/internal/app/server/open_saves.go
+++ b/internal/app/server/open_saves.go
@@ -27,20 +27,23 @@ import (
 	"github.com/googleforgames/open-saves/internal/pkg/cache/redis"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/blobref"
+	"github.com/googleforgames/open-saves/internal/pkg/metadb/blobref/chunkref"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/record"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/store"
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/emptypb"
 	empty "google.golang.org/protobuf/types/known/emptypb"
 )
 
 // TODO(hongalex): make this a configurable field for users.
 const (
-	maxRecordSizeToCache int = 10 * 1024 * 1024 // 10 MB
-	maxInlineBlobSize    int = 64 * 1024        // 64 KiB
-	streamBufferSize     int = 1 * 1024 * 1024  // 1 MiB
-	opaqueStringLimit    int = 32 * 1024        // 32 KiB
+	maxRecordSizeToCache int = 10 * 1024 * 1024       // 10 MB
+	maxInlineBlobSize    int = 64 * 1024              // 64 KiB
+	streamBufferSize     int = 1 * 1024 * 1024        // 1 MiB
+	opaqueStringLimit    int = 32 * 1024              // 32 KiB
+	chunkSizeLimit       int = 1 * 1024 * 1024 * 1024 // 1 GiB
 )
 
 type openSavesServer struct {
@@ -279,6 +282,13 @@ func (s *openSavesServer) blobRefFail(ctx context.Context, blobref *blobref.Blob
 	}
 }
 
+func (s *openSavesServer) chunkRefFail(ctx context.Context, chunk *chunkref.ChunkRef) {
+	chunk.Fail()
+	if err := s.metaDB.UpdateChunkRef(ctx, chunk); err != nil {
+		log.Errorf("Failed to mark ChunkRef (%v) as Failed: %v", chunk.Key, err)
+	}
+}
+
 func (s *openSavesServer) insertExternalBlob(ctx context.Context, stream pb.OpenSaves_CreateBlobServer, meta *pb.BlobMetadata) error {
 	log.Debugf("Inserting external blob: %v\n", meta)
 	// Create a blob reference based on the metadata.
@@ -457,6 +467,193 @@ func (s *openSavesServer) Ping(ctx context.Context, req *pb.PingRequest) (*pb.Pi
 	}, nil
 }
 
+func (s *openSavesServer) CreateChunkedBlob(ctx context.Context, req *pb.CreateChunkedBlobRequest) (*pb.CreateChunkedBlobResponse, error) {
+	b := blobref.NewChunkedBlobRef(req.GetStoreKey(), req.GetRecordKey())
+	b, err := s.metaDB.InsertBlobRef(ctx, b)
+	if err != nil {
+		log.Errorf("CreateChunkedBlob failed for record (%v), blob key (%v): %v", b.RecordKey, b.Key, err)
+		return nil, err
+	}
+	return &pb.CreateChunkedBlobResponse{
+		SessionId: b.Key.String(),
+	}, nil
+}
+
+func (s *openSavesServer) UploadChunk(stream pb.OpenSaves_UploadChunkServer) error {
+	ctx := stream.Context()
+	req, err := stream.Recv()
+	if err != nil {
+		log.Errorf("Recv() returned error: %v", err)
+		return err
+	}
+	meta := req.GetMetadata()
+	if meta == nil {
+		log.Error("GetMetadata() returned nil. First message should be a ChunkMetadata.")
+		return status.Error(codes.InvalidArgument, "First message should be a ChunkMetadata.")
+	}
+	blobKey, err := uuid.Parse(meta.GetSessionId())
+	if err != nil {
+		log.Errorf("SessionId is not a valid UUID string: %v", err)
+		return status.Errorf(codes.InvalidArgument, "SessionId is not a valid UUID string: %v", err)
+	}
+
+	// Create a chunk reference based on the metadata.
+	chunk := chunkref.New(blobKey, int32(meta.GetNumber()))
+	if err := s.metaDB.InsertChunkRef(ctx, chunk); err != nil {
+		return err
+	}
+
+	writer, err := s.blobStore.NewWriter(ctx, chunk.ObjectPath())
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if writer != nil {
+			writer.Close()
+			// This means an abnormal exit, so make sure to mark the blob as Fail.
+			s.chunkRefFail(ctx, chunk)
+		}
+	}()
+
+	written := 0
+	for {
+		req, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			log.Errorf("UploadChunk: stream recv error: %v", err)
+			return err
+		}
+		fragment := req.GetContent()
+		if fragment == nil {
+			return status.Error(codes.InvalidArgument, "Subsequent input messages must contain chunk content")
+		}
+		n, err := writer.Write(fragment)
+		if err != nil {
+			log.Errorf("UploadChunk: BlobStore write error: %v", err)
+			return err
+		}
+		written += n
+		// TODO(yuryu): This is not suitable for unit tests until we make the value
+		// configurable, or have a BlobStore mock.
+		if written > chunkSizeLimit {
+			err := status.Errorf(codes.ResourceExhausted, "UploadChunk: Received chunk size (%v) exceed the limit (%v)", written, chunkSizeLimit)
+			log.Error(err)
+			return err
+		}
+	}
+	err = writer.Close()
+	writer = nil
+	if err != nil {
+		log.Errorf("writer.Close() failed on chunk object %v: %v", chunk.ObjectPath(), err)
+		s.chunkRefFail(ctx, chunk)
+		// The uploaded object can be deleted immediately.
+		if derr := s.blobStore.Delete(ctx, chunk.ObjectPath()); derr != nil {
+			log.Errorf("Delete chunk failed after writer.Close() error: %v", derr)
+		}
+		return err
+	}
+
+	// Update the chunk size based on the actual bytes written
+	// MarkChunkRefReady commits the new change to Datastore
+	chunk.Size = int32(written)
+	if err := s.metaDB.MarkChunkRefReady(ctx, chunk); err != nil {
+		log.Errorf("Failed to update chunkref metadata (%v): %v", chunk.Key, err)
+		s.chunkRefFail(ctx, chunk)
+		// Do not delete the chunk object here. Leave it to the garbage collector.
+		return err
+	}
+	return stream.SendAndClose(meta)
+}
+
+func (s *openSavesServer) CommitChunkedUpload(ctx context.Context, req *pb.CommitChunkedUploadRequest) (*pb.BlobMetadata, error) {
+	blobKey, err := uuid.Parse(req.GetSessionId())
+	if err != nil {
+		log.Errorf("SessionId is not a valid UUID string: %v", err)
+		return nil, status.Errorf(codes.InvalidArgument, "SessionId is not a valid UUID string: %v", err)
+	}
+	blob, err := s.metaDB.GetBlobRef(ctx, blobKey)
+	if err != nil {
+		log.Errorf("Cannot retrieve chunked blob metadata for session (%v): %v", blobKey, err)
+		return nil, err
+	}
+	_, blob, err = s.metaDB.PromoteBlobRefToCurrent(ctx, blob)
+	if err != nil {
+		log.Errorf("PromoteBlobRefToCurrent failed for object %v: %v", blob.ObjectPath(), err)
+		// Do not delete the blob object here. Leave it to the garbage collector.
+		return nil, err
+	}
+	return blob.ToProto(), nil
+}
+
+func (s *openSavesServer) AbortChunkedUpload(ctx context.Context, req *pb.AbortChunkedUploadRequest) (*emptypb.Empty, error) {
+	id, err := uuid.Parse(req.GetSessionId())
+	if err != nil {
+		log.Errorf("SessionId is not a valid UUID: %v", err)
+		return new(empty.Empty), status.Errorf(codes.InvalidArgument, "SessionId is not a valid UUID: %v", err)
+	}
+	err = s.metaDB.MarkUncommittedBlobForDeletion(ctx, id)
+	if err != nil {
+		log.Errorf("AbortChunkedUpload failed for session (%v): %v", id, err)
+	}
+	return new(empty.Empty), err
+}
+
+func (s *openSavesServer) GetBlobChunk(req *pb.GetBlobChunkRequest, response pb.OpenSaves_GetBlobChunkServer) error {
+	ctx := response.Context()
+
+	chunk, err := s.metaDB.FindChunkRefByNumber(ctx, req.GetStoreKey(), req.GetRecordKey(), int32(req.GetChunkNumber()))
+	if err != nil {
+		log.Errorf("GetBlobChunk failed to get chunk metadata for store (%v), record (%v), number (%v): %v",
+			req.GetStoreKey(), req.GetRecordKey(), req.GetChunkNumber(), err)
+		return err
+	}
+
+	err = response.Send(&pb.GetBlobChunkResponse{
+		Response: &pb.GetBlobChunkResponse_Metadata{
+			Metadata: chunk.ToProto(),
+		},
+	})
+	if err != nil {
+		log.Errorf("Send failed for metadata: %v", err)
+	}
+
+	reader, err := s.blobStore.NewReader(ctx, chunk.ObjectPath())
+	if err != nil {
+		log.Errorf("GetBlobChunk failed to create a reader for chunk (%v): %v", chunk.ObjectPath(), err)
+		return err
+	}
+	defer reader.Close()
+
+	buf := make([]byte, streamBufferSize)
+	sent := 0
+	for {
+		n, err := reader.Read(buf)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			log.Errorf("GetBlobChunk: BlobStore Reader returned error for object (%v): %v", chunk.ObjectPath(), err)
+			return err
+		}
+		err = response.Send(&pb.GetBlobChunkResponse{Response: &pb.GetBlobChunkResponse_Content{
+			Content: buf[:n],
+		}})
+		if err != nil {
+			log.Errorf("GetBlobChunk: Stream send error for object (%v): %v", chunk.ObjectPath(), err)
+			return err
+		}
+		sent += n
+	}
+	if sent != int(chunk.Size) {
+		log.Errorf("GetBlobChunk: Chunk size sent (%v) and stored in the metadata (%v) don't match.", sent, chunk.Size)
+		return status.Errorf(codes.DataLoss,
+			"GetBlobChunk: Chunk size sent (%v) and stored in the metadata (%v) don't match.", sent, chunk.Size)
+	}
+	return nil
+}
+
 func (s *openSavesServer) getRecordAndCache(ctx context.Context, storeKey, key string, hint *pb.Hint) (*record.Record, error) {
 	if shouldCheckCache(hint) {
 		r := new(record.Record)
diff --git a/internal/pkg/metadb/blobref/blobref.go b/internal/pkg/metadb/blobref/blobref.go
index adf2dbf..3ca824c 100644
--- a/internal/pkg/metadb/blobref/blobref.go
+++ b/internal/pkg/metadb/blobref/blobref.go
@@ -17,6 +17,7 @@ package blobref
 import (
 	"cloud.google.com/go/datastore"
 	"github.com/google/uuid"
+	pb "github.com/googleforgames/open-saves/api"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/timestamps"
 )
 
@@ -101,3 +102,12 @@ func NewChunkedBlobRef(storeKey, recordKey string) *BlobRef {
 func (b *BlobRef) ObjectPath() string {
 	return b.Key.String()
 }
+
+// ToProto returns a BlobMetadata representation of the object.
+func (b *BlobRef) ToProto() *pb.BlobMetadata {
+	return &pb.BlobMetadata{
+		StoreKey:  b.StoreKey,
+		RecordKey: b.RecordKey,
+		Size:      b.Size,
+	}
+}
diff --git a/internal/pkg/metadb/blobref/chunkref/chunkref.go b/internal/pkg/metadb/blobref/chunkref/chunkref.go
index dd4302f..c6e077a 100644
--- a/internal/pkg/metadb/blobref/chunkref/chunkref.go
+++ b/internal/pkg/metadb/blobref/chunkref/chunkref.go
@@ -20,6 +20,7 @@ import (
 
 	"cloud.google.com/go/datastore"
 	"github.com/google/uuid"
+	pb "github.com/googleforgames/open-saves/api"
 	"github.com/googleforgames/open-saves/internal/pkg/cache"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/blobref"
 	"github.com/googleforgames/open-saves/internal/pkg/metadb/timestamps"
@@ -123,3 +124,13 @@ func (c *ChunkRef) DecodeBytes(by []byte) error {
 	d := gob.NewDecoder(b)
 	return d.Decode(c)
 }
+
+// ToProto converts returns a pb.ChunkMetadata representation of the
+// ChunkRef object.
+func (c *ChunkRef) ToProto() *pb.ChunkMetadata {
+	return &pb.ChunkMetadata{
+		SessionId: c.BlobRef.String(),
+		Number:    int64(c.Number),
+		Size:      int64(c.Size),
+	}
+}


Test patch:
diff --git a/internal/pkg/metadb/blobref/blobref_test.go b/internal/pkg/metadb/blobref/blobref_test.go
index deef29d..a3fc69e 100644
--- a/internal/pkg/metadb/blobref/blobref_test.go
+++ b/internal/pkg/metadb/blobref/blobref_test.go
@@ -118,3 +118,19 @@ func TestBlobRef_GetObjectPath(t *testing.T) {
 
 	assert.Equal(t, blob.Key.String(), blob.ObjectPath())
 }
+
+func TestBlobRef_ToProto(t *testing.T) {
+	const (
+		size   = int64(123)
+		store  = "store"
+		record = "record"
+	)
+	b := NewBlobRef(size, store, record)
+
+	proto := b.ToProto()
+	if assert.NotNil(t, proto) {
+		assert.Equal(t, b.StoreKey, proto.GetStoreKey())
+		assert.Equal(t, b.RecordKey, proto.GetRecordKey())
+		assert.Equal(t, b.Size, proto.GetSize())
+	}
+}
diff --git a/internal/pkg/metadb/blobref/chunkref/chunkref_test.go b/internal/pkg/metadb/blobref/chunkref/chunkref_test.go
index 7beb11e..61be155 100644
--- a/internal/pkg/metadb/blobref/chunkref/chunkref_test.go
+++ b/internal/pkg/metadb/blobref/chunkref/chunkref_test.go
@@ -95,3 +95,15 @@ func TestChunkRef_EncodeDecodeBytes(t *testing.T) {
 	}
 	assert.Equal(t, c, decoded)
 }
+
+func TestChunkRef_ToProto(t *testing.T) {
+	blobKey := uuid.New()
+	c := New(blobKey, 42)
+	c.Size = 12345
+	proto := c.ToProto()
+	if assert.NotNil(t, proto) {
+		assert.Equal(t, blobKey.String(), proto.GetSessionId())
+		assert.EqualValues(t, c.Size, proto.GetSize())
+		assert.EqualValues(t, c.Number, proto.GetNumber())
+	}
+}