edge/pkg/storage/rpc/client/blob_bucket.go

package client

import (
	"context"
	"io"

	"github.com/keegancsmith/rpc"

	"forge.cadoles.com/arcad/edge/pkg/storage"
	"forge.cadoles.com/arcad/edge/pkg/storage/rpc/server/blob"
	"github.com/pkg/errors"
)

type BlobBucket struct {
	name   string
	id     blob.BucketID
	client *rpc.Client
}

// Size implements storage.BlobBucket
func (b *BlobBucket) Size(ctx context.Context) (int64, error) {
	args := blob.GetBucketSizeArgs{
		BucketID: b.id,
	}

	reply := blob.GetBucketSizeReply{}

	if err := b.client.Call(ctx, "Service.GetBucketSize", args, &reply); err != nil {
		return 0, errors.WithStack(err)
	}

	return reply.Size, nil
}

// Name implements storage.BlobBucket
func (b *BlobBucket) Name() string {
	return b.name
}

// Close implements storage.BlobBucket
func (b *BlobBucket) Close() error {
	args := blob.CloseBucketArgs{
		BucketID: b.id,
	}

	reply := blob.CloseBucketReply{}

	if err := b.client.Call(context.Background(), "Service.CloseBucket", args, &reply); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Delete implements storage.BlobBucket
func (b *BlobBucket) Delete(ctx context.Context, id storage.BlobID) error {
	args := blob.DeleteBucketArgs{
		BucketName: b.name,
	}

	reply := blob.DeleteBucketReply{}

	if err := b.client.Call(context.Background(), "Service.DeleteBucket", args, &reply); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Get implements storage.BlobBucket
func (b *BlobBucket) Get(ctx context.Context, id storage.BlobID) (storage.BlobInfo, error) {
	args := blob.GetBlobInfoArgs{
		Name:   b.name,
		BlobID: id,
	}

	reply := blob.GetBlobInfoReply{}

	if err := b.client.Call(context.Background(), "Service.GetBlobInfo", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return reply.BlobInfo, nil
}

// List implements storage.BlobBucket
func (b *BlobBucket) List(ctx context.Context) ([]storage.BlobInfo, error) {
	args := blob.ListBlobInfoArgs{
		BucketID: b.id,
	}

	reply := blob.ListBlobInfoReply{}

	if err := b.client.Call(context.Background(), "Service.ListBlobInfo", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return reply.BlobInfos, nil
}

// NewReader implements storage.BlobBucket
func (b *BlobBucket) NewReader(ctx context.Context, id storage.BlobID) (io.ReadSeekCloser, error) {
	args := blob.NewBlobReaderArgs{
		BucketID: b.id,
		BlobID:   id,
	}

	reply := blob.NewBlobReaderReply{}

	if err := b.client.Call(context.Background(), "Service.NewBlobReader", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return &blobReaderCloser{
		readerID: reply.ReaderID,
		client:   b.client,
	}, nil
}

// NewWriter implements storage.BlobBucket
func (b *BlobBucket) NewWriter(ctx context.Context, id storage.BlobID) (io.WriteCloser, error) {
	args := blob.NewBlobWriterArgs{
		BucketID: b.id,
		BlobID:   id,
	}

	reply := blob.NewBlobWriterReply{}

	if err := b.client.Call(context.Background(), "Service.NewBlobWriter", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return &blobWriterCloser{
		blobID:   id,
		writerID: reply.WriterID,
		client:   b.client,
	}, nil
}

type blobWriterCloser struct {
	blobID   storage.BlobID
	writerID blob.WriterID
	client   *rpc.Client
}

// Write implements io.WriteCloser
func (bwc *blobWriterCloser) Write(data []byte) (int, error) {
	args := blob.WriteBlobArgs{
		WriterID: bwc.writerID,
		Data:     data,
	}

	reply := blob.WriteBlobReply{}

	if err := bwc.client.Call(context.Background(), "Service.WriteBlob", args, &reply); err != nil {
		return 0, errors.WithStack(err)
	}

	return reply.Written, nil
}

// Close implements io.WriteCloser
func (bwc *blobWriterCloser) Close() error {
	args := blob.CloseWriterArgs{
		WriterID: bwc.writerID,
	}

	reply := blob.CloseBucketReply{}

	if err := bwc.client.Call(context.Background(), "Service.CloseWriter", args, &reply); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

type blobReaderCloser struct {
	readerID blob.ReaderID
	client   *rpc.Client
}

// Read implements io.ReadSeekCloser
func (brc *blobReaderCloser) Read(p []byte) (int, error) {
	args := blob.ReadBlobArgs{
		ReaderID: brc.readerID,
		Length:   len(p),
	}

	reply := blob.ReadBlobReply{}

	if err := brc.client.Call(context.Background(), "Service.ReadBlob", args, &reply); err != nil {
		return 0, errors.WithStack(err)
	}

	copy(p, reply.Data)

	if reply.EOF {
		return reply.Read, io.EOF
	}

	return reply.Read, nil
}

// Seek implements io.ReadSeekCloser
func (brc *blobReaderCloser) Seek(offset int64, whence int) (int64, error) {
	args := blob.SeekBlobArgs{
		ReaderID: brc.readerID,
		Offset:   offset,
		Whence:   whence,
	}

	reply := blob.SeekBlobReply{}

	if err := brc.client.Call(context.Background(), "Service.SeekBlob", args, &reply); err != nil {
		return 0, errors.WithStack(err)
	}

	return reply.Read, nil
}

// Close implements io.ReadSeekCloser
func (brc *blobReaderCloser) Close() error {
	args := blob.CloseReaderArgs{
		ReaderID: brc.readerID,
	}

	reply := blob.CloseReaderReply{}

	if err := brc.client.Call(context.Background(), "Service.CloseReader", args, &reply); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

var (
	_ storage.BlobBucket = &BlobBucket{}
	_ storage.BlobInfo   = &BlobInfo{}
	_ io.WriteCloser     = &blobWriterCloser{}
	_ io.ReadSeekCloser  = &blobReaderCloser{}
)