cubefs/blobstore/common/kvstore/table.go

// Copyright 2022 The CubeFS Authors.
//
// 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 kvstore

import (
	"bytes"
	"errors"
	"fmt"
	"sync"

	rdb "github.com/tecbot/gorocksdb"
)

type KVTable interface {
	KVStorage
	GetDB() *rdb.DB
	GetCf() *rdb.ColumnFamilyHandle
}

type table struct {
	name string
	ins  *instance
	db   *rdb.DB
	ro   *rdb.ReadOptions
	wo   *rdb.WriteOptions
	fo   *rdb.FlushOptions
	cf   *rdb.ColumnFamilyHandle
}

/* table function list */
func (t *table) Get(key []byte) (data []byte, err error) {
	done := make(chan struct{})
	t.ins.rpool.Run(func() {
		defer close(done)
		cValue, err1 := t.db.GetCF(t.ro, t.cf, key)
		if err1 != nil {
			err = err1
			return
		}
		defer cValue.Free()

		if !cValue.Exists() {
			err = ErrNotFound
			return
		}
		data = make([]byte, cValue.Size())
		copy(data, cValue.Data())
	})
	<-done
	return
}

func (t *table) Put(kv KV) (err error) {
	task := &writeTask{
		typ:  cfPutEvent,
		data: cfPut{t.cf, kv.Key, kv.Value},
		err:  make(chan error, 1),
	}
	t.ins.wchan <- task
	return <-task.err
}

func (t *table) Delete(key []byte) (err error) {
	task := &writeTask{
		typ:  cfDeleteEvent,
		data: cfDel{t.cf, key},
		err:  make(chan error, 1),
	}
	t.ins.wchan <- task
	return <-task.err
}

func (t *table) DeleteRange(start, end []byte) (err error) {
	task := &writeTask{
		typ:  cfRangeDeleteEvent,
		data: cfRangeDelete{t.cf, start, end},
		err:  make(chan error, 1),
	}
	t.ins.wchan <- task
	return <-task.err
}

func (t *table) Name() string {
	return t.name
}

func (t *table) Flush() error {
	return t.ins.Flush()
}

func (t *table) DeleteBatch(keys [][]byte, safe bool) (err error) {
	b := []batch{}
	for _, key := range keys {
		if safe {
			_, err := t.Get(key)
			if err != nil {
				return err
			}
		}
		b = append(b, batch{cfDeleteEvent, cfDel{t.cf, key}})
	}
	if len(b) == 0 {
		return
	}
	task := &writeTask{
		typ:  batchEvent,
		data: b,
		err:  make(chan error, 1),
	}
	t.ins.wchan <- task
	return <-task.err
}

func (t *table) WriteBatch(kvs []KV, safe bool) (err error) {
	b := []batch{}
	for _, kv := range kvs {
		key := kv.Key
		if safe {
			v, err := t.Get(key)
			if err != nil {
				return err
			}
			if v != nil &&
				!bytes.Equal(kv.Value, v) {
				msg := fmt.Sprintf("table(%v): conflict data at row %v\n old buf: %v, new buf: %v",
					t.db.Name(), kv.Key, v, kv.Value)
				return errors.New(msg)
			}
		}
		b = append(b, batch{cfPutEvent, cfPut{t.cf, kv.Key, kv.Value}})
	}
	if len(b) == 0 {
		return
	}
	task := &writeTask{
		typ:  batchEvent,
		data: b,
		err:  make(chan error, 1),
	}
	t.ins.wchan <- task
	return <-task.err
}

func (t *table) NewSnapshot() *Snapshot {
	return (*Snapshot)(t.db.NewSnapshot())
}

func (s *table) ReleaseSnapshot(snapshot *Snapshot) {
	s.db.ReleaseSnapshot((*rdb.Snapshot)(snapshot))
}

func (t *table) NewIterator(snapshot *Snapshot, opts ...OpOption) Iterator {
	var op Op

	op.applyOpts(opts)
	ro := op.Ro

	if ro == nil {
		ro = NewReadOptions()
	}

	if snapshot != nil {
		ro.SetSnapshot((*rdb.Snapshot)(snapshot))
	}

	return &iterator{rpool: &t.ins.rpool, ro: ro.ReadOptions, iter: t.db.NewIteratorCF(ro.ReadOptions, t.cf), once: sync.Once{}}
}

func (t *table) GetDB() *rdb.DB {
	return t.db
}

func (t *table) GetCf() *rdb.ColumnFamilyHandle {
	return t.cf
}

func (t *table) NewWriteBatch() *WriteBatch {
	return &WriteBatch{}
}

func (t *table) DoBatch(batch *WriteBatch) error {
	return t.ins.DoBatch(batch)
}