github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/replica_sideload_inmem.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package kvserver
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"path/filepath"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    20  	"github.com/cockroachdb/cockroach/pkg/storage"
    21  )
    22  
    23  type slKey struct {
    24  	index, term uint64
    25  }
    26  
    27  type inMemSideloadStorage struct {
    28  	m      map[slKey][]byte
    29  	prefix string
    30  }
    31  
    32  func mustNewInMemSideloadStorage(
    33  	rangeID roachpb.RangeID, replicaID roachpb.ReplicaID, baseDir string,
    34  ) SideloadStorage {
    35  	ss, err := newInMemSideloadStorage(cluster.MakeTestingClusterSettings(), rangeID, replicaID, baseDir, nil)
    36  	if err != nil {
    37  		panic(err)
    38  	}
    39  	return ss
    40  }
    41  
    42  func newInMemSideloadStorage(
    43  	_ *cluster.Settings,
    44  	rangeID roachpb.RangeID,
    45  	replicaID roachpb.ReplicaID,
    46  	baseDir string,
    47  	eng storage.Engine,
    48  ) (SideloadStorage, error) {
    49  	return &inMemSideloadStorage{
    50  		prefix: filepath.Join(baseDir, fmt.Sprintf("%d.%d", rangeID, replicaID)),
    51  		m:      make(map[slKey][]byte),
    52  	}, nil
    53  }
    54  
    55  func (ss *inMemSideloadStorage) key(index, term uint64) slKey {
    56  	return slKey{index: index, term: term}
    57  }
    58  
    59  func (ss *inMemSideloadStorage) Dir() string {
    60  	// We could return ss.prefix but real code calling this would then take the
    61  	// result in look for it on the actual file system.
    62  	panic("unsupported")
    63  }
    64  
    65  func (ss *inMemSideloadStorage) Put(_ context.Context, index, term uint64, contents []byte) error {
    66  	key := ss.key(index, term)
    67  	ss.m[key] = contents
    68  	return nil
    69  }
    70  
    71  func (ss *inMemSideloadStorage) Get(_ context.Context, index, term uint64) ([]byte, error) {
    72  	key := ss.key(index, term)
    73  	data, ok := ss.m[key]
    74  	if !ok {
    75  		return nil, errSideloadedFileNotFound
    76  	}
    77  	return data, nil
    78  }
    79  
    80  func (ss *inMemSideloadStorage) Filename(_ context.Context, index, term uint64) (string, error) {
    81  	return filepath.Join(ss.prefix, fmt.Sprintf("i%d.t%d", index, term)), nil
    82  }
    83  
    84  func (ss *inMemSideloadStorage) Purge(_ context.Context, index, term uint64) (int64, error) {
    85  	k := ss.key(index, term)
    86  	if _, ok := ss.m[k]; !ok {
    87  		return 0, errSideloadedFileNotFound
    88  	}
    89  	size := int64(len(ss.m[k]))
    90  	delete(ss.m, k)
    91  	return size, nil
    92  }
    93  
    94  func (ss *inMemSideloadStorage) Clear(_ context.Context) error {
    95  	ss.m = make(map[slKey][]byte)
    96  	return nil
    97  }
    98  
    99  func (ss *inMemSideloadStorage) TruncateTo(
   100  	_ context.Context, index uint64,
   101  ) (freed, retained int64, _ error) {
   102  	// Not efficient, but this storage is for testing purposes only anyway.
   103  	for k, v := range ss.m {
   104  		if k.index < index {
   105  			freed += int64(len(v))
   106  			delete(ss.m, k)
   107  		} else {
   108  			retained += int64(len(v))
   109  		}
   110  	}
   111  	return freed, retained, nil
   112  }