seaweedfs/weed/filer/ydb/ydb_store.go
2022-05-12 12:23:02 +05:00

414 lines
12 KiB
Go

//go:build ydb
// +build ydb
package ydb
import (
"context"
"fmt"
"github.com/chrislusf/seaweedfs/weed/filer"
"github.com/chrislusf/seaweedfs/weed/filer/abstract_sql"
"github.com/chrislusf/seaweedfs/weed/glog"
"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
"github.com/chrislusf/seaweedfs/weed/util"
environ "github.com/ydb-platform/ydb-go-sdk-auth-environ"
"github.com/ydb-platform/ydb-go-sdk/v3"
"github.com/ydb-platform/ydb-go-sdk/v3/sugar"
"github.com/ydb-platform/ydb-go-sdk/v3/table"
"github.com/ydb-platform/ydb-go-sdk/v3/table/result"
"github.com/ydb-platform/ydb-go-sdk/v3/table/result/named"
"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
"os"
"path"
"strings"
"sync"
"time"
)
const (
defaultDialTimeOut = 10
)
var (
roTX = table.TxControl(
table.BeginTx(table.WithOnlineReadOnly()),
table.CommitTx(),
)
rwTX = table.DefaultTxControl()
)
type YdbStore struct {
DB ydb.Connection
dirBuckets string
tablePathPrefix string
SupportBucketTable bool
dbs map[string]bool
dbsLock sync.Mutex
}
func init() {
filer.Stores = append(filer.Stores, &YdbStore{})
}
func (store *YdbStore) GetName() string {
return "ydb"
}
func (store *YdbStore) Initialize(configuration util.Configuration, prefix string) (err error) {
return store.initialize(
configuration.GetString("filer.options.buckets_folder"),
configuration.GetString(prefix+"dsn"),
configuration.GetString(prefix+"prefix"),
configuration.GetBool(prefix+"useBucketPrefix"),
configuration.GetInt(prefix+"dialTimeOut"),
configuration.GetInt(prefix+"poolSizeLimit"),
)
}
func (store *YdbStore) initialize(dirBuckets string, dsn string, tablePathPrefix string, useBucketPrefix bool, dialTimeOut int, poolSizeLimit int) (err error) {
store.dirBuckets = dirBuckets
store.SupportBucketTable = useBucketPrefix
if store.SupportBucketTable {
glog.V(0).Infof("enabled BucketPrefix")
}
store.dbs = make(map[string]bool)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
if dialTimeOut == 0 {
dialTimeOut = defaultDialTimeOut
}
opts := []ydb.Option{
ydb.WithDialTimeout(time.Duration(dialTimeOut) * time.Second),
environ.WithEnvironCredentials(ctx),
}
if poolSizeLimit > 0 {
opts = append(opts, ydb.WithSessionPoolSizeLimit(poolSizeLimit))
}
if dsn == "" {
dsn = os.Getenv("YDB_CONNECTION_STRING")
}
store.DB, err = ydb.Open(ctx, dsn, opts...)
if err != nil {
if store.DB != nil {
_ = store.DB.Close(ctx)
store.DB = nil
}
return fmt.Errorf("can not connect to %s error: %v", dsn, err)
}
store.tablePathPrefix = path.Join(store.DB.Name(), tablePathPrefix)
if err = sugar.MakeRecursive(ctx, store.DB, store.tablePathPrefix); err != nil {
return fmt.Errorf("MakeRecursive %s : %v", store.tablePathPrefix, err)
}
if err = store.createTable(ctx, store.tablePathPrefix); err != nil {
glog.Errorf("createTable %s: %v", store.tablePathPrefix, err)
}
return err
}
func (store *YdbStore) doTxOrDB(ctx context.Context, query *string, params *table.QueryParameters, tc *table.TransactionControl, processResultFunc func(res result.Result) error) (err error) {
var res result.Result
if tx, ok := ctx.Value("tx").(table.Transaction); ok {
res, err = tx.Execute(ctx, *query, params)
if err != nil {
return fmt.Errorf("execute transaction: %v", err)
}
} else {
err = store.DB.Table().Do(ctx, func(ctx context.Context, s table.Session) (err error) {
_, res, err = s.Execute(ctx, tc, *query, params)
if err != nil {
return fmt.Errorf("execute statement: %v", err)
}
return nil
})
}
if err != nil {
return err
}
if res != nil {
defer func() { _ = res.Close() }()
if processResultFunc != nil {
if err = processResultFunc(res); err != nil {
return fmt.Errorf("process result: %v", err)
}
}
}
return err
}
func (store *YdbStore) insertOrUpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
dir, name := entry.FullPath.DirAndName()
meta, err := entry.EncodeAttributesAndChunks()
if err != nil {
return fmt.Errorf("encode %s: %s", entry.FullPath, err)
}
if len(entry.Chunks) > filer.CountEntryChunksForGzip {
meta = util.MaybeGzipData(meta)
}
tablePathPrefix, shortDir := store.getPrefix(ctx, &dir)
fileMeta := FileMeta{util.HashStringToLong(dir), name, *shortDir, meta}
return store.doTxOrDB(ctx, withPragma(tablePathPrefix, upsertQuery), fileMeta.queryParameters(entry.TtlSec), rwTX, nil)
}
func (store *YdbStore) InsertEntry(ctx context.Context, entry *filer.Entry) (err error) {
return store.insertOrUpdateEntry(ctx, entry)
}
func (store *YdbStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
return store.insertOrUpdateEntry(ctx, entry)
}
func (store *YdbStore) FindEntry(ctx context.Context, fullpath util.FullPath) (entry *filer.Entry, err error) {
dir, name := fullpath.DirAndName()
var data []byte
entryFound := false
tablePathPrefix, shortDir := store.getPrefix(ctx, &dir)
query := withPragma(tablePathPrefix, findQuery)
queryParams := table.NewQueryParameters(
table.ValueParam("$dir_hash", types.Int64Value(util.HashStringToLong(*shortDir))),
table.ValueParam("$name", types.UTF8Value(name)))
err = store.doTxOrDB(ctx, query, queryParams, roTX, func(res result.Result) error {
if !res.NextResultSet(ctx) || !res.HasNextRow() {
return nil
}
for res.NextRow() {
if err = res.ScanNamed(named.OptionalWithDefault("meta", &data)); err != nil {
return fmt.Errorf("scanNamed %s : %v", fullpath, err)
}
entryFound = true
return nil
}
return res.Err()
})
if err != nil {
return nil, err
}
if !entryFound {
return nil, filer_pb.ErrNotFound
}
entry = &filer.Entry{
FullPath: fullpath,
}
if err := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(data)); err != nil {
return nil, fmt.Errorf("decode %s : %v", fullpath, err)
}
return entry, nil
}
func (store *YdbStore) DeleteEntry(ctx context.Context, fullpath util.FullPath) (err error) {
dir, name := fullpath.DirAndName()
tablePathPrefix, shortDir := store.getPrefix(ctx, &dir)
query := withPragma(tablePathPrefix, deleteQuery)
queryParams := table.NewQueryParameters(
table.ValueParam("$dir_hash", types.Int64Value(util.HashStringToLong(*shortDir))),
table.ValueParam("$name", types.UTF8Value(name)))
return store.doTxOrDB(ctx, query, queryParams, rwTX, nil)
}
func (store *YdbStore) DeleteFolderChildren(ctx context.Context, fullpath util.FullPath) (err error) {
dir, _ := fullpath.DirAndName()
tablePathPrefix, shortDir := store.getPrefix(ctx, &dir)
query := withPragma(tablePathPrefix, deleteFolderChildrenQuery)
queryParams := table.NewQueryParameters(
table.ValueParam("$dir_hash", types.Int64Value(util.HashStringToLong(*shortDir))),
table.ValueParam("$directory", types.UTF8Value(*shortDir)))
return store.doTxOrDB(ctx, query, queryParams, rwTX, nil)
}
func (store *YdbStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "", eachEntryFunc)
}
func (store *YdbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
dir := string(dirPath)
tablePathPrefix, shortDir := store.getPrefix(ctx, &dir)
var query *string
if includeStartFile {
query = withPragma(tablePathPrefix, listInclusiveDirectoryQuery)
} else {
query = withPragma(tablePathPrefix, listDirectoryQuery)
}
truncated := true
eachEntryFuncIsNotBreake := true
entryCount := int64(0)
for truncated && eachEntryFuncIsNotBreake {
if lastFileName != "" {
startFileName = lastFileName
if includeStartFile {
query = withPragma(tablePathPrefix, listDirectoryQuery)
}
}
restLimit := limit - entryCount
queryParams := table.NewQueryParameters(
table.ValueParam("$dir_hash", types.Int64Value(util.HashStringToLong(*shortDir))),
table.ValueParam("$directory", types.UTF8Value(*shortDir)),
table.ValueParam("$start_name", types.UTF8Value(startFileName)),
table.ValueParam("$prefix", types.UTF8Value(prefix+"%")),
table.ValueParam("$limit", types.Uint64Value(uint64(restLimit))),
)
err = store.doTxOrDB(ctx, query, queryParams, roTX, func(res result.Result) error {
var name string
var data []byte
if !res.NextResultSet(ctx) || !res.HasNextRow() {
truncated = false
return nil
}
truncated = res.CurrentResultSet().Truncated()
for res.NextRow() {
if err := res.ScanNamed(
named.OptionalWithDefault("name", &name),
named.OptionalWithDefault("meta", &data)); err != nil {
return fmt.Errorf("list scanNamed %s : %v", dir, err)
}
lastFileName = name
entry := &filer.Entry{
FullPath: util.NewFullPath(dir, name),
}
if err = entry.DecodeAttributesAndChunks(util.MaybeDecompressData(data)); err != nil {
return fmt.Errorf("scan decode %s : %v", entry.FullPath, err)
}
if !eachEntryFunc(entry) {
eachEntryFuncIsNotBreake = false
break
}
entryCount += 1
}
return res.Err()
})
}
if err != nil {
return lastFileName, err
}
return lastFileName, nil
}
func (store *YdbStore) BeginTransaction(ctx context.Context) (context.Context, error) {
session, err := store.DB.Table().CreateSession(ctx)
if err != nil {
return ctx, err
}
tx, err := session.BeginTransaction(ctx, table.TxSettings(table.WithSerializableReadWrite()))
if err != nil {
return ctx, err
}
return context.WithValue(ctx, "tx", tx), nil
}
func (store *YdbStore) CommitTransaction(ctx context.Context) error {
if tx, ok := ctx.Value("tx").(table.Transaction); ok {
_, err := tx.CommitTx(ctx)
return err
}
return nil
}
func (store *YdbStore) RollbackTransaction(ctx context.Context) error {
if tx, ok := ctx.Value("tx").(table.Transaction); ok {
return tx.Rollback(ctx)
}
return nil
}
func (store *YdbStore) Shutdown() {
_ = store.DB.Close(context.Background())
}
func (store *YdbStore) CanDropWholeBucket() bool {
return store.SupportBucketTable
}
func (store *YdbStore) OnBucketCreation(bucket string) {
store.dbsLock.Lock()
defer store.dbsLock.Unlock()
if err := store.createTable(context.Background(),
path.Join(store.tablePathPrefix, bucket)); err != nil {
glog.Errorf("createTable %s: %v", bucket, err)
}
if store.dbs == nil {
return
}
store.dbs[bucket] = true
}
func (store *YdbStore) OnBucketDeletion(bucket string) {
store.dbsLock.Lock()
defer store.dbsLock.Unlock()
if err := store.deleteTable(context.Background(),
path.Join(store.tablePathPrefix, bucket)); err != nil {
glog.Errorf("deleteTable %s: %v", bucket, err)
}
if store.dbs == nil {
return
}
delete(store.dbs, bucket)
}
func (store *YdbStore) createTable(ctx context.Context, prefix string) error {
return store.DB.Table().Do(ctx, func(ctx context.Context, s table.Session) error {
return s.CreateTable(ctx, path.Join(prefix, abstract_sql.DEFAULT_TABLE), createTableOptions()...)
})
}
func (store *YdbStore) deleteTable(ctx context.Context, prefix string) error {
if !store.SupportBucketTable {
return nil
}
if err := store.DB.Table().Do(ctx, func(ctx context.Context, s table.Session) error {
return s.DropTable(ctx, path.Join(prefix, abstract_sql.DEFAULT_TABLE))
}); err != nil {
return err
}
glog.V(4).Infof("deleted table %s", prefix)
return nil
}
func (store *YdbStore) getPrefix(ctx context.Context, dir *string) (tablePathPrefix *string, shortDir *string) {
tablePathPrefix = &store.tablePathPrefix
shortDir = dir
if !store.SupportBucketTable {
return
}
prefixBuckets := store.dirBuckets + "/"
if strings.HasPrefix(*dir, prefixBuckets) {
// detect bucket
bucketAndDir := (*dir)[len(prefixBuckets):]
var bucket string
if t := strings.Index(bucketAndDir, "/"); t > 0 {
bucket = bucketAndDir[:t]
} else if t < 0 {
bucket = bucketAndDir
}
if bucket == "" {
return
}
store.dbsLock.Lock()
defer store.dbsLock.Unlock()
tablePathPrefixWithBucket := path.Join(store.tablePathPrefix, bucket)
if _, found := store.dbs[bucket]; !found {
if err := store.createTable(ctx, tablePathPrefixWithBucket); err == nil {
store.dbs[bucket] = true
glog.V(4).Infof("created table %s", tablePathPrefixWithBucket)
} else {
glog.Errorf("createTable %s: %v", tablePathPrefixWithBucket, err)
}
}
tablePathPrefix = &tablePathPrefixWithBucket
}
return
}