github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/server/sticky_engine.go (about)

     1  // Copyright 2019 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 server
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    17  	"github.com/cockroachdb/cockroach/pkg/storage"
    18  	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
    19  	"github.com/cockroachdb/cockroach/pkg/util/log"
    20  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    21  	"github.com/cockroachdb/errors"
    22  )
    23  
    24  // stickyInMemEngine extends a normal engine, but does not allow them to be
    25  // closed using the normal Close() method, instead keeping the engines in
    26  // memory until CloseAllStickyInMemEngines is called, hence being "sticky".
    27  // This prevents users of the in memory engine from having to special
    28  // case "sticky" engines on every instance of "Close".
    29  // It is intended for use in demos and/or tests, where we want in-memory
    30  // storage nodes to persist between killed nodes.
    31  type stickyInMemEngine struct {
    32  	// id is the unique identifier for this sticky engine.
    33  	id string
    34  	// closed indicates whether the current engine has been closed.
    35  	closed bool
    36  
    37  	// Engine extends the Engine interface.
    38  	storage.Engine
    39  }
    40  
    41  // stickyInMemEngine implements Engine.
    42  var _ storage.Engine = &stickyInMemEngine{}
    43  
    44  // Close overwrites the default Engine interface to not close the underlying
    45  // engine if called. We mark the state as closed to reflect a correct result
    46  // in Closed().
    47  func (e *stickyInMemEngine) Close() {
    48  	e.closed = true
    49  }
    50  
    51  // Closed overwrites the default Engine interface.
    52  func (e *stickyInMemEngine) Closed() bool {
    53  	return e.closed
    54  }
    55  
    56  // stickyInMemEnginesRegistryImpl is the bookkeeper for all active
    57  // sticky engines, keyed by their id.
    58  type stickyInMemEnginesRegistryImpl struct {
    59  	entries map[string]*stickyInMemEngine
    60  	mu      syncutil.Mutex
    61  }
    62  
    63  var stickyInMemEnginesRegistry = &stickyInMemEnginesRegistryImpl{
    64  	entries: map[string]*stickyInMemEngine{},
    65  }
    66  
    67  // getOrCreateStickyInMemEngine returns an engine associated with the given id.
    68  // It will create a new in-memory engine if one does not already exist.
    69  // At most one engine with a given id can be active in
    70  // "getOrCreateStickyInMemEngine" at any given time.
    71  // Note that if you re-create an existing sticky engine the new attributes
    72  // and cache size will be ignored.
    73  // One must Close() on the sticky engine before another can be fetched.
    74  func getOrCreateStickyInMemEngine(
    75  	ctx context.Context,
    76  	id string,
    77  	engineType enginepb.EngineType,
    78  	attrs roachpb.Attributes,
    79  	cacheSize int64,
    80  ) (storage.Engine, error) {
    81  	stickyInMemEnginesRegistry.mu.Lock()
    82  	defer stickyInMemEnginesRegistry.mu.Unlock()
    83  
    84  	if engine, ok := stickyInMemEnginesRegistry.entries[id]; ok {
    85  		if !engine.closed {
    86  			return nil, errors.Errorf("sticky engine %s has not been closed", id)
    87  		}
    88  
    89  		log.Infof(ctx, "re-using sticky in-mem engine %s", id)
    90  		engine.closed = false
    91  		return engine, nil
    92  	}
    93  
    94  	log.Infof(ctx, "creating new sticky in-mem engine %s", id)
    95  	engine := &stickyInMemEngine{
    96  		id:     id,
    97  		closed: false,
    98  		Engine: storage.NewInMem(ctx, engineType, attrs, cacheSize),
    99  	}
   100  	stickyInMemEnginesRegistry.entries[id] = engine
   101  	return engine, nil
   102  }
   103  
   104  // CloseStickyInMemEngine closes the underlying engine and
   105  // removes the sticky engine keyed by the given id.
   106  // It will error if it does not exist.
   107  func CloseStickyInMemEngine(id string) error {
   108  	stickyInMemEnginesRegistry.mu.Lock()
   109  	defer stickyInMemEnginesRegistry.mu.Unlock()
   110  
   111  	if engine, ok := stickyInMemEnginesRegistry.entries[id]; ok {
   112  		engine.closed = true
   113  		engine.Engine.Close()
   114  		delete(stickyInMemEnginesRegistry.entries, id)
   115  		return nil
   116  	}
   117  	return errors.Errorf("sticky in-mem engine %s does not exist", id)
   118  }
   119  
   120  // CloseAllStickyInMemEngines closes and removes all sticky in memory engines.
   121  func CloseAllStickyInMemEngines() {
   122  	stickyInMemEnginesRegistry.mu.Lock()
   123  	defer stickyInMemEnginesRegistry.mu.Unlock()
   124  
   125  	for _, engine := range stickyInMemEnginesRegistry.entries {
   126  		engine.closed = true
   127  		engine.Engine.Close()
   128  	}
   129  
   130  	for id := range stickyInMemEnginesRegistry.entries {
   131  		delete(stickyInMemEnginesRegistry.entries, id)
   132  	}
   133  }