first agent crack

Signed-off-by: Jason Hall <imjasonh@gmail.com>

reader.go

@@ -0,0 +1,127 @@
+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
+}
+
+type Reader struct {
+	ctx          context.Context
+	client       *storage.Client
+	bucket       string
+	object       string
+	pollInterval time.Duration
+
+	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,
+	}
+}
+
+func (r *Reader) Read(p []byte) (int, error) {
+	for {
+		objh := r.client.Bucket(r.bucket).Object(r.object)
+
+		attrs, err := 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
+		}
+
+		rc, err := objh.NewRangeReader(r.ctx, r.cursor, length)
+		if err != nil {
+			return 0, fmt.Errorf("creating range reader: %w", err)
+		}
+		n, err := io.ReadFull(rc, p[:length])
+		rc.Close()
+		if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
+			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
+}