gcsbuf/reader.go

package buf

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

	"cloud.google.com/go/storage"
	"google.golang.org/api/googleapi"
)

const defaultPollInterval = 1 * time.Second

// ReaderOptions configures a Reader.
type ReaderOptions struct {
	// Client is the GCS client used for all operations.
	Client *storage.Client
	// Bucket is the GCS bucket name.
	Bucket string
	// Object is the GCS object name.
	Object string
	// PollInterval is how often the reader polls for new data or object creation. Default 1s.
	PollInterval time.Duration
	// OnOp, if set, is called after each GCS API operation completes.
	OnOp OnOp
}

type Reader struct {
	ctx          context.Context
	client       *storage.Client
	bucket       string
	object       string
	pollInterval time.Duration
	onOp         OnOp

	cursor int64
	done   bool
}

func NewReader(ctx context.Context, opts ReaderOptions) *Reader {
	interval := opts.PollInterval
	if interval == 0 {
		interval = defaultPollInterval
	}
	return &Reader{
		ctx:          ctx,
		client:       opts.Client,
		bucket:       opts.Bucket,
		object:       opts.Object,
		pollInterval: interval,
		onOp:         opts.OnOp,
	}
}

func (r *Reader) Read(p []byte) (int, error) {
	for {
		objh := r.client.Bucket(r.bucket).Object(r.object)

		attrs, err := retry(r.ctx, r.onOp, "attrs", func() (*storage.ObjectAttrs, error) {
			return objh.Attrs(r.ctx)
		})
		if err != nil {
			if isNotFound(err) {
				// Object doesn't exist yet, poll.
				if err := r.sleep(); err != nil {
					return 0, err
				}
				continue
			}
			return 0, fmt.Errorf("getting object attrs: %w", err)
		}

		if r.cursor >= attrs.Size {
			// No new data. Check if done.
			if attrs.Metadata["done"] == "true" {
				r.done = true
				return 0, io.EOF
			}
			// Not done, poll for more data.
			if err := r.sleep(); err != nil {
				return 0, err
			}
			continue
		}

		// Read available data.
		length := int64(len(p))
		available := attrs.Size - r.cursor
		if length > available {
			length = available
		}

		n, err := retry(r.ctx, r.onOp, "read", func() (int, error) {
			rc, err := objh.NewRangeReader(r.ctx, r.cursor, length)
			if err != nil {
				return 0, err
			}
			n, err := io.ReadFull(rc, p[:length])
			rc.Close()
			if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
				return n, err
			}
			return n, nil
		})
		if err != nil {
			return n, fmt.Errorf("reading object data: %w", err)
		}

		r.cursor += int64(n)
		return n, nil
	}
}

func (r *Reader) Close() error {
	return nil
}

func (r *Reader) sleep() error {
	select {
	case <-time.After(r.pollInterval):
		return nil
	case <-r.ctx.Done():
		return r.ctx.Err()
	}
}

func isNotFound(err error) bool {
	if errors.Is(err, storage.ErrObjectNotExist) {
		return true
	}
	var apiErr *googleapi.Error
	if errors.As(err, &apiErr) && apiErr.Code == http.StatusNotFound {
		return true
	}
	return false
}