github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/chunks/memory_store.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2016 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package chunks
    23  
    24  import (
    25  	"context"
    26  	"fmt"
    27  	"sync"
    28  
    29  	"github.com/dolthub/dolt/go/store/constants"
    30  	"github.com/dolthub/dolt/go/store/d"
    31  	"github.com/dolthub/dolt/go/store/hash"
    32  )
    33  
    34  // MemoryStorage provides a "persistent" storage layer to back multiple
    35  // MemoryStoreViews. A MemoryStorage instance holds the ground truth for the
    36  // root and set of chunks that are visible to all MemoryStoreViews vended by
    37  // NewView(), allowing them to implement the transaction-style semantics that
    38  // ChunkStore requires.
    39  type MemoryStorage struct {
    40  	data     map[hash.Hash]Chunk
    41  	rootHash hash.Hash
    42  	mu       sync.RWMutex
    43  	version  string
    44  }
    45  
    46  // NewView vends a MemoryStoreView backed by this MemoryStorage. It's
    47  // initialized with the currently "persisted" root.
    48  func (ms *MemoryStorage) NewView() ChunkStore {
    49  	version := ms.version
    50  	if version == "" {
    51  		version = constants.Format718String
    52  	}
    53  
    54  	return &MemoryStoreView{storage: ms, rootHash: ms.rootHash, version: version}
    55  }
    56  
    57  // NewViewWithVersion vends a MemoryStoreView backed by this MemoryStorage. It's
    58  // initialized with the currently "persisted" root. Uses the default format.
    59  func (ms *MemoryStorage) NewViewWithDefaultFormat() ChunkStore {
    60  	return &MemoryStoreView{storage: ms, rootHash: ms.rootHash, version: constants.FormatDefaultString}
    61  }
    62  
    63  // Get retrieves the Chunk with the Hash h, returning EmptyChunk if it's not
    64  // present.
    65  func (ms *MemoryStorage) Get(ctx context.Context, h hash.Hash) (Chunk, error) {
    66  	ms.mu.RLock()
    67  	defer ms.mu.RUnlock()
    68  	if c, ok := ms.data[h]; ok {
    69  		return c, nil
    70  	}
    71  	return EmptyChunk, nil
    72  }
    73  
    74  // Has returns true if the Chunk with the Hash h is present in ms.data, false
    75  // if not.
    76  func (ms *MemoryStorage) Has(ctx context.Context, r hash.Hash) (bool, error) {
    77  	ms.mu.RLock()
    78  	defer ms.mu.RUnlock()
    79  	_, ok := ms.data[r]
    80  	return ok, nil
    81  }
    82  
    83  // Len returns the number of Chunks in ms.data.
    84  func (ms *MemoryStorage) Len() int {
    85  	ms.mu.RLock()
    86  	defer ms.mu.RUnlock()
    87  	return len(ms.data)
    88  }
    89  
    90  // Root returns the currently "persisted" root hash of this in-memory store.
    91  func (ms *MemoryStorage) Root(ctx context.Context) hash.Hash {
    92  	ms.mu.RLock()
    93  	defer ms.mu.RUnlock()
    94  	return ms.rootHash
    95  }
    96  
    97  // Update checks the "persisted" root against last and, iff it matches,
    98  // updates the root to current, adds all of novel to ms.data, and returns
    99  // true. Otherwise returns false.
   100  func (ms *MemoryStorage) Update(current, last hash.Hash, novel map[hash.Hash]Chunk) bool {
   101  	ms.mu.Lock()
   102  	defer ms.mu.Unlock()
   103  	if last != ms.rootHash {
   104  		return false
   105  	}
   106  	if ms.data == nil {
   107  		ms.data = map[hash.Hash]Chunk{}
   108  	}
   109  	for h, c := range novel {
   110  		ms.data[h] = c
   111  	}
   112  	ms.rootHash = current
   113  	return true
   114  }
   115  
   116  // MemoryStoreView is an in-memory implementation of store.ChunkStore. Useful
   117  // mainly for tests.
   118  // The proper way to get one:
   119  // storage := &MemoryStorage{}
   120  // ms := storage.NewView()
   121  type MemoryStoreView struct {
   122  	pending  map[hash.Hash]Chunk
   123  	rootHash hash.Hash
   124  	mu       sync.RWMutex
   125  	version  string
   126  
   127  	storage *MemoryStorage
   128  }
   129  
   130  var _ ChunkStore = &MemoryStoreView{}
   131  var _ ChunkStoreGarbageCollector = &MemoryStoreView{}
   132  
   133  func (ms *MemoryStoreView) Get(ctx context.Context, h hash.Hash) (Chunk, error) {
   134  	ms.mu.RLock()
   135  	defer ms.mu.RUnlock()
   136  	if c, ok := ms.pending[h]; ok {
   137  		return c, nil
   138  	}
   139  	return ms.storage.Get(ctx, h)
   140  }
   141  
   142  func (ms *MemoryStoreView) GetMany(ctx context.Context, hashes hash.HashSet, found func(*Chunk)) error {
   143  	for h := range hashes {
   144  		c, err := ms.Get(ctx, h)
   145  
   146  		if err != nil {
   147  			return err
   148  		}
   149  
   150  		if !c.IsEmpty() {
   151  			found(&c)
   152  		}
   153  	}
   154  
   155  	return nil
   156  }
   157  
   158  func (ms *MemoryStoreView) Has(ctx context.Context, h hash.Hash) (bool, error) {
   159  	ms.mu.RLock()
   160  	defer ms.mu.RUnlock()
   161  	if _, ok := ms.pending[h]; ok {
   162  		return true, nil
   163  	}
   164  	return ms.storage.Has(ctx, h)
   165  }
   166  
   167  func (ms *MemoryStoreView) HasMany(ctx context.Context, hashes hash.HashSet) (hash.HashSet, error) {
   168  	absent := hash.HashSet{}
   169  	for h := range hashes {
   170  		exists, err := ms.Has(ctx, h)
   171  		if err != nil {
   172  			return nil, err
   173  		} else if !exists {
   174  			absent.Insert(h)
   175  		}
   176  	}
   177  	return absent, nil
   178  }
   179  
   180  func (ms *MemoryStoreView) Version() string {
   181  	return ms.version
   182  }
   183  
   184  func (ms *MemoryStoreView) Put(ctx context.Context, c Chunk) error {
   185  	ms.mu.Lock()
   186  	defer ms.mu.Unlock()
   187  	if ms.pending == nil {
   188  		ms.pending = map[hash.Hash]Chunk{}
   189  	}
   190  	ms.pending[c.Hash()] = c
   191  
   192  	return nil
   193  }
   194  
   195  func (ms *MemoryStoreView) Len() int {
   196  	ms.mu.RLock()
   197  	defer ms.mu.RUnlock()
   198  	return len(ms.pending) + ms.storage.Len()
   199  }
   200  
   201  func (ms *MemoryStoreView) Rebase(ctx context.Context) error {
   202  	ms.mu.Lock()
   203  	defer ms.mu.Unlock()
   204  	ms.rootHash = ms.storage.Root(ctx)
   205  
   206  	return nil
   207  }
   208  
   209  func (ms *MemoryStoreView) Root(ctx context.Context) (hash.Hash, error) {
   210  	ms.mu.RLock()
   211  	defer ms.mu.RUnlock()
   212  	return ms.rootHash, nil
   213  }
   214  
   215  func (ms *MemoryStoreView) Commit(ctx context.Context, current, last hash.Hash) (bool, error) {
   216  	ms.mu.Lock()
   217  	defer ms.mu.Unlock()
   218  	if last != ms.rootHash {
   219  		return false, nil
   220  	}
   221  
   222  	success := ms.storage.Update(current, last, ms.pending)
   223  	if success {
   224  		ms.pending = nil
   225  	}
   226  	ms.rootHash = ms.storage.Root(ctx)
   227  	return success, nil
   228  }
   229  
   230  func (ms *MemoryStoreView) MarkAndSweepChunks(ctx context.Context, last hash.Hash, keepChunks <-chan []hash.Hash) error {
   231  	if last != ms.rootHash {
   232  		return fmt.Errorf("last does not match ms.Root()")
   233  	}
   234  
   235  	keepers := make(map[hash.Hash]Chunk, ms.storage.Len())
   236  
   237  LOOP:
   238  	for {
   239  		select {
   240  		case hs, ok := <-keepChunks:
   241  			if !ok {
   242  				break LOOP
   243  			}
   244  			for _, h := range hs {
   245  				c, err := ms.Get(ctx, h)
   246  				if err != nil {
   247  					return err
   248  				}
   249  				keepers[h] = c
   250  			}
   251  		case <-ctx.Done():
   252  			return ctx.Err()
   253  		}
   254  	}
   255  
   256  	ms.storage = &MemoryStorage{rootHash: ms.rootHash, data: keepers}
   257  	ms.pending = map[hash.Hash]Chunk{}
   258  	return nil
   259  }
   260  
   261  func (ms *MemoryStoreView) Stats() interface{} {
   262  	return nil
   263  }
   264  
   265  func (ms *MemoryStoreView) StatsSummary() string {
   266  	return "Unsupported"
   267  }
   268  
   269  func (ms *MemoryStoreView) Close() error {
   270  	return nil
   271  }
   272  
   273  type memoryStoreFactory struct {
   274  	stores map[string]*MemoryStorage
   275  	mu     *sync.Mutex
   276  }
   277  
   278  func NewMemoryStoreFactory() *memoryStoreFactory {
   279  	return &memoryStoreFactory{map[string]*MemoryStorage{}, &sync.Mutex{}}
   280  }
   281  
   282  func (f *memoryStoreFactory) CreateStoreFromCache(ctx context.Context, ns string) ChunkStore {
   283  	return f.CreateStore(ctx, ns)
   284  }
   285  
   286  func (f *memoryStoreFactory) CreateStore(ctx context.Context, ns string) ChunkStore {
   287  	f.mu.Lock()
   288  	defer f.mu.Unlock()
   289  
   290  	if f.stores == nil {
   291  		d.Panic("Cannot use memoryStoreFactory after Shutter().")
   292  	}
   293  	if ms, present := f.stores[ns]; present {
   294  		return ms.NewView()
   295  	}
   296  	f.stores[ns] = &MemoryStorage{}
   297  	return f.stores[ns].NewView()
   298  }
   299  
   300  func (f *memoryStoreFactory) Shutter() {
   301  	f.stores = nil
   302  }