You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
260 lines
6.0 KiB
260 lines
6.0 KiB
// Copyright 2015 PingCAP, Inc.
|
|
//
|
|
// Copyright 2015 Wenbin Xiao
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package kv
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"sync/atomic"
|
|
|
|
"github.com/pingcap/errors"
|
|
"github.com/pingcap/tidb/kv/memdb"
|
|
)
|
|
|
|
// memDbBuffer implements the MemBuffer interface.
|
|
type memDbBuffer struct {
|
|
sandbox *memdb.Sandbox
|
|
entrySizeLimit int
|
|
bufferLenLimit uint64
|
|
bufferSizeLimit uint64
|
|
}
|
|
|
|
type memDbIter struct {
|
|
iter memdb.Iterator
|
|
start []byte
|
|
end []byte
|
|
reverse bool
|
|
}
|
|
|
|
// NewMemDbBuffer creates a new memDbBuffer.
|
|
func NewMemDbBuffer() MemBuffer {
|
|
return &memDbBuffer{
|
|
sandbox: memdb.NewSandbox(),
|
|
entrySizeLimit: TxnEntrySizeLimit,
|
|
bufferSizeLimit: atomic.LoadUint64(&TxnTotalSizeLimit),
|
|
}
|
|
}
|
|
|
|
// Iter creates an Iterator.
|
|
func (m *memDbBuffer) Iter(k Key, upperBound Key) (Iterator, error) {
|
|
i := &memDbIter{
|
|
iter: m.sandbox.NewIterator(),
|
|
start: k,
|
|
end: upperBound,
|
|
reverse: false,
|
|
}
|
|
|
|
if k == nil {
|
|
i.iter.SeekToFirst()
|
|
} else {
|
|
i.iter.Seek(k)
|
|
}
|
|
return i, nil
|
|
}
|
|
|
|
func (m *memDbBuffer) IterReverse(k Key) (Iterator, error) {
|
|
i := &memDbIter{
|
|
iter: m.sandbox.NewIterator(),
|
|
end: k,
|
|
reverse: true,
|
|
}
|
|
if k == nil {
|
|
i.iter.SeekToLast()
|
|
} else {
|
|
i.iter.SeekForExclusivePrev(k)
|
|
}
|
|
return i, nil
|
|
}
|
|
|
|
// Get returns the value associated with key.
|
|
func (m *memDbBuffer) Get(ctx context.Context, k Key) ([]byte, error) {
|
|
v := m.sandbox.Get(k)
|
|
if v == nil {
|
|
return nil, ErrNotExist
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
// GetFlags returns the flags associated with key.
|
|
func (m *memDbBuffer) GetFlags(ctx context.Context, k Key) memdb.KeyFlags {
|
|
return m.sandbox.GetFlags(k)
|
|
}
|
|
|
|
// Set associates key with value.
|
|
func (m *memDbBuffer) Set(k Key, v []byte) error {
|
|
if len(v) == 0 {
|
|
return errors.Trace(ErrCannotSetNilValue)
|
|
}
|
|
if len(k)+len(v) > m.entrySizeLimit {
|
|
return ErrEntryTooLarge.GenWithStackByArgs(m.entrySizeLimit, len(k)+len(v))
|
|
}
|
|
|
|
m.sandbox.Put(k, v)
|
|
if m.Size() > int(m.bufferSizeLimit) {
|
|
return ErrTxnTooLarge.GenWithStackByArgs(atomic.LoadUint64(&TxnTotalSizeLimit))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Delete removes the entry from buffer with provided key.
|
|
func (m *memDbBuffer) Delete(k Key) error {
|
|
m.sandbox.Put(k, nil)
|
|
if m.Size() > int(m.bufferSizeLimit) {
|
|
return ErrTxnTooLarge.GenWithStackByArgs(atomic.LoadUint64(&TxnTotalSizeLimit))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// DeleteWithNeedLock delete key with a need lock mark
|
|
func (m *memDbBuffer) DeleteWithNeedLock(k Key) error {
|
|
m.sandbox.PutWithNeedLock(k, nil)
|
|
if m.Size() > int(m.bufferSizeLimit) {
|
|
return ErrTxnTooLarge.GenWithStackByArgs(atomic.LoadUint64(&TxnTotalSizeLimit))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Size returns sum of keys and values length.
|
|
func (m *memDbBuffer) Size() int {
|
|
return m.sandbox.Size()
|
|
}
|
|
|
|
// Len returns the number of entries in the DB.
|
|
func (m *memDbBuffer) Len() int {
|
|
return m.sandbox.Len()
|
|
}
|
|
|
|
func (m *memDbBuffer) NewStagingBuffer() MemBuffer {
|
|
return &memDbBuffer{
|
|
sandbox: m.sandbox.Derive(),
|
|
entrySizeLimit: TxnEntrySizeLimit,
|
|
bufferSizeLimit: m.bufferSizeLimit - uint64(m.sandbox.Size()),
|
|
}
|
|
}
|
|
|
|
func (m *memDbBuffer) Flush() (int, error) {
|
|
// There is no need to check size limit,
|
|
// because the size limit is maintain when derive this buffer.
|
|
return m.sandbox.Flush(), nil
|
|
}
|
|
|
|
func (m *memDbBuffer) Discard() {
|
|
m.sandbox.Discard()
|
|
}
|
|
|
|
// Next implements the Iterator Next.
|
|
func (i *memDbIter) Next() error {
|
|
if i.reverse {
|
|
i.iter.Prev()
|
|
} else {
|
|
i.iter.Next()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Valid implements the Iterator Valid.
|
|
func (i *memDbIter) Valid() bool {
|
|
if !i.reverse {
|
|
return i.iter.Valid() && (i.end == nil || bytes.Compare(i.Key(), i.end) < 0)
|
|
}
|
|
return i.iter.Valid()
|
|
}
|
|
|
|
// Key implements the Iterator Key.
|
|
func (i *memDbIter) Key() Key {
|
|
return i.iter.Key()
|
|
}
|
|
|
|
// Value implements the Iterator Value.
|
|
func (i *memDbIter) Value() []byte {
|
|
return i.iter.Value()
|
|
}
|
|
|
|
// Close Implements the Iterator Close.
|
|
func (i *memDbIter) Close() {
|
|
|
|
}
|
|
|
|
// WalkMemBuffer iterates all buffered kv pairs in memBuf
|
|
func WalkMemBuffer(memBuf MemBuffer, f func(k Key, v []byte) error) error {
|
|
iter, err := memBuf.Iter(nil, nil)
|
|
if err != nil {
|
|
return errors.Trace(err)
|
|
}
|
|
|
|
defer iter.Close()
|
|
for iter.Valid() {
|
|
if err = f(iter.Key(), iter.Value()); err != nil {
|
|
return errors.Trace(err)
|
|
}
|
|
err = iter.Next()
|
|
if err != nil {
|
|
return errors.Trace(err)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// BufferBatchGetter is the type for BatchGet with MemBuffer.
|
|
type BufferBatchGetter struct {
|
|
buffer MemBuffer
|
|
middle Getter
|
|
snapshot Snapshot
|
|
}
|
|
|
|
// NewBufferBatchGetter creates a new BufferBatchGetter.
|
|
func NewBufferBatchGetter(buffer MemBuffer, middleCache Getter, snapshot Snapshot) *BufferBatchGetter {
|
|
return &BufferBatchGetter{buffer: buffer, middle: middleCache, snapshot: snapshot}
|
|
}
|
|
|
|
// BatchGet implements the BatchGetter interface.
|
|
func (b *BufferBatchGetter) BatchGet(ctx context.Context, keys []Key) (map[string][]byte, error) {
|
|
if b.buffer.Len() == 0 {
|
|
return b.snapshot.BatchGet(ctx, keys)
|
|
}
|
|
bufferValues := make([][]byte, len(keys))
|
|
shrinkKeys := make([]Key, 0, len(keys))
|
|
for i, key := range keys {
|
|
val, err := b.buffer.Get(ctx, key)
|
|
if err == nil {
|
|
bufferValues[i] = val
|
|
continue
|
|
}
|
|
if !IsErrNotFound(err) {
|
|
return nil, errors.Trace(err)
|
|
}
|
|
if b.middle != nil {
|
|
val, err = b.middle.Get(ctx, key)
|
|
if err == nil {
|
|
bufferValues[i] = val
|
|
continue
|
|
}
|
|
}
|
|
shrinkKeys = append(shrinkKeys, key)
|
|
}
|
|
storageValues, err := b.snapshot.BatchGet(ctx, shrinkKeys)
|
|
if err != nil {
|
|
return nil, errors.Trace(err)
|
|
}
|
|
for i, key := range keys {
|
|
if len(bufferValues[i]) == 0 {
|
|
continue
|
|
}
|
|
storageValues[string(key)] = bufferValues[i]
|
|
}
|
|
return storageValues, nil
|
|
}
|
|
|