ceremonyclient/pebble/objstorage/objstorageprovider/remote_readable.go
Cassandra Heart 2e2a1e4789
v1.2.0 (#31)
2024-01-03 01:31:42 -06:00

163 lines
4.6 KiB
Go

// Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.
package objstorageprovider
import (
"context"
"io"
"github.com/cockroachdb/pebble/internal/base"
"github.com/cockroachdb/pebble/objstorage"
"github.com/cockroachdb/pebble/objstorage/objstorageprovider/sharedcache"
"github.com/cockroachdb/pebble/objstorage/remote"
)
const remoteMaxReadaheadSize = 1024 * 1024 /* 1MB */
// remoteReadable is a very simple implementation of Readable on top of the
// ReadCloser returned by remote.Storage.CreateObject.
type remoteReadable struct {
objReader remote.ObjectReader
size int64
fileNum base.DiskFileNum
provider *provider
}
var _ objstorage.Readable = (*remoteReadable)(nil)
func (p *provider) newRemoteReadable(
objReader remote.ObjectReader, size int64, fileNum base.DiskFileNum,
) *remoteReadable {
return &remoteReadable{
objReader: objReader,
size: size,
fileNum: fileNum,
provider: p,
}
}
// ReadAt is part of the objstorage.Readable interface.
func (r *remoteReadable) ReadAt(ctx context.Context, p []byte, offset int64) error {
return r.readInternal(ctx, p, offset, false /* forCompaction */)
}
// readInternal performs a read for the object, using the cache when
// appropriate.
func (r *remoteReadable) readInternal(
ctx context.Context, p []byte, offset int64, forCompaction bool,
) error {
if cache := r.provider.remote.cache; cache != nil {
flags := sharedcache.ReadFlags{
// Don't add data to the cache if this read is for a compaction.
ReadOnly: forCompaction,
}
return r.provider.remote.cache.ReadAt(ctx, r.fileNum, p, offset, r.objReader, r.size, flags)
}
return r.objReader.ReadAt(ctx, p, offset)
}
func (r *remoteReadable) Close() error {
defer func() { r.objReader = nil }()
return r.objReader.Close()
}
func (r *remoteReadable) Size() int64 {
return r.size
}
func (r *remoteReadable) NewReadHandle(_ context.Context) objstorage.ReadHandle {
// TODO(radu): use a pool.
rh := &remoteReadHandle{readable: r}
rh.readahead.state = makeReadaheadState(remoteMaxReadaheadSize)
return rh
}
type remoteReadHandle struct {
readable *remoteReadable
readahead struct {
state readaheadState
data []byte
offset int64
}
forCompaction bool
}
var _ objstorage.ReadHandle = (*remoteReadHandle)(nil)
// ReadAt is part of the objstorage.ReadHandle interface.
func (r *remoteReadHandle) ReadAt(ctx context.Context, p []byte, offset int64) error {
readaheadSize := r.maybeReadahead(offset, len(p))
// Check if we already have the data from a previous read-ahead.
if rhSize := int64(len(r.readahead.data)); rhSize > 0 {
if r.readahead.offset <= offset && r.readahead.offset+rhSize > offset {
n := copy(p, r.readahead.data[offset-r.readahead.offset:])
if n == len(p) {
// All data was available.
return nil
}
// Use the data that we had and do a shorter read.
offset += int64(n)
p = p[n:]
readaheadSize -= n
}
}
if readaheadSize > len(p) {
// Don't try to read past EOF.
if offset+int64(readaheadSize) > r.readable.size {
readaheadSize = int(r.readable.size - offset)
if readaheadSize <= 0 {
// This shouldn't happen in practice (Pebble should never try to read
// past EOF).
return io.EOF
}
}
r.readahead.offset = offset
// TODO(radu): we need to somehow account for this memory.
if cap(r.readahead.data) >= readaheadSize {
r.readahead.data = r.readahead.data[:readaheadSize]
} else {
r.readahead.data = make([]byte, readaheadSize)
}
if err := r.readable.readInternal(ctx, r.readahead.data, offset, r.forCompaction); err != nil {
// Make sure we don't treat the data as valid next time.
r.readahead.data = r.readahead.data[:0]
return err
}
copy(p, r.readahead.data)
return nil
}
return r.readable.readInternal(ctx, p, offset, r.forCompaction)
}
func (r *remoteReadHandle) maybeReadahead(offset int64, len int) int {
if r.forCompaction {
return remoteMaxReadaheadSize
}
return int(r.readahead.state.maybeReadahead(offset, int64(len)))
}
// Close is part of the objstorage.ReadHandle interface.
func (r *remoteReadHandle) Close() error {
r.readable = nil
r.readahead.data = nil
return nil
}
// SetupForCompaction is part of the objstorage.ReadHandle interface.
func (r *remoteReadHandle) SetupForCompaction() {
r.forCompaction = true
}
// RecordCacheHit is part of the objstorage.ReadHandle interface.
func (r *remoteReadHandle) RecordCacheHit(_ context.Context, offset, size int64) {
if !r.forCompaction {
r.readahead.state.recordCacheHit(offset, size)
}
}