mirror of
https://github.com/gravitational/teleport.git
synced 2026-11-01 05:27:24 +08:00
rosstimothy
93d7831ced
`btree.BTree` has been a backward compatability wrapper over a `btree.BTreeG` since `github.com/google/btree@v1.1.0`. This converts all usages to the new generic variant.
508 lines
12 KiB
Go
508 lines
12 KiB
Go
/*
|
|
Copyright 2019 Gravitational, Inc.
|
|
|
|
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,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
*/
|
|
|
|
package memory
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/google/btree"
|
|
"github.com/gravitational/trace"
|
|
"github.com/jonboulle/clockwork"
|
|
log "github.com/sirupsen/logrus"
|
|
|
|
"github.com/gravitational/teleport"
|
|
"github.com/gravitational/teleport/api/types"
|
|
"github.com/gravitational/teleport/lib/backend"
|
|
)
|
|
|
|
// GetName is a part of backend API and it returns in-memory backend type
|
|
// as it appears in `storage/type` section of Teleport YAML
|
|
func GetName() string {
|
|
return "in-memory"
|
|
}
|
|
|
|
const (
|
|
// defaultBTreeDegree is a default degree of a B-Tree
|
|
defaultBTreeDegree = 8
|
|
)
|
|
|
|
// Config holds configuration for the backend
|
|
type Config struct {
|
|
// Context is a context for opening the
|
|
// database
|
|
Context context.Context
|
|
// BTreeDegree is a degree of B-Tree, 2 for example, will create a
|
|
// 2-3-4 tree (each node contains 1-3 items and 2-4 children).
|
|
BTreeDegree int
|
|
// Clock is a clock for time-related operations
|
|
Clock clockwork.Clock
|
|
// Component is a logging component
|
|
Component string
|
|
// EventsOff turns off events generation
|
|
EventsOff bool
|
|
// BufferSize sets up event buffer size
|
|
BufferSize int
|
|
// Mirror mode is used when the memory backend is used for caching. In mirror
|
|
// mode, record IDs for Put and PutRange requests are re-used (instead of
|
|
// generating fresh ones) and expiration is turned off.
|
|
Mirror bool
|
|
}
|
|
|
|
// CheckAndSetDefaults checks and sets default values
|
|
func (cfg *Config) CheckAndSetDefaults() error {
|
|
if cfg.Context == nil {
|
|
cfg.Context = context.Background()
|
|
}
|
|
if cfg.BufferSize == 0 {
|
|
cfg.BufferSize = backend.DefaultBufferCapacity
|
|
}
|
|
if cfg.BTreeDegree <= 0 {
|
|
cfg.BTreeDegree = defaultBTreeDegree
|
|
}
|
|
if cfg.Clock == nil {
|
|
cfg.Clock = clockwork.NewRealClock()
|
|
}
|
|
if cfg.Component == "" {
|
|
cfg.Component = teleport.ComponentMemory
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// New creates a new memory backend
|
|
func New(cfg Config) (*Memory, error) {
|
|
if err := cfg.CheckAndSetDefaults(); err != nil {
|
|
return nil, trace.Wrap(err)
|
|
}
|
|
ctx, cancel := context.WithCancel(cfg.Context)
|
|
buf := backend.NewCircularBuffer(
|
|
backend.BufferCapacity(cfg.BufferSize),
|
|
)
|
|
buf.SetInit()
|
|
m := &Memory{
|
|
Mutex: &sync.Mutex{},
|
|
Entry: log.WithFields(log.Fields{
|
|
trace.Component: teleport.ComponentMemory,
|
|
}),
|
|
Config: cfg,
|
|
tree: btree.NewG(cfg.BTreeDegree, func(a, b *btreeItem) bool {
|
|
return a.Less(b)
|
|
}),
|
|
heap: newMinHeap(),
|
|
cancel: cancel,
|
|
ctx: ctx,
|
|
buf: buf,
|
|
}
|
|
return m, nil
|
|
}
|
|
|
|
// Memory is a memory B-Tree based backend
|
|
type Memory struct {
|
|
// nextID is a next record ID
|
|
// intentionally placed first to ensure 64-bit alignment
|
|
nextID int64
|
|
|
|
*sync.Mutex
|
|
*log.Entry
|
|
Config
|
|
// tree is a BTree with items
|
|
tree *btree.BTreeG[*btreeItem]
|
|
// heap is a min heap with expiry records
|
|
heap *minHeap
|
|
// cancel is a function that cancels
|
|
// all operations
|
|
cancel context.CancelFunc
|
|
// ctx is a context signaling close
|
|
ctx context.Context
|
|
buf *backend.CircularBuffer
|
|
}
|
|
|
|
// Close closes memory backend
|
|
func (m *Memory) Close() error {
|
|
m.cancel()
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.buf.Close()
|
|
return nil
|
|
}
|
|
|
|
// CloseWatchers closes all the watchers
|
|
// without closing the backend
|
|
func (m *Memory) CloseWatchers() {
|
|
m.buf.Clear()
|
|
}
|
|
|
|
// Clock returns clock used by this backend
|
|
func (m *Memory) Clock() clockwork.Clock {
|
|
return m.Config.Clock
|
|
}
|
|
|
|
// Create creates item if it does not exist
|
|
func (m *Memory) Create(ctx context.Context, i backend.Item) (*backend.Lease, error) {
|
|
if len(i.Key) == 0 {
|
|
return nil, trace.BadParameter("missing parameter key")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
if m.tree.Has(&btreeItem{Item: i}) {
|
|
return nil, trace.AlreadyExists("key %q already exists", string(i.Key))
|
|
}
|
|
event := backend.Event{
|
|
Type: types.OpPut,
|
|
Item: i,
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
return m.newLease(i), nil
|
|
}
|
|
|
|
// Get returns a single item or not found error
|
|
func (m *Memory) Get(ctx context.Context, key []byte) (*backend.Item, error) {
|
|
if len(key) == 0 {
|
|
return nil, trace.BadParameter("missing parameter key")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
i, found := m.tree.Get(&btreeItem{Item: backend.Item{Key: key}})
|
|
if !found {
|
|
return nil, trace.NotFound("key %q is not found", string(key))
|
|
}
|
|
return &i.Item, nil
|
|
}
|
|
|
|
// Update updates item if it exists, or returns NotFound error
|
|
func (m *Memory) Update(ctx context.Context, i backend.Item) (*backend.Lease, error) {
|
|
if len(i.Key) == 0 {
|
|
return nil, trace.BadParameter("missing parameter key")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
if !m.tree.Has(&btreeItem{Item: i}) {
|
|
return nil, trace.NotFound("key %q is not found", string(i.Key))
|
|
}
|
|
if !m.Mirror {
|
|
i.ID = m.generateID()
|
|
}
|
|
event := backend.Event{
|
|
Type: types.OpPut,
|
|
Item: i,
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
return m.newLease(i), nil
|
|
}
|
|
|
|
// Put puts value into backend (creates if it does not
|
|
// exist, updates it otherwise)
|
|
func (m *Memory) Put(ctx context.Context, i backend.Item) (*backend.Lease, error) {
|
|
if len(i.Key) == 0 {
|
|
return nil, trace.BadParameter("missing parameter key")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
if !m.Mirror {
|
|
i.ID = m.generateID()
|
|
}
|
|
event := backend.Event{
|
|
Type: types.OpPut,
|
|
Item: i,
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
return m.newLease(i), nil
|
|
}
|
|
|
|
// PutRange puts range of items into backend (creates if items do not
|
|
// exist, updates it otherwise)
|
|
func (m *Memory) PutRange(ctx context.Context, items []backend.Item) error {
|
|
for i := range items {
|
|
if items[i].Key == nil {
|
|
return trace.BadParameter("missing parameter key in item %v", i)
|
|
}
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
for _, item := range items {
|
|
event := backend.Event{
|
|
Type: types.OpPut,
|
|
Item: item,
|
|
}
|
|
if !m.Mirror {
|
|
event.Item.ID = m.generateID()
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Delete deletes item by key, returns NotFound error
|
|
// if item does not exist
|
|
func (m *Memory) Delete(ctx context.Context, key []byte) error {
|
|
if len(key) == 0 {
|
|
return trace.BadParameter("missing parameter key")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
if !m.tree.Has(&btreeItem{Item: backend.Item{Key: key}}) {
|
|
return trace.NotFound("key %q is not found", string(key))
|
|
}
|
|
event := backend.Event{
|
|
Type: types.OpDelete,
|
|
Item: backend.Item{
|
|
Key: key,
|
|
},
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// DeleteRange deletes range of items with keys between startKey and endKey
|
|
// Note that elements deleted by range do not produce any events
|
|
func (m *Memory) DeleteRange(ctx context.Context, startKey, endKey []byte) error {
|
|
if len(startKey) == 0 {
|
|
return trace.BadParameter("missing parameter startKey")
|
|
}
|
|
if len(endKey) == 0 {
|
|
return trace.BadParameter("missing parameter endKey")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
re := m.getRange(ctx, startKey, endKey, backend.NoLimit)
|
|
for _, item := range re.Items {
|
|
event := backend.Event{
|
|
Type: types.OpDelete,
|
|
Item: item,
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetRange returns query range
|
|
func (m *Memory) GetRange(ctx context.Context, startKey []byte, endKey []byte, limit int) (*backend.GetResult, error) {
|
|
if len(startKey) == 0 {
|
|
return nil, trace.BadParameter("missing parameter startKey")
|
|
}
|
|
if len(endKey) == 0 {
|
|
return nil, trace.BadParameter("missing parameter endKey")
|
|
}
|
|
if limit <= 0 {
|
|
limit = backend.DefaultRangeLimit
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
re := m.getRange(ctx, startKey, endKey, limit)
|
|
if len(re.Items) == backend.DefaultRangeLimit {
|
|
m.Warnf("Range query hit backend limit. (this is a bug!) startKey=%q,limit=%d", startKey, backend.DefaultRangeLimit)
|
|
}
|
|
return &re, nil
|
|
}
|
|
|
|
// KeepAlive updates TTL on the lease
|
|
func (m *Memory) KeepAlive(ctx context.Context, lease backend.Lease, expires time.Time) error {
|
|
if lease.IsEmpty() {
|
|
return trace.BadParameter("lease is empty")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
i, found := m.tree.Get(&btreeItem{Item: backend.Item{Key: lease.Key}})
|
|
if !found {
|
|
return trace.NotFound("key %q is not found", string(lease.Key))
|
|
}
|
|
item := i.Item
|
|
item.Expires = expires
|
|
if !m.Mirror {
|
|
// ID is updated on keep alive for consistency with other backends
|
|
item.ID = m.generateID()
|
|
}
|
|
event := backend.Event{
|
|
Type: types.OpPut,
|
|
Item: item,
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CompareAndSwap compares item with existing item and replaces it with replaceWith item
|
|
func (m *Memory) CompareAndSwap(ctx context.Context, expected backend.Item, replaceWith backend.Item) (*backend.Lease, error) {
|
|
if len(expected.Key) == 0 {
|
|
return nil, trace.BadParameter("missing parameter Key")
|
|
}
|
|
if len(replaceWith.Key) == 0 {
|
|
return nil, trace.BadParameter("missing parameter Key")
|
|
}
|
|
if !bytes.Equal(expected.Key, replaceWith.Key) {
|
|
return nil, trace.BadParameter("expected and replaceWith keys should match")
|
|
}
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
m.removeExpired()
|
|
i, found := m.tree.Get(&btreeItem{Item: expected})
|
|
if !found {
|
|
return nil, trace.CompareFailed("key %q is not found", string(expected.Key))
|
|
}
|
|
existingItem := i.Item
|
|
if !bytes.Equal(existingItem.Value, expected.Value) {
|
|
return nil, trace.CompareFailed("current value does not match expected for %v", string(expected.Key))
|
|
}
|
|
event := backend.Event{
|
|
Type: types.OpPut,
|
|
Item: replaceWith,
|
|
}
|
|
m.processEvent(event)
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
return m.newLease(replaceWith), nil
|
|
}
|
|
|
|
// NewWatcher returns a new event watcher
|
|
func (m *Memory) NewWatcher(ctx context.Context, watch backend.Watch) (backend.Watcher, error) {
|
|
if m.EventsOff {
|
|
return nil, trace.BadParameter("events are turned off for this backend")
|
|
}
|
|
return m.buf.NewWatcher(ctx, watch)
|
|
}
|
|
|
|
func (m *Memory) generateID() int64 {
|
|
return atomic.AddInt64(&m.nextID, 1)
|
|
}
|
|
|
|
func (m *Memory) getRange(ctx context.Context, startKey, endKey []byte, limit int) backend.GetResult {
|
|
var res backend.GetResult
|
|
m.tree.AscendRange(&btreeItem{Item: backend.Item{Key: startKey}}, &btreeItem{Item: backend.Item{Key: endKey}}, func(item *btreeItem) bool {
|
|
res.Items = append(res.Items, item.Item)
|
|
if limit > 0 && len(res.Items) >= limit {
|
|
return false
|
|
}
|
|
return true
|
|
})
|
|
return res
|
|
}
|
|
|
|
func (m *Memory) newLease(item backend.Item) *backend.Lease {
|
|
var lease backend.Lease
|
|
if item.Expires.IsZero() {
|
|
return &lease
|
|
}
|
|
lease.Key = item.Key
|
|
return &lease
|
|
}
|
|
|
|
// removeExpired makes a pass through map and removes expired elements
|
|
// returns the number of expired elements removed
|
|
func (m *Memory) removeExpired() int {
|
|
// In mirror mode, don't expire any elements. This allows the cache to setup
|
|
// a watch and expire elements as the events roll in.
|
|
if m.Mirror {
|
|
return 0
|
|
}
|
|
|
|
removed := 0
|
|
now := m.Clock().Now().UTC()
|
|
for {
|
|
if len(*m.heap) == 0 {
|
|
break
|
|
}
|
|
item := m.heap.PeekEl()
|
|
if now.Before(item.Expires) {
|
|
break
|
|
}
|
|
m.heap.PopEl()
|
|
m.tree.Delete(item)
|
|
m.Debugf("Removed expired %v %v item.", string(item.Key), item.Expires)
|
|
removed++
|
|
|
|
event := backend.Event{
|
|
Type: types.OpDelete,
|
|
Item: backend.Item{
|
|
Key: item.Key,
|
|
},
|
|
}
|
|
if !m.EventsOff {
|
|
m.buf.Emit(event)
|
|
}
|
|
}
|
|
if removed > 0 {
|
|
m.Debugf("Removed %v expired items.", removed)
|
|
}
|
|
return removed
|
|
}
|
|
|
|
func (m *Memory) processEvent(event backend.Event) {
|
|
switch event.Type {
|
|
case types.OpPut:
|
|
item := &btreeItem{Item: event.Item, index: -1}
|
|
treeItem, found := m.tree.Get(item)
|
|
var existingItem *btreeItem
|
|
if found {
|
|
existingItem = treeItem
|
|
}
|
|
|
|
// new item is added, but it has not expired yet
|
|
if existingItem != nil && existingItem.index >= 0 {
|
|
m.heap.RemoveEl(existingItem)
|
|
}
|
|
if !item.Expires.IsZero() {
|
|
m.heap.PushEl(item)
|
|
}
|
|
m.tree.ReplaceOrInsert(item)
|
|
case types.OpDelete:
|
|
item, found := m.tree.Get(&btreeItem{Item: event.Item})
|
|
if !found {
|
|
return
|
|
}
|
|
|
|
m.tree.Delete(item)
|
|
if item.index >= 0 {
|
|
m.heap.RemoveEl(item)
|
|
}
|
|
default:
|
|
// skip unsupported record
|
|
}
|
|
}
|