github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/remotestorage/map_chunk_cache.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  package remotestorage
    16  
    17  import (
    18  	"sync"
    19  
    20  	"github.com/dolthub/dolt/go/store/hash"
    21  	"github.com/dolthub/dolt/go/store/nbs"
    22  )
    23  
    24  // mapChunkCache is a ChunkCache implementation that stores everything in an in memory map.
    25  type mapChunkCache struct {
    26  	mu          *sync.Mutex
    27  	hashToChunk map[hash.Hash]nbs.CompressedChunk
    28  	toFlush     map[hash.Hash]nbs.CompressedChunk
    29  }
    30  
    31  func newMapChunkCache() *mapChunkCache {
    32  	return &mapChunkCache{
    33  		&sync.Mutex{},
    34  		make(map[hash.Hash]nbs.CompressedChunk),
    35  		make(map[hash.Hash]nbs.CompressedChunk),
    36  	}
    37  }
    38  
    39  // Put puts a slice of chunks into the cache.
    40  func (mcc *mapChunkCache) Put(chnks []nbs.CompressedChunk) {
    41  	mcc.mu.Lock()
    42  	defer mcc.mu.Unlock()
    43  
    44  	for i := 0; i < len(chnks); i++ {
    45  		c := chnks[i]
    46  		h := c.Hash()
    47  
    48  		if curr, ok := mcc.hashToChunk[h]; ok {
    49  			if !curr.IsEmpty() {
    50  				continue
    51  			}
    52  		}
    53  
    54  		mcc.hashToChunk[h] = c
    55  
    56  		if !c.IsEmpty() {
    57  			mcc.toFlush[h] = c
    58  		}
    59  	}
    60  }
    61  
    62  // Get gets a map of hash to chunk for a set of hashes.  In the event that a chunk is not in the cache, chunks.Empty.
    63  // is put in it's place
    64  func (mcc *mapChunkCache) Get(hashes hash.HashSet) map[hash.Hash]nbs.CompressedChunk {
    65  	hashToChunk := make(map[hash.Hash]nbs.CompressedChunk)
    66  
    67  	mcc.mu.Lock()
    68  	defer mcc.mu.Unlock()
    69  
    70  	for h := range hashes {
    71  		if c, ok := mcc.hashToChunk[h]; ok {
    72  			hashToChunk[h] = c
    73  		} else {
    74  			hashToChunk[h] = nbs.EmptyCompressedChunk
    75  		}
    76  	}
    77  
    78  	return hashToChunk
    79  }
    80  
    81  // Has takes a set of hashes and returns the set of hashes that the cache currently does not have in it.
    82  func (mcc *mapChunkCache) Has(hashes hash.HashSet) (absent hash.HashSet) {
    83  	absent = make(hash.HashSet)
    84  
    85  	mcc.mu.Lock()
    86  	defer mcc.mu.Unlock()
    87  
    88  	for h := range hashes {
    89  		if _, ok := mcc.hashToChunk[h]; !ok {
    90  			absent[h] = struct{}{}
    91  		}
    92  	}
    93  
    94  	return absent
    95  }
    96  
    97  // PutChunk puts a single chunk in the cache.  true returns in the event that the chunk was cached successfully
    98  // and false is returned if that chunk is already is the cache.
    99  func (mcc *mapChunkCache) PutChunk(ch nbs.CompressedChunk) bool {
   100  	mcc.mu.Lock()
   101  	defer mcc.mu.Unlock()
   102  
   103  	h := ch.Hash()
   104  	if existing, ok := mcc.hashToChunk[h]; !ok || existing.IsEmpty() {
   105  		mcc.hashToChunk[h] = ch
   106  		mcc.toFlush[h] = ch
   107  		return true
   108  	}
   109  
   110  	return false
   111  }
   112  
   113  // GetAndClearChunksToFlush gets a map of hash to chunk which includes all the chunks that were put in the cache
   114  // between the last time GetAndClearChunksToFlush was called and now.
   115  func (mcc *mapChunkCache) GetAndClearChunksToFlush() map[hash.Hash]nbs.CompressedChunk {
   116  	newToFlush := make(map[hash.Hash]nbs.CompressedChunk)
   117  
   118  	mcc.mu.Lock()
   119  	defer mcc.mu.Unlock()
   120  
   121  	toFlush := mcc.toFlush
   122  	mcc.toFlush = newToFlush
   123  
   124  	return toFlush
   125  }