Merge pull request #15 from base-org/gzip-s3

Allow Storing Compressed Blobs

common/flags/config.go

@@ -28,6 +28,7 @@ type S3Config struct {
 	S3CredentialType S3CredentialType
 	AccessKey        string
 	SecretAccessKey  string
+	Compress         bool
 }
 
 func (c S3Config) check() error {
@@ -104,6 +105,7 @@ func readS3Config(ctx *cli.Context) S3Config {
 		UseHttps:         ctx.Bool(S3EndpointHttpsFlagName),
 		Bucket:           ctx.String(S3BucketFlagName),
 		S3CredentialType: toS3CredentialType(ctx.String(S3CredentialTypeFlagName)),
+		Compress:         ctx.Bool(S3CompressFlagName),
 	}
 }
 

common/flags/flags.go

@@ -12,6 +12,7 @@ const (
 	S3CredentialTypeFlagName        = "s3-credential-type"
 	S3EndpointFlagName              = "s3-endpoint"
 	S3EndpointHttpsFlagName         = "s3-endpoint-https"
+	S3CompressFlagName              = "s3-compress"
 	S3AccessKeyFlagName             = "s3-access-key"
 	S3SecretAccessKeyFlagName       = "s3-secret-access-key"
 	S3BucketFlagName                = "s3-bucket"
@@ -51,6 +52,12 @@ func CLIFlags(envPrefix string) []cli.Flag {
 			Value:   true,
 			EnvVars: opservice.PrefixEnvVar(envPrefix, "S3_ENDPOINT_HTTPS"),
 		},
+		&cli.BoolFlag{
+			Name:    S3CompressFlagName,
+			Usage:   "Whether to compress data before storing in S3",
+			Value:   false,
+			EnvVars: opservice.PrefixEnvVar(envPrefix, "S3_COMPRESS"),
+		},
 		&cli.StringFlag{
 			Name:    S3AccessKeyFlagName,
 			Usage:   "The S3 access key for the bucket",

common/storage/s3.go

@@ -2,8 +2,10 @@ package storage
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
+	"io"
 
 	"github.com/base-org/blob-archiver/common/flags"
 	"github.com/ethereum/go-ethereum/common"
@@ -13,9 +15,10 @@ import (
 )
 
 type S3Storage struct {
-	s3     *minio.Client
-	bucket string
-	log    log.Logger
+	s3       *minio.Client
+	bucket   string
+	log      log.Logger
+	compress bool
 }
 
 func NewS3Storage(cfg flags.S3Config, l log.Logger) (*S3Storage, error) {
@@ -36,9 +39,10 @@ func NewS3Storage(cfg flags.S3Config, l log.Logger) (*S3Storage, error) {
 	}
 
 	return &S3Storage{
-		s3:     client,
-		bucket: cfg.Bucket,
-		log:    l,
+		s3:       client,
+		bucket:   cfg.Bucket,
+		log:      l,
+		compress: cfg.Compress,
 	}, nil
 }
 
@@ -63,7 +67,7 @@ func (s *S3Storage) Read(ctx context.Context, hash common.Hash) (BlobData, error
 		return BlobData{}, ErrStorage
 	}
 	defer res.Close()
-	_, err = res.Stat()
+	stat, err := res.Stat()
 	if err != nil {
 		errResponse := minio.ToErrorResponse(err)
 		if errResponse.Code == "NoSuchKey" {
@@ -75,8 +79,19 @@ func (s *S3Storage) Read(ctx context.Context, hash common.Hash) (BlobData, error
 		}
 	}
 
+	var reader io.ReadCloser = res
+	defer reader.Close()
+
+	if stat.Metadata.Get("Content-Encoding") == "gzip" {
+		reader, err = gzip.NewReader(reader)
+		if err != nil {
+			s.log.Warn("error creating gzip reader", "hash", hash.String(), "err", err)
+			return BlobData{}, ErrMarshaling
+		}
+	}
+
 	var data BlobData
-	err = json.NewDecoder(res).Decode(&data)
+	err = json.NewDecoder(reader).Decode(&data)
 	if err != nil {
 		s.log.Warn("error decoding blob", "hash", hash.String(), "err", err)
 		return BlobData{}, ErrMarshaling
@@ -92,10 +107,22 @@ func (s *S3Storage) Write(ctx context.Context, data BlobData) error {
 		return ErrMarshaling
 	}
 
-	reader := bytes.NewReader(b)
-	_, err = s.s3.PutObject(ctx, s.bucket, data.Header.BeaconBlockHash.String(), reader, int64(len(b)), minio.PutObjectOptions{
+	options := minio.PutObjectOptions{
 		ContentType: "application/json",
-	})
+	}
+
+	if s.compress {
+		b, err = compress(b)
+		if err != nil {
+			s.log.Warn("error compressing blob", "err", err)
+			return ErrCompress
+		}
+		options.ContentEncoding = "gzip"
+	}
+
+	reader := bytes.NewReader(b)
+
+	_, err = s.s3.PutObject(ctx, s.bucket, data.Header.BeaconBlockHash.String(), reader, int64(len(b)), options)
 
 	if err != nil {
 		s.log.Warn("error writing blob", "err", err)
@@ -105,3 +132,17 @@ func (s *S3Storage) Write(ctx context.Context, data BlobData) error {
 	s.log.Info("wrote blob", "hash", data.Header.BeaconBlockHash.String())
 	return nil
 }
+
+func compress(in []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+	_, err := gz.Write(in)
+	if err != nil {
+		return nil, err
+	}
+	err = gz.Close()
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}

common/storage/storage.go

@@ -21,6 +21,8 @@ var (
 	ErrStorage = errors.New("error accessing storage")
 	// ErrMarshaling is returned when there is an error in (un)marshaling the blob
 	ErrMarshaling = errors.New("error encoding/decoding blob")
+	// ErrCompress is returned when there is an error gzipping the data
+	ErrCompress = errors.New("error compressing blob")
 )
 
 type Header struct {

validator/cmd/main.go

@@ -57,7 +57,7 @@ func Main() cliapp.LifecycleAction {
 		}
 
 		beaconClient := service.NewBlobSidecarClient(cfg.BeaconConfig.BeaconURL)
-		blobClient := service.NewBlobSidecarClient(cfg.BeaconConfig.BeaconURL)
+		blobClient := service.NewBlobSidecarClient(cfg.BlobConfig.BeaconURL)
 
 		return service.NewValidator(l, headerClient, beaconClient, blobClient, closeApp), nil
 	}