github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/objstorage/remote/mem.go (about)

     1  // Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
     2  // of this source code is governed by a BSD-style license that can be found in
     3  // the LICENSE file.
     4  
     5  package remote
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"io"
    11  	"os"
    12  	"strings"
    13  	"sync"
    14  )
    15  
    16  // NewInMem returns an in-memory implementation of the remote.Storage
    17  // interface (for testing).
    18  func NewInMem() Storage {
    19  	store := &inMemStore{}
    20  	store.mu.objects = make(map[string]*inMemObj)
    21  	return store
    22  }
    23  
    24  // inMemStore is an in-memory implementation of the remote.Storage interface
    25  // (for testing).
    26  type inMemStore struct {
    27  	mu struct {
    28  		sync.Mutex
    29  		objects map[string]*inMemObj
    30  	}
    31  }
    32  
    33  var _ Storage = (*inMemStore)(nil)
    34  
    35  type inMemObj struct {
    36  	name string
    37  	data []byte
    38  }
    39  
    40  func (s *inMemStore) Close() error {
    41  	*s = inMemStore{}
    42  	return nil
    43  }
    44  
    45  func (s *inMemStore) ReadObject(
    46  	ctx context.Context, objName string,
    47  ) (_ ObjectReader, objSize int64, _ error) {
    48  	obj, err := s.getObj(objName)
    49  	if err != nil {
    50  		return nil, 0, err
    51  	}
    52  	return &inMemReader{data: obj.data}, int64(len(obj.data)), nil
    53  }
    54  
    55  type inMemReader struct {
    56  	data []byte
    57  }
    58  
    59  var _ ObjectReader = (*inMemReader)(nil)
    60  
    61  func (r *inMemReader) ReadAt(ctx context.Context, p []byte, offset int64) error {
    62  	if offset+int64(len(p)) > int64(len(r.data)) {
    63  		return io.EOF
    64  	}
    65  	copy(p, r.data[offset:])
    66  	return nil
    67  }
    68  
    69  func (r *inMemReader) Close() error {
    70  	r.data = nil
    71  	return nil
    72  }
    73  
    74  func (s *inMemStore) CreateObject(objName string) (io.WriteCloser, error) {
    75  	return &inMemWriter{
    76  		store: s,
    77  		name:  objName,
    78  	}, nil
    79  }
    80  
    81  type inMemWriter struct {
    82  	store *inMemStore
    83  	name  string
    84  	buf   bytes.Buffer
    85  }
    86  
    87  var _ io.WriteCloser = (*inMemWriter)(nil)
    88  
    89  func (o *inMemWriter) Write(p []byte) (n int, err error) {
    90  	if o.store == nil {
    91  		panic("Write after Close")
    92  	}
    93  	return o.buf.Write(p)
    94  }
    95  
    96  func (o *inMemWriter) Close() error {
    97  	if o.store != nil {
    98  		o.store.addObj(&inMemObj{
    99  			name: o.name,
   100  			data: o.buf.Bytes(),
   101  		})
   102  		o.store = nil
   103  	}
   104  	return nil
   105  }
   106  
   107  func (s *inMemStore) List(prefix, delimiter string) ([]string, error) {
   108  	if delimiter != "" {
   109  		panic("delimiter unimplemented")
   110  	}
   111  
   112  	s.mu.Lock()
   113  	defer s.mu.Unlock()
   114  	res := make([]string, 0, len(s.mu.objects))
   115  	for name := range s.mu.objects {
   116  		if strings.HasPrefix(name, prefix) {
   117  			res = append(res, name)
   118  		}
   119  	}
   120  	return res, nil
   121  }
   122  
   123  func (s *inMemStore) Delete(objName string) error {
   124  	s.rmObj(objName)
   125  	return nil
   126  }
   127  
   128  // Size returns the length of the named object in bytesWritten.
   129  func (s *inMemStore) Size(objName string) (int64, error) {
   130  	obj, err := s.getObj(objName)
   131  	if err != nil {
   132  		return 0, err
   133  	}
   134  	return int64(len(obj.data)), nil
   135  }
   136  
   137  func (s *inMemStore) IsNotExistError(err error) bool {
   138  	return err == os.ErrNotExist
   139  }
   140  
   141  func (s *inMemStore) getObj(name string) (*inMemObj, error) {
   142  	s.mu.Lock()
   143  	defer s.mu.Unlock()
   144  	obj, ok := s.mu.objects[name]
   145  	if !ok {
   146  		return nil, os.ErrNotExist
   147  	}
   148  	return obj, nil
   149  }
   150  
   151  func (s *inMemStore) addObj(o *inMemObj) {
   152  	s.mu.Lock()
   153  	defer s.mu.Unlock()
   154  	s.mu.objects[o.name] = o
   155  }
   156  
   157  func (s *inMemStore) rmObj(name string) {
   158  	s.mu.Lock()
   159  	defer s.mu.Unlock()
   160  	delete(s.mu.objects, name)
   161  }