GCSのReader実装を拡張してio.Seekerを実装する

という実装のコード片をGithubのissueコメントから拾ったので自分で完成版を作ってみた。

内部的にNewRangeReaderを何度も呼び出しているのでストレージに対するReadの量は増える可能性があるのでそれだけ注意したい。

もしもシーク位置が比較的局所的な仕様ならばsunfish-shogi/bufseekioという便利なものがあるので、これでラップしてやるとReadトラフィックが増えるみたいな懸念はある程度払しょくできるかもしれない。

package storage

import (
    "context"
    "errors"
    "io"

    "cloud.google.com/go/storage"
)

var _ io.ReadSeekCloser = &SeekableReader{}

type SeekableReader struct {
    object   *storage.ObjectHandle
    ctx      context.Context
    reader   *storage.Reader
    offset   int64 // initial offset
    fileSize int64 // if this is known and set, it enables io.SeekEnd
}

func NewSeekableReader(ctx context.Context, object *storage.ObjectHandle) (*SeekableReader, error) {
    metaReader, err := object.NewRangeReader(ctx, 0, 0)
    if err != nil {
        return nil, err
    }
    if metaReader.Attrs.Size <= 0 {
        return nil, errors.New("fileSize should not be zero or negative value")
    }
    fileSize := metaReader.Attrs.Size

    return &SeekableReader{
        object:   object,
        ctx:      ctx,
        reader:   nil,
        offset:   0,
        fileSize: fileSize,
    }, nil
}

func (r *SeekableReader) Read(p []byte) (int, error) {
    var err error

    if r.reader == nil {
        r.reader, err = r.object.NewRangeReader(r.ctx, r.offset, int64(len(p)))
        if err != nil {
            return 0, err
        }
    }

    n, err := r.reader.Read(p)
    if err != nil {
        return 0, err
    }

    return n, nil
}

func (r *SeekableReader) Seek(offset int64, whence int) (int64, error) {
    var newOffset int64

    switch whence {
    case io.SeekStart:
        newOffset = offset
    case io.SeekCurrent:
        newOffset = r.offset + offset
    case io.SeekEnd:
        newOffset = r.fileSize - offset
    }

    r.Close()
    r.reader = nil
    r.offset = newOffset
    return r.offset, nil
}

func (r *SeekableReader) Close() error {
    if r.reader != nil {
        return r.reader.Close()
    }
    return nil
}