ceremonyclient/pebble/objstorage/remote/mem.go

162 lines
3.0 KiB
Go
Raw Permalink Normal View History

2024-01-03 07:31:42 +00:00
// 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 remote
import (
"bytes"
"context"
"io"
"os"
"strings"
"sync"
)
// NewInMem returns an in-memory implementation of the remote.Storage
// interface (for testing).
func NewInMem() Storage {
store := &inMemStore{}
store.mu.objects = make(map[string]*inMemObj)
return store
}
// inMemStore is an in-memory implementation of the remote.Storage interface
// (for testing).
type inMemStore struct {
mu struct {
sync.Mutex
objects map[string]*inMemObj
}
}
var _ Storage = (*inMemStore)(nil)
type inMemObj struct {
name string
data []byte
}
func (s *inMemStore) Close() error {
*s = inMemStore{}
return nil
}
func (s *inMemStore) ReadObject(
ctx context.Context, objName string,
) (_ ObjectReader, objSize int64, _ error) {
obj, err := s.getObj(objName)
if err != nil {
return nil, 0, err
}
return &inMemReader{data: obj.data}, int64(len(obj.data)), nil
}
type inMemReader struct {
data []byte
}
var _ ObjectReader = (*inMemReader)(nil)
func (r *inMemReader) ReadAt(ctx context.Context, p []byte, offset int64) error {
if offset+int64(len(p)) > int64(len(r.data)) {
return io.EOF
}
copy(p, r.data[offset:])
return nil
}
func (r *inMemReader) Close() error {
r.data = nil
return nil
}
func (s *inMemStore) CreateObject(objName string) (io.WriteCloser, error) {
return &inMemWriter{
store: s,
name: objName,
}, nil
}
type inMemWriter struct {
store *inMemStore
name string
buf bytes.Buffer
}
var _ io.WriteCloser = (*inMemWriter)(nil)
func (o *inMemWriter) Write(p []byte) (n int, err error) {
if o.store == nil {
panic("Write after Close")
}
return o.buf.Write(p)
}
func (o *inMemWriter) Close() error {
if o.store != nil {
o.store.addObj(&inMemObj{
name: o.name,
data: o.buf.Bytes(),
})
o.store = nil
}
return nil
}
func (s *inMemStore) List(prefix, delimiter string) ([]string, error) {
if delimiter != "" {
panic("delimiter unimplemented")
}
s.mu.Lock()
defer s.mu.Unlock()
res := make([]string, 0, len(s.mu.objects))
for name := range s.mu.objects {
if strings.HasPrefix(name, prefix) {
res = append(res, name)
}
}
return res, nil
}
func (s *inMemStore) Delete(objName string) error {
s.rmObj(objName)
return nil
}
// Size returns the length of the named object in bytesWritten.
func (s *inMemStore) Size(objName string) (int64, error) {
obj, err := s.getObj(objName)
if err != nil {
return 0, err
}
return int64(len(obj.data)), nil
}
func (s *inMemStore) IsNotExistError(err error) bool {
return err == os.ErrNotExist
}
func (s *inMemStore) getObj(name string) (*inMemObj, error) {
s.mu.Lock()
defer s.mu.Unlock()
obj, ok := s.mu.objects[name]
if !ok {
return nil, os.ErrNotExist
}
return obj, nil
}
func (s *inMemStore) addObj(o *inMemObj) {
s.mu.Lock()
defer s.mu.Unlock()
s.mu.objects[o.name] = o
}
func (s *inMemStore) rmObj(name string) {
s.mu.Lock()
defer s.mu.Unlock()
delete(s.mu.objects, name)
}