Line data    Source code

       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           2 : func NewInMem() Storage {
      19           2 :         store := &inMemStore{}
      20           2 :         store.mu.objects = make(map[string]*inMemObj)
      21           2 :         return store
      22           2 : }
      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           0 : func (s *inMemStore) Close() error {
      41           0 :         *s = inMemStore{}
      42           0 :         return nil
      43           0 : }
      44             : 
      45             : func (s *inMemStore) ReadObject(
      46             :         ctx context.Context, objName string,
      47           2 : ) (_ ObjectReader, objSize int64, _ error) {
      48           2 :         obj, err := s.getObj(objName)
      49           2 :         if err != nil {
      50           1 :                 return nil, 0, err
      51           1 :         }
      52           2 :         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           2 : func (r *inMemReader) ReadAt(ctx context.Context, p []byte, offset int64) error {
      62           2 :         if offset+int64(len(p)) > int64(len(r.data)) {
      63           0 :                 return io.EOF
      64           0 :         }
      65           2 :         copy(p, r.data[offset:])
      66           2 :         return nil
      67             : }
      68             : 
      69           2 : func (r *inMemReader) Close() error {
      70           2 :         r.data = nil
      71           2 :         return nil
      72           2 : }
      73             : 
      74           2 : func (s *inMemStore) CreateObject(objName string) (io.WriteCloser, error) {
      75           2 :         return &inMemWriter{
      76           2 :                 store: s,
      77           2 :                 name:  objName,
      78           2 :         }, nil
      79           2 : }
      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           2 : func (o *inMemWriter) Write(p []byte) (n int, err error) {
      90           2 :         if o.store == nil {
      91           0 :                 panic("Write after Close")
      92             :         }
      93           2 :         return o.buf.Write(p)
      94             : }
      95             : 
      96           2 : func (o *inMemWriter) Close() error {
      97           2 :         if o.store != nil {
      98           2 :                 o.store.addObj(&inMemObj{
      99           2 :                         name: o.name,
     100           2 :                         data: o.buf.Bytes(),
     101           2 :                 })
     102           2 :                 o.store = nil
     103           2 :         }
     104           2 :         return nil
     105             : }
     106             : 
     107           2 : func (s *inMemStore) List(prefix, delimiter string) ([]string, error) {
     108           2 :         if delimiter != "" {
     109           0 :                 panic("delimiter unimplemented")
     110             :         }
     111             : 
     112           2 :         s.mu.Lock()
     113           2 :         defer s.mu.Unlock()
     114           2 :         res := make([]string, 0, len(s.mu.objects))
     115           2 :         for name := range s.mu.objects {
     116           2 :                 if strings.HasPrefix(name, prefix) {
     117           2 :                         res = append(res, name)
     118           2 :                 }
     119             :         }
     120           2 :         return res, nil
     121             : }
     122             : 
     123           2 : func (s *inMemStore) Delete(objName string) error {
     124           2 :         s.rmObj(objName)
     125           2 :         return nil
     126           2 : }
     127             : 
     128             : // Size returns the length of the named object in bytesWritten.
     129           2 : func (s *inMemStore) Size(objName string) (int64, error) {
     130           2 :         obj, err := s.getObj(objName)
     131           2 :         if err != nil {
     132           1 :                 return 0, err
     133           1 :         }
     134           2 :         return int64(len(obj.data)), nil
     135             : }
     136             : 
     137           1 : func (s *inMemStore) IsNotExistError(err error) bool {
     138           1 :         return err == os.ErrNotExist
     139           1 : }
     140             : 
     141           2 : func (s *inMemStore) getObj(name string) (*inMemObj, error) {
     142           2 :         s.mu.Lock()
     143           2 :         defer s.mu.Unlock()
     144           2 :         obj, ok := s.mu.objects[name]
     145           2 :         if !ok {
     146           1 :                 return nil, os.ErrNotExist
     147           1 :         }
     148           2 :         return obj, nil
     149             : }
     150             : 
     151           2 : func (s *inMemStore) addObj(o *inMemObj) {
     152           2 :         s.mu.Lock()
     153           2 :         defer s.mu.Unlock()
     154           2 :         s.mu.objects[o.name] = o
     155           2 : }
     156             : 
     157           2 : func (s *inMemStore) rmObj(name string) {
     158           2 :         s.mu.Lock()
     159           2 :         defer s.mu.Unlock()
     160           2 :         delete(s.mu.objects, name)
     161           2 : }