edge/pkg/storage/driver/cache/reader.go

package cache

import (
	"bytes"
	"context"
	"fmt"
	"io"

	"github.com/hashicorp/golang-lru/v2/expirable"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

type readCacher struct {
	reader io.ReadSeekCloser
	cache  *expirable.LRU[string, []byte]
	key    string
	buffer bytes.Buffer
}

// Close implements io.ReadSeekCloser.
func (r *readCacher) Close() error {
	if err := r.reader.Close(); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Read implements io.ReadSeekCloser.
func (r *readCacher) Read(p []byte) (n int, err error) {
	length, err := r.reader.Read(p)
	if err != nil {
		if err == io.EOF {
			if length > 0 {
				if _, err := r.buffer.Write(p[:length]); err != nil {
					logger.Error(context.Background(), "could not write to buffer", logger.CapturedE(errors.WithStack(err)))
					return length, io.EOF
				}
			}

			logger.Debug(context.Background(), "caching blob", logger.F("cacheKey", r.key))
			r.cache.Add(r.key, r.buffer.Bytes())

			return length, io.EOF
		}

		return length, errors.WithStack(err)
	}

	if length > 0 {
		if _, err := r.buffer.Write(p[:length]); err != nil {
			logger.Error(context.Background(), "could not write to buffer", logger.CapturedE(errors.WithStack(err)))
		}
	}

	return length, nil
}

// Seek implements io.ReadSeekCloser.
func (r *readCacher) Seek(offset int64, whence int) (int64, error) {
	length, err := r.reader.Seek(offset, whence)
	if err != nil {
		return length, errors.WithStack(err)
	}

	return length, nil
}

var _ io.ReadSeekCloser = &readCacher{}

type cachedReader struct {
	buffer []byte
	offset int64
}

// Read implements io.ReadSeekCloser.
func (r *cachedReader) Read(p []byte) (n int, err error) {
	available := len(r.buffer) - int(r.offset)
	if available == 0 {
		return 0, io.EOF
	}

	size := len(p)
	if size > available {
		size = available
	}

	copy(p, r.buffer[r.offset:r.offset+int64(size)])
	r.offset += int64(size)

	return size, nil
}

// Close implements io.ReadSeekCloser.
func (r *cachedReader) Close() error {
	return nil
}

// Seek implements io.ReadSeekCloser.
func (r *cachedReader) 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 = int64(len(r.buffer)) + offset
	default:
		return 0, errors.Errorf("unknown seek whence '%d'", whence)
	}

	if newOffset > int64(len(r.buffer)) || newOffset < 0 {
		return 0, fmt.Errorf("invalid offset %d", offset)
	}

	r.offset = newOffset

	return newOffset, nil
}

var _ io.ReadSeekCloser = &cachedReader{}