github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/storage/union_store.go (about)

     1  // Copyright 2022 zGraph Authors. All rights reserved.
     2  //
     3  // Copyright 2021 PingCAP, Inc.
     4  //
     5  // Licensed under the Apache License, Version 2.0 (the "License");
     6  // you may not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //     http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  package storage
    18  
    19  import (
    20  	"context"
    21  
    22  	"github.com/pingcap/errors"
    23  	"github.com/vescale/zgraph/storage/kv"
    24  )
    25  
    26  // UnionStore is an in-memory Store which contains a buffer for write and a
    27  // snapshot for read.
    28  type UnionStore struct {
    29  	memBuffer *MemDB
    30  	snapshot  kv.Snapshot
    31  }
    32  
    33  // NewUnionStore builds a new unionStore.
    34  func NewUnionStore(snapshot kv.Snapshot) *UnionStore {
    35  	return &UnionStore{
    36  		snapshot:  snapshot,
    37  		memBuffer: newMemDB(),
    38  	}
    39  }
    40  
    41  // MemBuffer return the MemBuffer binding to this unionStore.
    42  func (us *UnionStore) MemBuffer() *MemDB {
    43  	return us.memBuffer
    44  }
    45  
    46  // Get implements the Retriever interface.
    47  func (us *UnionStore) Get(ctx context.Context, k kv.Key) ([]byte, error) {
    48  	v, err := us.memBuffer.Get(ctx, k)
    49  	if errors.Cause(err) == kv.ErrNotExist {
    50  		v, err = us.snapshot.Get(ctx, k)
    51  	}
    52  	if err != nil {
    53  		return v, err
    54  	}
    55  	if len(v) == 0 {
    56  		return nil, kv.ErrNotExist
    57  	}
    58  	return v, nil
    59  }
    60  
    61  // Iter implements the Retriever interface.
    62  func (us *UnionStore) Iter(lowerBound, upperBound kv.Key) (kv.Iterator, error) {
    63  	bufferIt, err := us.memBuffer.Iter(lowerBound, upperBound)
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  	retrieverIt, err := us.snapshot.Iter(lowerBound, upperBound)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	return NewUnionIter(bufferIt, retrieverIt, false)
    72  }
    73  
    74  // IterReverse implements the Retriever interface.
    75  func (us *UnionStore) IterReverse(lowerBound, upperBound kv.Key) (kv.Iterator, error) {
    76  	bufferIt, err := us.memBuffer.IterReverse(lowerBound, upperBound)
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  	retrieverIt, err := us.snapshot.IterReverse(lowerBound, upperBound)
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  	return NewUnionIter(bufferIt, retrieverIt, true)
    85  }
    86  
    87  // HasPresumeKeyNotExists gets the key exist error info for the lazy check.
    88  func (us *UnionStore) HasPresumeKeyNotExists(k []byte) bool {
    89  	flags, err := us.memBuffer.GetFlags(k)
    90  	if err != nil {
    91  		return false
    92  	}
    93  	return flags.HasPresumeKeyNotExists()
    94  }
    95  
    96  // UnmarkPresumeKeyNotExists deletes the key exist error info for the lazy check.
    97  func (us *UnionStore) UnmarkPresumeKeyNotExists(k []byte) {
    98  	us.memBuffer.UpdateFlags(k, kv.DelPresumeKeyNotExists)
    99  }
   100  
   101  // SetEntrySizeLimit sets the size limit for each entry and total buffer.
   102  func (us *UnionStore) SetEntrySizeLimit(entryLimit, bufferLimit uint64) {
   103  	us.memBuffer.entrySizeLimit = entryLimit
   104  	us.memBuffer.bufferSizeLimit = bufferLimit
   105  }