scriptoriumadmin/featurebase
1
0
Fork 0
mirror of https://github.com/featurebasedb/featurebase.git synced 2026-08-28 02:44:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #2207 from FeatureBaseDB/commit-sync-1-19-2022

Browse source 
Commit sync 1-19-2023
This commit is contained in:
Joseph Friedrich 2023-01-19 20:19:26 -06:00 • committed by GitHub
commit ceb91cff51
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
151 changed files with 10508 additions and 6457 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
Dockerfile-clustertests
View file

@ -20,12 +20,12 @@ RUN apt install -y docker.io
ADD https://github.com/docker/compose/releases/latest/download/docker-compose-Linux-x86_64 /usr/local/bin/docker-compose
RUN chmod +x /usr/local/bin/docker-compose
WORKDIR /go/src/github.com/molecula/featurebase/cmd/featurebase
WORKDIR /go/src/github.com/featurebasedb/featurebase/cmd/featurebase
# generate an instrumented binary to allow for calculating code coverage for clustertests
# the entrypoint for the binary is TestRunMain, which is wrapper for main
RUN go test -covermode=atomic -coverpkg=../../... -c -tags testrunmain -o featurebase
RUN cp /go/src/github.com/molecula/featurebase/cmd/featurebase/featurebase /featurebase
RUN cp /go/src/github.com/featurebasedb/featurebase/cmd/featurebase/featurebase /featurebase
COPY NOTICE /NOTICE

6
Dockerfile-clustertests-client
View file

@ -19,10 +19,10 @@ RUN apt install -y docker.io
ADD https://github.com/docker/compose/releases/latest/download/docker-compose-Linux-x86_64 /usr/local/bin/docker-compose
RUN chmod +x /usr/local/bin/docker-compose
WORKDIR /go/src/github.com/molecula/featurebase/cmd/featurebase
WORKDIR /go/src/github.com/featurebasedb/featurebase/cmd/featurebase
RUN go test -covermode=atomic -coverpkg=../../... -c -tags testrunmain -o featurebase
RUN cp /go/src/github.com/molecula/featurebase/cmd/featurebase/featurebase /featurebase
RUN cp /go/src/github.com/featurebasedb/featurebase/cmd/featurebase/featurebase /featurebase
COPY NOTICE /NOTICE
@ -32,6 +32,6 @@ COPY ./internal/clustertests /go/src/github.com/featurebasedb/featurebase/intern
EXPOSE 10101
VOLUME /data
WORKDIR /go/src/github.com/molecula/featurebase
WORKDIR /go/src/github.com/featurebasedb/featurebase
CMD ["/featurebase", "-test.run=TestRunMain", "-test.coverprofile=/results/coverage.out", "server", "--data-dir", "/data", "--bind", "http://0.0.0.0:10101"]

102
api.go
View file

@ -29,12 +29,12 @@ import (
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/featurebasedb/featurebase/v3/disco"
"github.com/featurebasedb/featurebase/v3/rbf"
"github.com/prometheus/client_golang/prometheus"
//"github.com/featurebasedb/featurebase/v3/pg"
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/roaring"
planner_types "github.com/featurebasedb/featurebase/v3/sql3/planner/types"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/tracing"
"github.com/pkg/errors"
"golang.org/x/sync/errgroup"
@ -278,7 +278,7 @@ func (api *API) CreateIndex(ctx context.Context, indexName string, options Index
return nil, errors.Wrap(err, "creating index")
}
api.holder.Stats.Count(MetricCreateIndex, 1, 1.0)
CounterCreateIndex.Inc()
return index, nil
}
@ -319,7 +319,7 @@ func (api *API) DeleteDataframe(ctx context.Context, indexName string) error {
api.server.logger.Errorf("problem sending DeleteIndex message: %s", err)
return errors.Wrap(err, "sending DeleteIndex message")
}
api.holder.Stats.Count(MetricDeleteDataframe, 1, 1.0)
CounterDeleteDataframe.Inc()
return nil
}
@ -354,7 +354,7 @@ func (api *API) DeleteIndex(ctx context.Context, indexName string) error {
return errors.Wrap(err, "deleting id allocation for index")
}
}
api.holder.Stats.Count(MetricDeleteIndex, 1, 1.0)
CounterDeleteIndex.Inc()
return nil
}
@ -406,7 +406,7 @@ func (api *API) CreateField(ctx context.Context, indexName string, fieldName str
return nil, errors.Wrap(err, "sending CreateField message")
}
api.holder.Stats.CountWithCustomTags(MetricCreateField, 1, 1.0, []string{fmt.Sprintf("index:%s", indexName)})
CounterCreateField.With(prometheus.Labels{"index": indexName})
return field, nil
}
@ -750,7 +750,7 @@ func (api *API) DeleteField(ctx context.Context, indexName string, fieldName str
api.server.logger.Errorf("problem sending DeleteField message: %s", err)
return errors.Wrap(err, "sending DeleteField message")
}
api.holder.Stats.CountWithCustomTags(MetricDeleteField, 1, 1.0, []string{fmt.Sprintf("index:%s", indexName)})
CounterDeleteField.With(prometheus.Labels{"index": indexName})
return nil
}
@ -782,7 +782,7 @@ func (api *API) DeleteAvailableShard(_ context.Context, indexName, fieldName str
api.server.logger.Errorf("problem sending DeleteAvailableShard message: %s", err)
return errors.Wrap(err, "sending DeleteAvailableShard message")
}
api.holder.Stats.CountWithCustomTags(MetricDeleteAvailableShard, 1, 1.0, []string{fmt.Sprintf("index:%s", indexName)})
CounterDeleteAvailableShard.With(prometheus.Labels{"index": indexName}).Inc()
return nil
}
@ -2096,15 +2096,6 @@ func (api *API) AvailableShards(ctx context.Context, indexName string) (*roaring
return index.AvailableShards(false), nil
}
// StatsWithTags returns an instance of whatever implementation of StatsClient
// pilosa is using with the given tags.
func (api *API) StatsWithTags(tags []string) stats.StatsClient {
if api.holder == nil || api.cluster == nil {
return nil
}
return api.holder.Stats.WithTags(tags...)
}
// LongQueryTime returns the configured threshold for logging/statting
// long running queries.
func (api *API) LongQueryTime() time.Duration {
@ -2391,19 +2382,19 @@ func (api *API) StartTransaction(ctx context.Context, id string, timeout time.Du
switch err {
case nil:
if exclusive {
api.holder.Stats.Count(MetricExclusiveTransactionRequest, 1, 1.0)
CounterExclusiveTransactionRequest.Inc()
} else {
api.holder.Stats.Count(MetricTransactionStart, 1, 1.0)
CounterTransactionStart.Inc()
}
case ErrTransactionExclusive:
if exclusive {
api.holder.Stats.Count(MetricExclusiveTransactionBlocked, 1, 1.0)
CounterExclusiveTransactionBlocked.Inc()
} else {
api.holder.Stats.Count(MetricTransactionBlocked, 1, 1.0)
CounterTransactionBlocked.Inc()
}
}
if exclusive && t != nil && t.Active {
api.holder.Stats.Count(MetricExclusiveTransactionActive, 1, 1.0)
CounterExclusiveTransactionActive.Inc()
}
return t, err
}
@ -2415,9 +2406,9 @@ func (api *API) FinishTransaction(ctx context.Context, id string, remote bool) (
t, err := api.server.FinishTransaction(ctx, id, remote)
if err == nil {
if t.Exclusive {
api.holder.Stats.Count(MetricExclusiveTransactionEnd, 1, 1.0)
CounterExclusiveTransactionEnd.Inc()
} else {
api.holder.Stats.Count(MetricTransactionEnd, 1, 1.0)
CounterTransactionEnd.Inc()
}
}
return t, err
@ -2437,7 +2428,7 @@ func (api *API) GetTransaction(ctx context.Context, id string, remote bool) (*Tr
t, err := api.server.GetTransaction(ctx, id, remote)
if err == nil {
if t.Exclusive && t.Active {
api.holder.Stats.Count(MetricExclusiveTransactionActive, 1, 1.0)
CounterExclusiveTransactionActive.Inc()
}
}
return t, err
@ -3046,6 +3037,10 @@ func (api *API) CompilePlan(ctx context.Context, q string) (planner_types.PlanOp
return api.server.CompileExecutionPlan(ctx, q)
}
func (api *API) RehydratePlanOperator(ctx context.Context, reader io.Reader) (planner_types.PlanOperator, error) {
return api.server.RehydratePlanOperator(ctx, reader)
}
func (api *API) RBFDebugInfo() map[string]*rbf.DebugInfo {
infos := make(map[string]*rbf.DebugInfo)
@ -3343,6 +3338,7 @@ type SystemAPI interface {
ClusterState() string
DataDir() string
NodeID() string
ClusterNodes() []ClusterNode
}
@ -3358,6 +3354,9 @@ type QueryAPI interface {
Query(ctx context.Context, req *QueryRequest) (QueryResponse, error)
}
// Ensure type implements interface.
var _ SystemAPI = (*FeatureBaseSystemAPI)(nil)
// FeatureBaseSystemAPI is a wrapper around pilosa.API. It implements the
// SystemAPI interface
type FeatureBaseSystemAPI struct {
@ -3414,6 +3413,10 @@ func (fsapi *FeatureBaseSystemAPI) DataDir() string {
return fsapi.server.dataDir
}
func (fsapi *FeatureBaseSystemAPI) NodeID() string {
return fsapi.cluster.Node.ID
}
func (fsapi *FeatureBaseSystemAPI) ClusterNodes() []ClusterNode {
result := make([]ClusterNode, 0)
@ -3432,3 +3435,54 @@ func (fsapi *FeatureBaseSystemAPI) ClusterNodes() []ClusterNode {
return result
}
// Ensure type implements interface.
var _ SystemAPI = (*NopSystemAPI)(nil)
// NopSystemAPI is a no-op implementation of the SystemAPI.
type NopSystemAPI struct{}
func (napi *NopSystemAPI) ClusterName() string {
return ""
}
func (napi *NopSystemAPI) Version() string {
return ""
}
func (napi *NopSystemAPI) PlatformDescription() string {
return ""
}
func (napi *NopSystemAPI) PlatformVersion() string {
return ""
}
func (napi *NopSystemAPI) ClusterNodeCount() int {
return 0
}
func (napi *NopSystemAPI) ClusterReplicaCount() int {
return 0
}
func (napi *NopSystemAPI) ShardWidth() int {
return 0
}
func (napi *NopSystemAPI) ClusterState() string {
return ""
}
func (napi *NopSystemAPI) DataDir() string {
return ""
}
func (napi *NopSystemAPI) NodeID() string {
return ""
}
func (napi *NopSystemAPI) ClusterNodes() []ClusterNode {
result := make([]ClusterNode, 0)
return result
}

8
api_directive_test.go
View file

@ -20,10 +20,10 @@ func TestAPI_Directive(t *testing.T) {
api := c.GetPrimary().API
ctx := context.Background()
qual := dax.NewTableQualifier("acme", "db1")
tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "tbl1", 12, false)
tbl2 := daxtest.TestQualifiedTableWithID(t, qual, "2", "tbl2", 12, false)
tbl3 := daxtest.TestQualifiedTableWithID(t, qual, "3", "tbl3", 12, false)
qdbid := dax.NewQualifiedDatabaseID("acme", "db1")
tbl1 := daxtest.TestQualifiedTableWithID(t, qdbid, "1", "tbl1", 12, false)
tbl2 := daxtest.TestQualifiedTableWithID(t, qdbid, "2", "tbl2", 12, false)
tbl3 := daxtest.TestQualifiedTableWithID(t, qdbid, "3", "tbl3", 12, false)
t.Run("Schema", func(t *testing.T) {

117
apply.go
View file

@ -7,6 +7,7 @@ import (
"io"
"os"
"path/filepath"
"sort"
"strings"
"sync"
@ -260,14 +261,14 @@ func (e *executor) executeApplyShard(ctx context.Context, qcx *Qcx, index string
func NewShardFile(ctx context.Context, name string, mem memory.Allocator, e *executor) (*ShardFile, error) {
if !e.dataFrameExists(name) {
return &ShardFile{dest: name, executor: e}, nil
return &ShardFile{dest: name, executor: e, strings: make(map[key][]string)}, nil
}
// else read in existing
table, err := e.getDataTable(ctx, name, mem)
if err != nil {
return nil, err
}
return &ShardFile{table: table, schema: table.Schema(), dest: name, executor: e}, nil
return &ShardFile{table: table, schema: table.Schema(), dest: name, executor: e, strings: make(map[key][]string)}, nil
}
type NameType struct {
@ -291,6 +292,8 @@ func cast(v interface{}) arrow.DataType {
return arrow.PrimitiveTypes.Float64
case float64:
return arrow.PrimitiveTypes.Float64
case *arrow.StringType:
return arrow.BinaryTypes.String
default:
vprint.VV("%T .... %v", v, v)
}
@ -305,6 +308,11 @@ func (cr *ChangesetRequest) ArrowSchema() *arrow.Schema {
return arrow.NewSchema(fields, nil)
}
type key struct {
col int
chunk int
}
type ShardFile struct {
table arrow.Table
schema *arrow.Schema
@ -313,6 +321,7 @@ type ShardFile struct {
columns []interface{}
dest string
executor *executor
strings map[key][]string
}
func compareSchema(s1, s2 *arrow.Schema) bool {
@ -365,6 +374,8 @@ func (sf *ShardFile) buildAppenders(maxid int64) {
sf.columns[i] = make([]int64, newSize)
case arrow.PrimitiveTypes.Float64:
sf.columns[i] = make([]float64, newSize)
case arrow.BinaryTypes.String:
sf.columns[i] = make([]string, newSize)
}
}
sf.added = newSize
@ -381,6 +392,11 @@ func (sf *ShardFile) SetFloatValue(col int, row int64, val float64) {
v[row-sf.beforeRows] = val
}
func (sf *ShardFile) SetStringValue(col int, row int64, val string) {
v := sf.columns[col].([]string)
v[row-sf.beforeRows] = val
}
func (sf *ShardFile) Process(cs *ChangesetRequest) error {
err := sf.process(cs)
if err != nil {
@ -394,9 +410,34 @@ func (sf *ShardFile) Process(cs *ChangesetRequest) error {
return os.Rename(rtemp+sf.executor.TableExtension(), sf.dest+sf.executor.TableExtension())
}
func (sf *ShardFile) LoadBlobs() error {
for col := 0; col < len(sf.schema.Fields()); col++ {
column := sf.table.Column(col)
switch column.DataType() {
case arrow.BinaryTypes.String:
for i, chunk := range column.Data().Chunks() {
stringData := chunk.(*array.String)
k := key{col: col, chunk: i}
for j := 0; j < stringData.Len(); j++ {
v := stringData.Value(j)
sf.strings[k] = append(sf.strings[k], v)
}
}
}
}
return nil
}
func (sf *ShardFile) ReplaceString(col, chunk, l int, s string) {
sf.strings[key{col: col, chunk: chunk}][l] = s
}
func (sf *ShardFile) process(cs *ChangesetRequest) error {
offset := 0
if sf.table != nil {
// need to load blobs prior
sf.LoadBlobs()
column := sf.table.Column(0)
resolver := dataframe.NewChunkResolver(column)
for i, rowid := range cs.ShardIds {
@ -414,6 +455,10 @@ func (sf *ShardFile) process(cs *ChangesetRequest) error {
case arrow.PrimitiveTypes.Float64:
v := column.Data().Chunk(chunk).(*array.Float64).Float64Values()
v[l] = cs.Columns[col].([]float64)[i]
case arrow.BinaryTypes.String:
// TODO(twg) 2023/01/09 How to update existing?
new := cs.Columns[col].([]string)[i]
sf.ReplaceString(col, chunk, l, new)
default:
panic(fmt.Sprintf("Unknown Type %v", column.DataType()))
}
@ -433,6 +478,8 @@ func (sf *ShardFile) process(cs *ChangesetRequest) error {
sf.SetIntValue(col, rowid, cs.Columns[col].([]int64)[i])
case arrow.PrimitiveTypes.Float64:
sf.SetFloatValue(col, rowid, cs.Columns[col].([]float64)[i])
case arrow.BinaryTypes.String:
sf.SetStringValue(col, rowid, cs.Columns[col].([]string)[i])
default:
panic(fmt.Sprintf("2 Unknown Type %v", sf.schema.Field(col).Type))
}
@ -443,15 +490,65 @@ func (sf *ShardFile) process(cs *ChangesetRequest) error {
return nil
}
type twoSlices struct {
id_slice []int
lists_slice [][]string
}
type SortByOther twoSlices
func (sbo SortByOther) Len() int {
return len(sbo.id_slice)
}
func (sbo SortByOther) Swap(i, j int) {
sbo.id_slice[i], sbo.id_slice[j] = sbo.id_slice[j], sbo.id_slice[i]
sbo.lists_slice[i], sbo.lists_slice[j] = sbo.lists_slice[j], sbo.lists_slice[i]
}
func (sbo SortByOther) Less(i, j int) bool {
return sbo.id_slice[i] < sbo.id_slice[j]
}
func (sf *ShardFile) buildFromStrings(idx int, mem memory.Allocator) []arrow.Array {
ids := make([]int, 0)
lists := make([][]string, 0)
for k, v := range sf.strings {
if k.col == idx { // ugh not ordered :(
ids = append(ids, k.chunk)
lists = append(lists, v)
}
}
// sort ids/lists
parts := twoSlices{id_slice: ids, lists_slice: lists}
sort.Sort(SortByOther(parts))
builder := array.NewStringBuilder(mem)
chunks := make([]arrow.Array, 0)
for _, v := range parts.lists_slice {
builder.AppendValues(v, nil)
newChunk := builder.NewArray()
chunks = append(chunks, newChunk)
}
return chunks
}
func (sf *ShardFile) Save(name string) error {
parts := make([]arrow.Array, 0)
mem := memory.NewGoAllocator()
for col := 0; col < len(sf.schema.Fields()); col++ {
chunks := make([]arrow.Array, 0)
if sf.table != nil {
// we append if there was existing parquet file
// we append if there was existing file
column := sf.table.Column(col)
chunks = append(chunks, column.Data().Chunks()...)
// if primative type
switch column.DataType() {
case arrow.BinaryTypes.String:
chunks = sf.buildFromStrings(col, mem)
default:
chunks = append(chunks, column.Data().Chunks()...)
}
// else binary type
}
switch sf.schema.Field(col).Type {
case arrow.PrimitiveTypes.Int64:
@ -480,6 +577,18 @@ func (sf *ShardFile) Save(name string) error {
return err
}
parts = append(parts, record)
case arrow.BinaryTypes.String:
if sf.added > 0 {
fbuild := array.NewStringBuilder(mem)
fbuild.AppendValues(sf.columns[col].([]string), nil) // TODO(twg) 2022/09/28 need to handle null
newChunk := fbuild.NewArray()
chunks = append(chunks, newChunk)
}
record, err := array.Concatenate(chunks, mem)
if err != nil {
return err
}
parts = append(parts, record)
default:
vprint.VV("UNKNOWN %T", sf.schema.Field(col).Type)
}

10
arrow.go
View file

@ -194,6 +194,8 @@ func (st *basicTable) Get(column, row int) interface{} {
case *arrow.Float64Type:
v := chunk.(*array.Float64).Float64Values()
return v[i]
case *arrow.StringType:
return chunk.(*array.String).Value(i)
}
return 0
}
@ -223,8 +225,10 @@ func builderFrom(mem memory.Allocator, dt arrow.DataType, size int64) array.Buil
bldr = array.NewFloat32Builder(mem)
case *arrow.Float64Type:
bldr = array.NewFloat64Builder(mem)
case *arrow.StringType:
bldr = array.NewStringBuilder(mem)
default:
panic(fmt.Errorf("npy2root: invalid Arrow type %v", dt))
panic(fmt.Errorf("builderFrom: invalid Arrow type %v", dt))
}
bldr.Reserve(int(size))
return bldr
@ -254,8 +258,10 @@ func appendData(bldr array.Builder, v interface{}) {
bldr.Append(v.(float32))
case *array.Float64Builder:
bldr.Append(v.(float64))
case *array.StringBuilder:
bldr.Append(v.(string))
default:
panic(fmt.Errorf("npy2root: invalid Arrow builder type %T", bldr))
panic(fmt.Errorf("appendData: invalid Arrow builder type %T", bldr))
}
}

4
batch/Dockerfile-test
View file

@ -2,10 +2,10 @@ ARG GO_VERSION=1.19
FROM golang:${GO_VERSION}
WORKDIR /go/src/github.com/molecula/featurebase/
WORKDIR /go/src/github.com/featurebasedb/featurebase/
COPY . .
WORKDIR /go/src/github.com/molecula/featurebase/batch/
WORKDIR /go/src/github.com/featurebasedb/featurebase/batch/
CMD ["go","test","-v","-mod=vendor","-tags=odbc,dynamic","./..."]

12
batch/batch.go
View file

@ -755,7 +755,7 @@ func (b *Batch) Import() error {
}()
}
defer func() {
b.importer.StatsTiming(MetricBatchImportDurationSeconds, time.Since(start), 1.0)
featurebase.SummaryBatchImportDurationSeconds.Observe(time.Since(start).Seconds())
}()
size := len(b.ids)
@ -828,7 +828,7 @@ func (b *Batch) Flush() error {
if err != nil {
b.log.Errorf("error finishing transaction: %v. trns: %+v", err, trnsl)
}
b.importer.StatsTiming(MetricBatchFlushDurationSeconds, time.Since(start), 1.0)
featurebase.SummaryBatchFlushDurationSeconds.Observe(time.Since(start).Seconds())
}()
importStart := time.Now()
@ -1188,7 +1188,7 @@ func (b *Batch) doImportShardTransactional(frags, clearFrags fragments) error {
}
}
b.importer.StatsTiming(MetricBatchShardImportBuildRequestsSeconds, time.Since(start), 1.0)
featurebase.SummaryBatchShardImportBuildRequestsSeconds.Observe(time.Since(start).Seconds())
start = time.Now()
eg := egpool.Group{PoolSize: 20}
for shard, request := range requests {
@ -1200,7 +1200,7 @@ func (b *Batch) doImportShardTransactional(frags, clearFrags fragments) error {
}
err := eg.Wait()
dur := time.Since(start)
b.importer.StatsTiming(MetricBatchShardImportDurationSeconds, dur, 1.0)
featurebase.SummaryBatchImportDurationSeconds.Observe(dur.Seconds())
b.log.Printf("import shard took: %v\n", dur)
return errors.Wrap(err, "doing shard-transactional imports")
}
@ -1245,7 +1245,7 @@ func (b *Batch) doImport(frags, clearFrags fragments) error {
}
ferr := b.importer.ImportRoaringBitmap(ctx, b.tbl.ID, fld, shard, viewMap, false)
b.log.Debugf("imp-roar %s,shard:%d,views:%d %v", field, shard, len(clearViewMap), time.Since(starty))
b.log.Debugf("imp-roar field: %s, shard:%d, views:%d %v", field, shard, len(clearViewMap), time.Since(starty))
return errors.Wrapf(ferr, "importing data for %s", field)
})
}
@ -1703,7 +1703,7 @@ func (b *Batch) importValueData() error {
start := time.Now()
fld := featurebase.FieldInfoToField(field)
err := b.importer.DoImport(ctx, b.tbl.ID, fld, shard, path, data)
b.log.Debugf("imp-vals %s,shard:%d,data:%d %v", field, shard, len(data), time.Since(start))
b.log.Debugf("imp-vals field: %s, shard: %d, data: %d %v", field.Name, shard, len(data), time.Since(start))
return errors.Wrapf(err, "importing values for field = %s", field.Name)
})
startIdx = i

25
batch/metrics.go
View file

@ -1,28 +1,3 @@
// Copyright 2022 Molecula Corp. (DBA FeatureBase).
// SPDX-License-Identifier: Apache-2.0
package batch
const (
// MetricBatchImportDurationSeconds records the full time of the
// RecordBatch.Import call. This includes starting and finishing a
// transaction, doing key translation, building fragments locally,
// importing all data, and resetting internal structures.
MetricBatchImportDurationSeconds = "batch_import_duration_seconds"
// MetricBatchFlushDurationSeconds records the full time for
// RecordBatch.Flush (if splitBatchMode is in use). This includes
// starting and finishing a transaction, importing all data, and
// resetting internal structures.
MetricBatchFlushDurationSeconds = "batch_flush_duration_seconds"
// MetricBatchShardImportBuildRequestsSeconds is the time it takes
// after making fragments to build the shard-transactional request
// objects (but not actually import them or do any network activity).
MetricBatchShardImportBuildRequestsSeconds = "batch_shard_import_build_requests_seconds"
// MetricBatchShardImportDurationSeconds is the time it takes to
// import all data for all shards in the batch using the
// shard-transactional endpoint. This does not include the time it
// takes to build the requests locally.
MetricBatchShardImportDurationSeconds = "batch_shard_import_duration_seconds"
)

51
cache.go
View file

@ -13,7 +13,6 @@ import (
"github.com/featurebasedb/featurebase/v3/lru"
pb "github.com/featurebasedb/featurebase/v3/proto"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/pkg/errors"
)
@ -41,9 +40,6 @@ type cache interface {
// Returns an ordered list of the top ranked bitmaps.
Top() []bitmapPair
// SetStats defines the stats client used in the cache.
SetStats(s stats.StatsClient)
// Clear removes everything from the cache. If possible it should leave allocated structures in place to be reused.
Clear()
}
@ -52,7 +48,6 @@ type cache interface {
type lruCache struct {
cache *lru.Cache
counts map[uint64]uint64
stats stats.StatsClient
// maxEntries is saved to support Clear which recreates the cache.
maxEntries uint32
}
@ -62,7 +57,6 @@ func newLRUCache(maxEntries uint32) *lruCache {
c := &lruCache{
cache: lru.New(int(maxEntries)),
counts: make(map[uint64]uint64),
stats: stats.NopStatsClient,
maxEntries: maxEntries,
}
c.cache.OnEvicted = c.onEvicted
@ -120,11 +114,6 @@ func (c *lruCache) Top() []bitmapPair {
return a
}
// SetStats defines the stats client used in the cache.
func (c *lruCache) SetStats(s stats.StatsClient) {
c.stats = s
}
func (c *lruCache) Clear() {
for k := range c.counts {
delete(c.counts, k)
@ -158,8 +147,6 @@ type rankCache struct {
// thresholdValue is the value of the last item in the cache
thresholdValue uint64
stats stats.StatsClient
}
// NewRankCache returns a new instance of RankCache.
@ -168,7 +155,6 @@ func NewRankCache(maxEntries uint32) *rankCache {
maxEntries: maxEntries,
thresholdBuffer: int(thresholdFactor * float64(maxEntries)),
entries: make(map[uint64]uint64),
stats: stats.NopStatsClient,
}
}
@ -229,7 +215,7 @@ func (c *rankCache) BulkAdd(id uint64, n uint64) {
// as this can take up an upbounded amount of memory. This is especially
// true when restoring shards as all rows will be added.
if len(c.entries) > int(2*c.maxEntries) {
c.stats.Count(MetricRecalculateCache, 1, 1.0)
CounterRecalculateCache.Inc()
c.recalculate()
}
}
@ -274,7 +260,7 @@ func (c *rankCache) Invalidate() {
func (c *rankCache) Recalculate() {
c.mu.Lock()
defer c.mu.Unlock()
c.stats.Count(MetricRecalculateCache, 1, 1.0)
CounterRecalculateCache.Inc()
c.recalculate()
}
@ -286,12 +272,12 @@ func (c *rankCache) invalidate() {
// This is somewhat necessary for now since recalculation is not cheap.
// The cache will remain flagged as dirty and will be recalculated if Top is called.
// This may cause unexpected memory growth, so record it in metrics for debugging purposes.
c.stats.Count(MetricInvalidateCacheSkipped, 1, 1.0)
CounterInvalidateCacheSkipped.Inc()
// Ensure that we're marked as dirty even if we weren't otherwise.
c.dirty = true
return
}
c.stats.Count(MetricInvalidateCache, 1, 1.0)
CounterInvalidateCache.Inc()
c.recalculate()
}
@ -317,7 +303,7 @@ func (c *rankCache) recalculate() {
// Store the count of the item at the threshold index.
length := len(c.rankings)
c.stats.Gauge(MetricRankCacheLength, float64(length), 1.0)
GaugeRankCacheLength.Set(float64(length))
var removeItems []bitmapPair // cached, ordered list
if length > int(c.maxEntries) {
@ -333,7 +319,7 @@ func (c *rankCache) recalculate() {
// If size is larger than the threshold then trim it.
if len(c.entries) > c.thresholdBuffer {
c.stats.Count(MetricCacheThresholdReached, 1, 1.0)
CounterCacheThresholdReached.Inc()
for _, pair := range removeItems {
delete(c.entries, pair.ID)
}
@ -343,11 +329,6 @@ func (c *rankCache) recalculate() {
c.dirty = false
}
// SetStats defines the stats client used in the cache.
func (c *rankCache) SetStats(s stats.StatsClient) {
c.stats = s
}
// Top returns an ordered list of pairs.
func (c *rankCache) Top() []bitmapPair {
c.mu.Lock()
@ -355,7 +336,7 @@ func (c *rankCache) Top() []bitmapPair {
if c.dirty {
// The cache is dirty, so we need to recalculate it to get a consistent view.
c.stats.Count(MetricReadDirtyCache, 1, 1.0)
CounterReadDirtyCache.Inc()
c.recalculate()
}
@ -606,25 +587,21 @@ func (p uint64Slice) Len() int { return len(p) }
func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }
// nopCache represents a no-op Cache implementation.
type nopCache struct {
stats stats.StatsClient
}
type nopCache struct{}
// Ensure NopCache implements Cache.
var globalNopCache cache = nopCache{
stats: stats.NopStatsClient,
}
var globalNopCache cache = nopCache{}
func (c nopCache) Add(uint64, uint64) {}
func (c nopCache) BulkAdd(uint64, uint64) {}
func (c nopCache) Get(uint64) uint64 { return 0 }
func (c nopCache) IDs() []uint64 { return []uint64{} }
func (c nopCache) Invalidate() {}
func (c nopCache) Len() int { return 0 }
func (c nopCache) Recalculate() {}
func (c nopCache) SetStats(stats.StatsClient) {}
func (c nopCache) Clear() {}
func (c nopCache) Invalidate() {}
func (c nopCache) Len() int { return 0 }
func (c nopCache) Recalculate() {}
func (c nopCache) Clear() {}
func (c nopCache) Top() []bitmapPair {
return []bitmapPair{}

222
ctl/cli.go → cli/cli.go
View file

@ -1,9 +1,7 @@
package ctl
package cli
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
@ -16,9 +14,7 @@ import (
"github.com/jedib0t/go-pretty/table"
"github.com/jedib0t/go-pretty/text"
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/featurebasedb/featurebase/v3/dax"
queryerhttp "github.com/featurebasedb/featurebase/v3/dax/queryer/http"
"github.com/featurebasedb/featurebase/v3/fbcloud"
"github.com/featurebasedb/featurebase/v3/cli/fbcloud"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/pkg/errors"
)
@ -32,6 +28,12 @@ const (
nullValue string = "NULL"
)
var (
Stdin io.ReadCloser = os.Stdin
Stdout io.Writer = os.Stdout
Stderr io.Writer = os.Stderr
)
var (
splash string = fmt.Sprintf(`FeatureBase CLI (%s)
Type "exit" to quit.
@ -55,35 +57,52 @@ type CLICommand struct {
OrganizationID string `json:"org-id"`
DatabaseID string `json:"db-id"`
queryer FBQueryer
Queryer Queryer `json:"-"`
Stdin io.ReadCloser `json:"-"`
Stdout io.Writer `json:"-"`
Stderr io.Writer `json:"-"`
}
func NewCLICommand(logdest logger.Logger) *CLICommand {
return &CLICommand{
Host: defaultHost,
HistoryPath: "",
OrganizationID: "",
DatabaseID: "",
Stdin: Stdin,
Stdout: Stdout,
Stderr: Stderr,
}
}
func (cmd *CLICommand) setupHistory() {
// If HistoryPath has already been configured (i.e. with a command flag),
// don't bother setting up the default in the home directory.
if cmd.HistoryPath != "" {
return
}
historyPath := ""
home, err := os.UserHomeDir()
if err != nil {
fmt.Printf("Error getting home directory, command history persistence will be disabled: %v\n", err)
if home, err := os.UserHomeDir(); err != nil {
cmd.Printf("Error getting home directory, command history persistence will be disabled: %v\n", err)
} else {
historyDir := filepath.Join(home, ".featurebase")
err := os.MkdirAll(historyDir, 0750)
err := os.MkdirAll(historyDir, 0o750)
if err != nil {
fmt.Printf("Creating directory for history: %v\n", err)
cmd.Printf("Creating directory for history: %v\n", err)
} else {
historyPath = filepath.Join(historyDir, "cli_history")
}
}
return &CLICommand{
Host: defaultHost,
HistoryPath: historyPath,
OrganizationID: "",
DatabaseID: "",
}
cmd.HistoryPath = historyPath
}
// printQualifiers displays the currently set OrganizationID and DatabaseID.
func (cmd *CLICommand) printQualifiers() {
fmt.Printf(" Host: %s\n Org: %s\n DB: %s\n",
cmd.Printf(" Host: %s\n Org: %s\n DB: %s\n",
hostPort(cmd.Host, cmd.Port),
cmd.OrganizationID,
cmd.DatabaseID,
@ -91,6 +110,12 @@ func (cmd *CLICommand) printQualifiers() {
}
func (cmd *CLICommand) setupClient() error {
// If the Queryer has already been set (in tests for example), don't bother
// trying to detect it.
if cmd.Queryer != nil {
return nil
}
if strings.TrimSpace(cmd.Host) == "" {
return errors.Errorf("no host provided")
}
@ -106,20 +131,20 @@ func (cmd *CLICommand) setupClient() error {
switch typ {
case featurebaseTypeStandard:
fmt.Println("Detected standard deployment")
cmd.queryer = &standardQueryer{
cmd.Printf("Detected standard deployment\n")
cmd.Queryer = &standardQueryer{
Host: cmd.Host,
Port: cmd.Port,
}
case featurebaseTypeDAX:
fmt.Println("Detected dax deployment")
cmd.queryer = &daxQueryer{
cmd.Printf("Detected dax deployment\n")
cmd.Queryer = &daxQueryer{
Host: cmd.Host,
Port: cmd.Port,
}
case featurebaseTypeCloud:
fmt.Println("Detected cloud deployment")
cmd.queryer = &fbcloud.Queryer{
cmd.Printf("Detected cloud deployment\n")
cmd.Queryer = &fbcloud.Queryer{
Host: hostPort(cmd.Host, cmd.Port),
ClientID: cmd.ClientID,
@ -222,9 +247,9 @@ func (cmd *CLICommand) detectFBType() (featurebaseType, error) {
func (cmd *CLICommand) Run(ctx context.Context) error {
// Print the splash message.
fmt.Print(splash)
err := cmd.setupClient()
if err != nil {
cmd.Printf(splash)
cmd.setupHistory()
if err := cmd.setupClient(); err != nil {
return errors.Wrap(err, "setting up client")
}
cmd.printQualifiers()
@ -234,6 +259,10 @@ func (cmd *CLICommand) Run(ctx context.Context) error {
HistoryFile: cmd.HistoryPath,
HistoryLimit: 100000,
DisableAutoSaveHistory: true,
Stdin: cmd.Stdin,
Stdout: cmd.Stdout,
Stderr: cmd.Stderr,
})
if err != nil {
return errors.Wrap(err, "getting readline")
@ -254,7 +283,7 @@ func (cmd *CLICommand) Run(ctx context.Context) error {
} else {
rl.SetPrompt(promptBegin)
// Add some white space before each new prompt.
fmt.Println()
cmd.Printf("\n")
}
// Read user provided input.
@ -263,6 +292,27 @@ func (cmd *CLICommand) Run(ctx context.Context) error {
return errors.Wrap(err, "reading line")
}
// We append a line feed at the end of each line because at this point
// we have effectively stripped any intentional line feeds (since we are
// reading a line at a time), and we don't want to do that. An example
// of an intentional line feed is in a BULK INSERT CSV STREAM like this
// example:
//
// bulk replace
// into foo (_id, age)
// map (0 id, 1 int)
// from
// x'3,33
// 4,44
// 5,55'
// with
// format 'CSV'
// input 'STREAM';
//
// We want to preserve the line feeds that are contained in the x''
// block; those are intentional as they demarc records within the csv.
line += "\n"
if !inMidCommand {
// Handle the exit command.
if line == exitCommand || line == exitCommand+terminationChar {
@ -284,7 +334,7 @@ func (cmd *CLICommand) Run(ctx context.Context) error {
for i, part := range parts {
partIsFinal := i == len(parts)-1
partIsBlank := part == ""
partIsBlank := strings.TrimSpace(part) == ""
if partIsBlank && partIsFinal {
continue
@ -315,7 +365,7 @@ func (cmd *CLICommand) Run(ctx context.Context) error {
err = rl.SaveHistory(strings.Join(cmd.commands, "; ") + ";")
if err != nil {
fmt.Printf("Couldn't save history: %v\n", err)
cmd.Printf("Couldn't save history: %v\n", err)
}
if err := cmd.executeCommands(ctx); err != nil {
@ -330,14 +380,10 @@ func appendCommand(orig string, part string) string {
if orig == "" {
return part
} else {
return orig + " " + part
return orig + part
}
}
type FBQueryer interface {
Query(org, db, sql string) (*featurebase.WireQueryResponse, error)
}
func (cmd *CLICommand) executeCommands(ctx context.Context) error {
// Clear out the buffered commands on any exit from this method.
defer func() {
@ -352,12 +398,12 @@ func (cmd *CLICommand) executeCommands(ctx context.Context) error {
continue
}
sqlResponse, err := cmd.queryer.Query(cmd.OrganizationID, cmd.DatabaseID, sql)
sqlResponse, err := cmd.Queryer.Query(cmd.OrganizationID, cmd.DatabaseID, sql)
if err != nil {
fmt.Printf("making query: %v\n", err)
cmd.Printf("making query: %v\n", err)
continue
}
err = writeOut(sqlResponse, os.Stdout)
err = writeOut(sqlResponse, cmd.Stdout, cmd.Stderr)
if err != nil {
return errors.Wrap(err, "writing out response")
}
@ -366,6 +412,12 @@ func (cmd *CLICommand) executeCommands(ctx context.Context) error {
return nil
}
// Printf is a helper method which sends the given payload to stdout.
func (cmd *CLICommand) Printf(format string, a ...any) {
out := fmt.Sprintf(format, a...)
cmd.Stdout.Write([]byte(out))
}
// handleIfNonSQLCommand will handle special case command like "SET ..." and
// "USE ...". If the sql command matches one of these conditions and is handled,
// the bool returned will be true;
@ -431,19 +483,19 @@ func writeWarnings(r *featurebase.WireQueryResponse, w io.Writer) error {
return nil
}
func writeOut(r *featurebase.WireQueryResponse, w io.Writer) error {
func writeOut(r *featurebase.WireQueryResponse, wOut io.Writer, wErr io.Writer) error {
if r == nil {
return errors.New("attempt to write out nil response")
}
if r.Error != "" {
if _, err := w.Write([]byte("Error: " + r.Error + "\n")); err != nil {
if _, err := wErr.Write([]byte("Error: " + r.Error + "\n")); err != nil {
return errors.Wrapf(err, "writing error: %s", r.Error)
}
return writeWarnings(r, w)
return writeWarnings(r, wOut)
}
t := table.NewWriter()
t.SetOutputMirror(w)
t.SetOutputMirror(wOut)
// Don't uppercase the header values.
t.Style().Format.Header = text.FormatDefault
@ -461,7 +513,7 @@ func writeOut(r *featurebase.WireQueryResponse, w io.Writer) error {
}
t.Render()
err := writeWarnings(r, w)
err := writeWarnings(r, wOut)
if err != nil {
return err
}
@ -474,7 +526,7 @@ func writeOut(r *featurebase.WireQueryResponse, w io.Writer) error {
lifeAffirmingMessage = " (Sorry! That took longer than expected 😭)"
}
if _, err := w.Write([]byte(fmt.Sprintf("\nExecution time: %dμs%s\n", r.ExecutionTime, lifeAffirmingMessage))); err != nil {
if _, err := wOut.Write([]byte(fmt.Sprintf("\nExecution time: %dμs%s\n", r.ExecutionTime, lifeAffirmingMessage))); err != nil {
return errors.Wrapf(err, "writing execution time: %s", r.Error)
}
@ -488,81 +540,3 @@ func schemaToRow(schema featurebase.WireQuerySchema) []interface{} {
}
return ret
}
// Ensure type implements interface.
var _ FBQueryer = (*standardQueryer)(nil)
// standardQueryer supports a standard featurebase deployment hitting the /sql
// endpoint with a payload containing only the sql statement.
type standardQueryer struct {
Host string
Port string
}
func (qryr *standardQueryer) Query(org, db, sql string) (*featurebase.WireQueryResponse, error) {
buf := bytes.Buffer{}
url := fmt.Sprintf("%s/sql", hostPort(qryr.Host, qryr.Port))
buf.Write([]byte(sql))
resp, err := http.Post(url, "application/json", &buf)
if err != nil {
return nil, errors.Wrapf(err, "posting query")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading response")
}
sqlResponse := &featurebase.WireQueryResponse{}
// TODO(tlt): switch this back once all responses are typed
// if err := json.Unmarshal(fullbod, sqlResponse); err != nil {
if err := sqlResponse.UnmarshalJSONTyped(fullbod, true); err != nil {
return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod)
}
return sqlResponse, nil
}
// Ensure type implements interface.
var _ FBQueryer = (*daxQueryer)(nil)
// daxQueryer is similar to the standardQueryer except that it hits a different
// endpoint, and its payload is a json object which includes, in addition to the
// sql statement, things like org and db.
type daxQueryer struct {
Host string
Port string
}
func (qryr *daxQueryer) Query(org, db, sql string) (*featurebase.WireQueryResponse, error) {
buf := bytes.Buffer{}
url := fmt.Sprintf("%s/queryer/sql", hostPort(qryr.Host, qryr.Port))
sqlReq := &queryerhttp.SQLRequest{
OrganizationID: dax.OrganizationID(org),
DatabaseID: dax.DatabaseID(db),
SQL: sql,
}
if err := json.NewEncoder(&buf).Encode(sqlReq); err != nil {
return nil, errors.Wrapf(err, "encoding sql request: %s", sql)
}
resp, err := http.Post(url, "application/json", &buf)
if err != nil {
return nil, errors.Wrapf(err, "posting query")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading response")
}
sqlResponse := &featurebase.WireQueryResponse{}
// TODO(tlt): switch this back once all responses are typed
// if err := json.Unmarshal(fullbod, sqlResponse); err != nil {
if err := sqlResponse.UnmarshalJSONTyped(fullbod, true); err != nil {
return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod)
}
return sqlResponse, nil
}

199
cli/cli_test.go Normal file
View file

@ -0,0 +1,199 @@
package cli_test
import (
"context"
"io"
"strings"
"sync"
"testing"
"time"
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/featurebasedb/featurebase/v3/cli"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/pkg/errors"
"github.com/stretchr/testify/assert"
)
func TestCLI(t *testing.T) {
t.Run("Input", func(t *testing.T) {
ctx := context.Background()
capture := newCapture(t)
cli := cli.NewCLICommand(logger.StderrLogger)
cli.Stdin = capture
cli.Stdout = capture
cli.Queryer = capture
go func() {
assert.NoError(t, cli.Run(ctx))
}()
none := []string{}
// One statement, one line.
capture.Assert("one;", []string{`one`})
// One statement, multiple lines.
capture.Assert("one", none)
capture.Assert(" two ", none)
capture.Assert("three;", []string{`one
two
three`})
// Multiple statements, one line.
capture.Assert("foo; bar;", []string{`foo`, `bar`})
// Multiple statements, multiple lines.
capture.Assert("a1", none)
capture.Assert("a2; b1", []string{`a1
a2`})
capture.Assert("b2;", []string{`b1
b2`})
// Blank lines.
capture.Assert("one", none)
capture.Assert("", none)
capture.Assert("three;", []string{`one
three`})
// Just a semi-colon.
capture.Assert(";", none)
// Multi-line with just a semi-colon.
capture.Assert("one", none)
capture.Assert(";", []string{`one`})
// Ensure a clean exit with no errors.
assert.NoError(t, capture.Exit())
})
}
////////////////////////////////////////////////////////
// Ensure type implementes interface.
var _ io.ReadCloser = (*capture)(nil)
var _ io.Writer = (*capture)(nil)
var _ cli.Queryer = (*capture)(nil)
// capture implements the various CLI interfaces in order to capture test input
// and submit it as though that input were being read from the command line. It
// also captures calls made to the Queryer.Query method and ensures the sql the
// contain is expected.
type capture struct {
t *testing.T
// ch is a channel of strings (one line at a time) of CLI input.
ch chan string
mu sync.RWMutex
sqls []string
// queryDone will receive an event any time the Query method is called and
// has completed. This is to tell the Assert method that it's safe to
// compare the sqls slice.
queryDone chan struct{}
asserting chan struct{}
err error
}
func newCapture(t *testing.T) *capture {
return &capture{
t: t,
ch: make(chan string),
sqls: make([]string, 0),
queryDone: make(chan struct{}),
}
}
func (c *capture) Exit() error {
c.sendLine("exit")
c.mu.RLock()
defer c.mu.RUnlock()
return c.err
}
func (c *capture) Assert(in string, out []string) {
c.asserting = make(chan struct{})
c.sendLine(in)
// Wait for the CLI command to complete processing the input and send the
// sql to Query() by blocking on the queryDone channel. Because Query gets
// called for every sql statement in the input, an input resulting in
// multiple sql statements needs to wait for all expected queries to
// complete. A timeout is included to this so it doesn't deadlock in the
// case where Query is expected to be called, but isn't; after the timeout,
// the test should fail completely. In summary: we wait on queryDone the
// number of sql statements we expect. If we receive fewer than expected,
// the timeout will occur. If we receive more than expected, the Query()
// method will effectively deadlock, reach its own timout, then write to
// capture.err, which will be reported upon Exit().
for range out {
select {
case <-c.queryDone:
case <-time.After(2 * time.Second):
c.t.Fatalf("expected Query() to be called")
}
}
close(c.asserting)
c.mu.Lock()
defer c.mu.Unlock()
assert.Equal(c.t, out, c.sqls)
// Reset the slice.
c.sqls = c.sqls[:0]
}
// sendLine sends the given string as a line input to the CLI command. It
// appends a line feed to the end of string in order to mimic the user hitting
// the return key.
func (c *capture) sendLine(s string) {
// Add a line feed before putting s on the channel in order to mimic the
// user hitting the return key.
c.ch <- s + "\n"
}
// Read is read by the CLI in place of user input. It effectively sends lines of
// input to the CLI, getting each line to be sent off the channel.
func (c *capture) Read(b []byte) (n int, err error) {
s := <-c.ch
return strings.NewReader(s).Read(b)
}
func (c *capture) Close() error {
close(c.ch)
return nil
}
// Write is called with anything written to output. This would included results
// from calling Query() under normal, non-testing conditions, as well as other
// informational text sent to output, such as the splash message.
func (c *capture) Write(b []byte) (n int, err error) {
return 0, nil
}
// Query is called by the CLI command once a full SQL statement is received
// (signified by the terminator: `;`).
func (c *capture) Query(org, db, sql string) (*featurebase.WireQueryResponse, error) {
c.mu.Lock()
c.sqls = append(c.sqls, sql)
c.mu.Unlock()
select {
case c.queryDone <- struct{}{}:
case <-c.asserting:
c.mu.Lock()
c.err = errors.Errorf("unexpected query: %s", sql)
c.mu.Unlock()
}
return &featurebase.WireQueryResponse{}, nil
}

0
fbcloud/auth.go → cli/fbcloud/auth.go
View file

0
fbcloud/client.go → cli/fbcloud/client.go
View file

96
cli/queryer.go Normal file
View file

@ -0,0 +1,96 @@
package cli
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/featurebasedb/featurebase/v3/dax"
queryerhttp "github.com/featurebasedb/featurebase/v3/dax/queryer/http"
"github.com/pkg/errors"
)
type Queryer interface {
Query(org, db, sql string) (*featurebase.WireQueryResponse, error)
}
// Ensure type implements interface.
var _ Queryer = (*standardQueryer)(nil)
// standardQueryer supports a standard featurebase deployment hitting the /sql
// endpoint with a payload containing only the sql statement.
type standardQueryer struct {
Host string
Port string
}
func (qryr *standardQueryer) Query(org, db, sql string) (*featurebase.WireQueryResponse, error) {
buf := bytes.Buffer{}
url := fmt.Sprintf("%s/sql", hostPort(qryr.Host, qryr.Port))
buf.Write([]byte(sql))
resp, err := http.Post(url, "application/json", &buf)
if err != nil {
return nil, errors.Wrapf(err, "posting query")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading response")
}
sqlResponse := &featurebase.WireQueryResponse{}
// TODO(tlt): switch this back once all responses are typed
// if err := json.Unmarshal(fullbod, sqlResponse); err != nil {
if err := sqlResponse.UnmarshalJSONTyped(fullbod, true); err != nil {
return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod)
}
return sqlResponse, nil
}
// Ensure type implements interface.
var _ Queryer = (*daxQueryer)(nil)
// daxQueryer is similar to the standardQueryer except that it hits a different
// endpoint, and its payload is a json object which includes, in addition to the
// sql statement, things like org and db.
type daxQueryer struct {
Host string
Port string
}
func (qryr *daxQueryer) Query(org, db, sql string) (*featurebase.WireQueryResponse, error) {
buf := bytes.Buffer{}
url := fmt.Sprintf("%s/queryer/sql", hostPort(qryr.Host, qryr.Port))
sqlReq := &queryerhttp.SQLRequest{
OrganizationID: dax.OrganizationID(org),
DatabaseID: dax.DatabaseID(db),
SQL: sql,
}
if err := json.NewEncoder(&buf).Encode(sqlReq); err != nil {
return nil, errors.Wrapf(err, "encoding sql request: %s", sql)
}
resp, err := http.Post(url, "application/json", &buf)
if err != nil {
return nil, errors.Wrapf(err, "posting query")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading response")
}
sqlResponse := &featurebase.WireQueryResponse{}
// TODO(tlt): switch this back once all responses are typed
// if err := json.Unmarshal(fullbod, sqlResponse); err != nil {
if err := sqlResponse.UnmarshalJSONTyped(fullbod, true); err != nil {
return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod)
}
return sqlResponse, nil
}

16
client/client.go
View file

@ -29,7 +29,6 @@ import (
"github.com/featurebasedb/featurebase/v3/pb"
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/vprint"
"github.com/opentracing/opentracing-go"
"github.com/pkg/errors"
@ -52,7 +51,6 @@ type Client struct {
manualFragmentNode *fragmentNode
manualServerURI *pnet.URI
tracer opentracing.Tracer
Stats stats.StatsClient
// An exponential backoff algorithm retries requests exponentially (if an HTTP request fails),
// increasing the waiting time between retries up to a maximum backoff time.
maxBackoff time.Duration
@ -211,11 +209,6 @@ func newClientWithOptions(options *ClientOptions) *Client {
} else {
c.tracer = options.tracer
}
if options.stats == nil {
c.Stats = stats.NopStatsClient
} else {
c.Stats = options.stats
}
c.maxRetries = *options.retries
c.maxBackoff = 2 * time.Minute
@ -1359,7 +1352,6 @@ type ClientOptions struct {
manualServerAddress bool
tracer opentracing.Tracer
retries *int
stats stats.StatsClient
nat map[pnet.URI]pnet.URI
pathPrefix string
}
@ -1445,14 +1437,6 @@ func OptClientRetries(retries int) ClientOption {
}
}
// OptClientStatsClient sets a stats client, such as Prometheus
func OptClientStatsClient(stats stats.StatsClient) ClientOption {
return func(options *ClientOptions) error {
options.stats = stats
return nil
}
}
// OptClientNAT sets a NAT map used to translate the advertised URI to something
// else (for example, when accessing pilosa running in docker).
func OptClientNAT(nat map[string]string) ClientOption {

4
client/importer.go
View file

@ -328,7 +328,3 @@ func (i *importer) EncodeImport(ctx context.Context, tid dax.TableID, fld *dax.F
func (i *importer) DoImport(ctx context.Context, tid dax.TableID, fld *dax.Field, shard uint64, path string, data []byte) error {
return i.client.DoImport(string(tid), shard, path, data)
}
func (i *importer) StatsTiming(name string, value time.Duration, rate float64) {
i.client.Stats.Timing(name, value, rate)
}

32
cmd/cli.go
View file

@ -2,34 +2,34 @@
package cmd
import (
"github.com/featurebasedb/featurebase/v3/ctl"
"github.com/featurebasedb/featurebase/v3/cli"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/spf13/cobra"
)
var cli *ctl.CLICommand
var cliCmd *cli.CLICommand
// newCLICommand runs the FeatureBase CLI subcommand for ingesting bulk data.
func newCLICommand(logdest logger.Logger) *cobra.Command {
cli = ctl.NewCLICommand(logdest)
cliCmd := &cobra.Command{
cliCmd = cli.NewCLICommand(logdest)
cobraCmd := &cobra.Command{
Use: "cli",
Short: "Query FB with SQL3 from the command line",
Long: ``,
RunE: usageErrorWrapper(cli),
RunE: usageErrorWrapper(cliCmd),
}
flags := cliCmd.Flags()
flags.StringVarP(&cli.Host, "host", "", cli.Host, "hostname of FeatureBase.")
flags.StringVarP(&cli.Port, "port", "", cli.Port, "port of FeatureBase.")
flags.StringVar(&cli.HistoryPath, "history-path", cli.HistoryPath, "path for history files.")
flags.StringVar(&cli.OrganizationID, "org-id", cli.OrganizationID, "OrganizationID.")
flags.StringVar(&cli.DatabaseID, "db-id", cli.DatabaseID, "DatabaseID.")
flags := cobraCmd.Flags()
flags.StringVarP(&cliCmd.Host, "host", "", cliCmd.Host, "hostname of FeatureBase.")
flags.StringVarP(&cliCmd.Port, "port", "", cliCmd.Port, "port of FeatureBase.")
flags.StringVar(&cliCmd.HistoryPath, "history-path", cliCmd.HistoryPath, "path for history files.")
flags.StringVar(&cliCmd.OrganizationID, "org-id", cliCmd.OrganizationID, "OrganizationID.")
flags.StringVar(&cliCmd.DatabaseID, "db-id", cliCmd.DatabaseID, "DatabaseID.")
flags.StringVar(&cli.ClientID, "client-id", cli.ClientID, "Cognito Client ID for FeatureBase Cloud access.")
flags.StringVar(&cli.Region, "region", cli.Region, "Cloud region for FeatureBase Cloud access (e.g. us-east-2).")
flags.StringVar(&cli.Email, "email", cli.Email, "Email address for FeatureBase Cloud access.")
flags.StringVar(&cli.Password, "password", cli.Password, "Password for FeatureBase Cloud access.")
flags.StringVar(&cliCmd.ClientID, "client-id", cliCmd.ClientID, "Cognito Client ID for FeatureBase Cloud access.")
flags.StringVar(&cliCmd.Region, "region", cliCmd.Region, "Cloud region for FeatureBase Cloud access (e.g. us-east-2).")
flags.StringVar(&cliCmd.Email, "email", cliCmd.Email, "Email address for FeatureBase Cloud access.")
flags.StringVar(&cliCmd.Password, "password", cliCmd.Password, "Password for FeatureBase Cloud access.")
return cliCmd
return cobraCmd
}

14
ctl/dataframe-csv-loader.go
View file

@ -31,6 +31,7 @@ var (
func init() {
gob.Register(arrow.PrimitiveTypes.Int64)
gob.Register(arrow.PrimitiveTypes.Float64)
gob.Register(arrow.BinaryTypes.String)
}
// TODO(rdp): add refresh token to this as well
@ -154,7 +155,7 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) {
}
fields := make([]arrow.Field, 0)
fields = append(fields, arrow.Field{Name: "_ID", Type: arrow.PrimitiveTypes.Int64})
fileScanner := bufio.NewScanner(readFile)
fileScanner := bufio.NewScanner(readFile) // TODO(twg) 2023/01/11 need to convert to the go CSV reader for more robust string support
fileScanner.Split(bufio.ScanLines)
// need for really long csv lines
var buf []byte
@ -176,6 +177,8 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) {
fields = append(fields, arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Int64})
} else if strings.HasSuffix(col, "__F") {
fields = append(fields, arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Float64})
} else if strings.HasSuffix(col, "__S") {
fields = append(fields, arrow.Field{Name: name, Type: arrow.BinaryTypes.String})
} else {
return errors.New("invalid format for type")
}
@ -268,6 +271,8 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) {
continue
}
shardFile.SetFloatValue(i, shardRow, val)
case arrow.BinaryTypes.String:
shardFile.SetStringValue(i, shardRow, rec)
default:
return errors.New("unhandled arrow type type")
}
@ -314,6 +319,11 @@ func (s *ShardDiff) SetFloatValue(col int, row int64, val float64) {
s.columns[col] = append(slice, val)
}
func (s *ShardDiff) SetStringValue(col int, row int64, val string) {
slice := s.columns[col].([]string)
s.columns[col] = append(slice, val)
}
func (s *ShardDiff) SetNulll(col int, row uint64) {
s.null[pair{col: col, row: row}] = struct{}{}
}
@ -329,6 +339,8 @@ func (s *ShardDiff) Setup(schema *arrow.Schema) {
s.columns = append(s.columns, make([]int64, 0))
case arrow.PrimitiveTypes.Float64:
s.columns = append(s.columns, make([]float64, 0))
case arrow.BinaryTypes.String:
s.columns = append(s.columns, make([]string, 0))
}
}
}

27
ctl/dataframe-csv-loader_test.go
View file

@ -66,6 +66,33 @@ func TestDataframeCsvLoaderCommand(t *testing.T) {
cm.Path = file.Name()
cm.Index = index
err = cm.Run(ctx)
if err != nil {
t.Fatalf("DataframeCsvLoader Run doesn't work: %s", err)
}
})
t.Run("strings", func(t *testing.T) {
cmLog := logger.NewStandardLogger(io.Discard)
cm := NewDataframeCsvLoaderCommand(cmLog)
file, err := testhook.TempFile(t, "import_string.csv")
if err != nil {
t.Fatalf("creating tempfile: %v", err)
}
_, err = file.Write([]byte("id,val__S\nA,ab\nB,cd\nC,ef"))
if err != nil {
t.Fatalf("writing to tempfile: %v", err)
}
ctx := context.Background()
if err != nil {
t.Fatal(err)
}
index := "strings"
cmd.API.CreateIndex(ctx, index, pilosa.IndexOptions{Keys: true})
cm.Host = cmd.API.Node().URI.HostPort()
cm.Path = file.Name()
cm.Index = index
err = cm.Run(ctx)
if err != nil {
t.Fatalf("DataframeCsvLoader Run doesn't work: %s", err)

45
dataframe_test.go
View file

@ -51,19 +51,20 @@ func TestExecutor_Apply(t *testing.T) {
}
t.Run("dataframe ingest", func(t *testing.T) {
// func (c *Client) ApplyDataframeChangeset(indexName string, cr *pilosa.ChangesetRequest, shard uint64) (map[string]interface{}, error) {
cr := &pilosa.ChangesetRequest{}
// for each row a list of columns
cr.Columns = []interface{}{
[]int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
[]int64{2, 4, 6, 8, 10, 12, 14, 16, 18, 20},
[]float64{1, 1.414, 1.732, 2, 2.236, 2.449, 2.646, 2.828, 3, 3.162},
[]string{"A", "B", "C", "D", "E", "F", "G", "H", "I", "J"},
}
cr.ShardIds = []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
cr.SimpleSchema = []pilosa.NameType{
{Name: "_ID", DataType: arrow.PrimitiveTypes.Int64},
{Name: "ival", DataType: arrow.PrimitiveTypes.Int64},
{Name: "fval", DataType: arrow.PrimitiveTypes.Float64},
{Name: "sval", DataType: arrow.BinaryTypes.String},
}
shard := uint64(0)
err := api.ApplyDataframeChangeset(ctx, indexName, cr, shard)
@ -72,7 +73,7 @@ func TestExecutor_Apply(t *testing.T) {
}
})
t.Run("dataframe schema", func(t *testing.T) {
expectedJSON := `[{"Name":"_ID","Type":"int64"},{"Name":"ival","Type":"int64"},{"Name":"fval","Type":"float64"}]`
expectedJSON := `[{"Name":"_ID","Type":"int64"},{"Name":"ival","Type":"int64"},{"Name":"fval","Type":"float64"},{"Name":"sval","Type":"utf8"}]`
parts, err := api.GetDataframeSchema(ctx, indexName)
if err != nil {
t.Fatal(err)
@ -138,7 +139,7 @@ func TestExecutor_Apply(t *testing.T) {
if res, err := api.Query(ctx, &pilosa.QueryRequest{Index: indexName, Query: pql}); err != nil {
t.Fatal(err)
} else {
expectedJSON := `{"Results":[{"_ID":[2,4,6],"fval":[1.414,2,2.449],"ival":[4,8,12]}],"Err":null,"Profile":null}`
expectedJSON := `{"Results":[{"_ID":[2,4,6],"fval":[1.414,2,2.449],"ival":[4,8,12],"sval":["B","D","F"]}],"Err":null,"Profile":null}`
w := new(bytes.Buffer)
if err := json.NewEncoder(w).Encode(res); err != nil {
t.Fatal(err)
@ -165,6 +166,44 @@ func TestExecutor_Apply(t *testing.T) {
}
}
})
t.Run("dataframe ingest update", func(t *testing.T) {
cr := &pilosa.ChangesetRequest{}
// for each row a list of columns
cr.Columns = []interface{}{
[]int64{1},
[]int64{20},
[]float64{10},
[]string{"A2"},
}
cr.ShardIds = []int64{1}
cr.SimpleSchema = []pilosa.NameType{
{Name: "_ID", DataType: arrow.PrimitiveTypes.Int64},
{Name: "ival", DataType: arrow.PrimitiveTypes.Int64},
{Name: "fval", DataType: arrow.PrimitiveTypes.Float64},
{Name: "sval", DataType: arrow.BinaryTypes.String},
}
shard := uint64(0)
err := api.ApplyDataframeChangeset(ctx, indexName, cr, shard)
if err != nil {
t.Fatal(err)
}
})
t.Run("dataframe arrow filter with header", func(t *testing.T) {
pql := `Arrow(ConstRow(columns=[1]),header=["ival","fval","sval"])`
if res, err := api.Query(ctx, &pilosa.QueryRequest{Index: indexName, Query: pql}); err != nil {
t.Fatal(err)
} else {
expectedJSON := `{"Results":[{"_ID":[1],"fval":[10],"ival":[20],"sval":["A2"]}],"Err":null,"Profile":null}`
w := new(bytes.Buffer)
if err := json.NewEncoder(w).Encode(res); err != nil {
t.Fatal(err)
}
got := strings.Trim(w.String(), "\t \n")
if strings.Compare(got, expectedJSON) != 0 {
t.Fatalf("expected: %v got: %v", expectedJSON, got)
}
}
})
t.Run("dataframe delete", func(t *testing.T) {
err := api.DeleteDataframe(ctx, indexName)
if err != nil {

12
dax/README.md
View file

@ -10,7 +10,7 @@ the Controller, as well as FeatureBase and IDK-based ingesters.
The DAX test currently requires docker images for: `featurebase` and `datagen`.
If at any point you run into problems with go mod failing to reference a private
repo, make sure that you have `gitlab.com/molecula` in your `GOPRIVATE`
repo, make sure that you have `gitlab.com/featurebasedb` in your `GOPRIVATE`
environment variable.
Note that during the docker image build step, `go mod vendor` is run, which
@ -24,7 +24,7 @@ These may no longer be relevant.
I needed to but this in my `~/.profile` file:
```export GOPRIVATE=github.com/molecula,gitlab.com/molecula```
```export GOPRIVATE=github.com/featurebasedb,gitlab.com/featurebasedb```
And this in my `~/.gitconfig`
@ -41,9 +41,9 @@ Then `make docker` ran successfully.
### Build the FeatureBase docker image
- Check out the
[dax](https://github.com/molecula/featurebase/tree/dax)
[dax](https://github.com/featurebasedb/featurebase/tree/dax)
branch of the
[featurebase](https://github.com/molecula/featurebase) repository.
[featurebase](https://github.com/featurebasedb/featurebase) repository.
- Run `make docker-image-featurebase` to build the docker image
- You should now have an image in docker named `dax/featurebase` with the tag `latest`.
@ -57,8 +57,8 @@ Then `make docker` ran successfully.
## Running the tests
- Check out the
[dax](https://github.com/molecula/featurebase/tree/dax)
[dax](https://github.com/featurebasedb/featurebase/tree/dax)
branch of the
[featurebase](https://github.com/molecula/featurebase) repository.
[featurebase](https://github.com/featurebasedb/featurebase) repository.
- Change into the `dax` directory: `cd dax`
- Run `make test-integration`.

7
dax/address.go
View file

@ -148,6 +148,13 @@ func (a addr) hostPortPath() string {
return ret
}
// Addresses is a sortable slice of Address.
type Addresses []Address
func (a Addresses) Len() int { return len(a) }
func (a Addresses) Less(i, j int) bool { return a[i] < a[j] }
func (a Addresses) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
// AddressManager is an interface for any service which needs to maintain a list
// of addresses, and receive add/remove address requests from other services.
type AddressManager interface {

7
dax/boltdb/boltdb.go
View file

@ -125,6 +125,7 @@ func (db *DB) Close() (err error) {
// BeginTx starts a transaction and returns a wrapper Tx type. This type
// provides a reference to the database and a fixed timestamp at the start of
// the transaction. The timestamp allows us to mock time during tests as well.
// The wrapper also contains the context.
func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) {
tx, err := db.db.Begin(writable)
if err != nil {
@ -134,6 +135,7 @@ func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) {
// Return wrapper Tx that includes the transaction start time.
return &Tx{
Tx: tx,
ctx: ctx,
db: db,
now: db.Now().UTC().Truncate(time.Second),
}, nil
@ -142,10 +144,15 @@ func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) {
// Tx wraps the SQL Tx object to provide a timestamp at the start of the transaction.
type Tx struct {
*bolt.Tx
ctx context.Context
db *DB
now time.Time
}
func (tx *Tx) Context() context.Context {
return tx.ctx
}
func (db *DB) Path() string {
return db.filePath
}

16
dax/boltdb/directiveversion.go
View file

@ -1,7 +1,6 @@
package boltdb
import (
"context"
"encoding/binary"
"github.com/featurebasedb/featurebase/v3/dax"
@ -32,14 +31,13 @@ func NewDirectiveVersion(db *DB) *DirectiveVersion {
}
}
func (d *DirectiveVersion) Increment(ctx context.Context, delta uint64) (uint64, error) {
tx, err := d.db.BeginTx(ctx, true)
if err != nil {
return 0, errors.Wrap(err, "getting transaction")
func (d *DirectiveVersion) Increment(tx dax.Transaction, delta uint64) (uint64, error) {
txx, ok := tx.(*Tx)
if !ok {
return 0, dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
bkt := tx.Bucket(bucketDirective)
bkt := txx.Bucket(bucketDirective)
if bkt == nil {
return 0, errors.Errorf(ErrFmtBucketNotFound, bucketDirective)
}
@ -58,9 +56,5 @@ func (d *DirectiveVersion) Increment(ctx context.Context, delta uint64) (uint64,
return 0, errors.Wrap(err, "putting next directive version")
}
if err := tx.Commit(); err != nil {
return 0, err
}
return nextVersion, nil
}

51
dax/boltdb/node.go
View file

@ -2,7 +2,6 @@ package boltdb
import (
"bytes"
"context"
"encoding/json"
"fmt"
@ -39,14 +38,13 @@ func NewNodeService(db *DB, logger logger.Logger) *NodeService {
}
}
func (s *NodeService) CreateNode(ctx context.Context, addr dax.Address, node *dax.Node) error {
tx, err := s.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "getting transaction")
func (s *NodeService) CreateNode(tx dax.Transaction, addr dax.Address, node *dax.Node) error {
txx, ok := tx.(*Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNodes)
bkt := txx.Bucket(bucketNodes)
if bkt == nil {
return errors.Errorf(ErrFmtBucketNotFound, bucketNodes)
}
@ -60,17 +58,16 @@ func (s *NodeService) CreateNode(ctx context.Context, addr dax.Address, node *da
return errors.Wrap(err, "putting node")
}
return tx.Commit()
return nil
}
func (s *NodeService) ReadNode(ctx context.Context, addr dax.Address) (*dax.Node, error) {
tx, err := s.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
func (s *NodeService) ReadNode(tx dax.Transaction, addr dax.Address) (*dax.Node, error) {
txx, ok := tx.(*Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNodes)
bkt := txx.Bucket(bucketNodes)
if bkt == nil {
return nil, errors.Errorf(ErrFmtBucketNotFound, bucketNodes)
}
@ -88,14 +85,13 @@ func (s *NodeService) ReadNode(ctx context.Context, addr dax.Address) (*dax.Node
return node, nil
}
func (s *NodeService) DeleteNode(ctx context.Context, addr dax.Address) error {
tx, err := s.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "beginning tx")
func (s *NodeService) DeleteNode(tx dax.Transaction, addr dax.Address) error {
txx, ok := tx.(*Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNodes)
bkt := txx.Bucket(bucketNodes)
if bkt == nil {
return errors.Errorf(ErrFmtBucketNotFound, bucketNodes)
}
@ -104,17 +100,16 @@ func (s *NodeService) DeleteNode(ctx context.Context, addr dax.Address) error {
return errors.Wrapf(err, "deleting node key: %s", addressKey(addr))
}
return tx.Commit()
return nil
}
func (s *NodeService) Nodes(ctx context.Context) ([]*dax.Node, error) {
tx, err := s.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "getting tx")
func (s *NodeService) Nodes(tx dax.Transaction) ([]*dax.Node, error) {
txx, ok := tx.(*Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
nodes, err := s.getNodes(ctx, tx)
nodes, err := s.getNodes(txx)
if err != nil {
return nil, errors.Wrap(err, "getting nodes")
}
@ -122,7 +117,7 @@ func (s *NodeService) Nodes(ctx context.Context) ([]*dax.Node, error) {
return nodes, nil
}
func (s *NodeService) getNodes(ctx context.Context, tx *Tx) ([]*dax.Node, error) {
func (s *NodeService) getNodes(tx *Tx) ([]*dax.Node, error) {
c := tx.Bucket(bucketNodes).Cursor()
// Deserialize rows into Node objects.

14
dax/boltdb/node_test.go
View file

@ -35,21 +35,27 @@ func TestNodeService(t *testing.T) {
},
}
tx, err := db.BeginTx(ctx, true)
assert.NoError(t, err)
defer tx.Rollback()
// Create node.
assert.NoError(t, ns.CreateNode(ctx, node1.Address, node1))
assert.NoError(t, ns.CreateNode(tx, node1.Address, node1))
// Read node.
n, err := ns.ReadNode(ctx, node1.Address)
n, err := ns.ReadNode(tx, node1.Address)
assert.NoError(t, err)
assert.Equal(t, node1, n)
// Delete node.
assert.NoError(t, ns.DeleteNode(ctx, node1.Address))
assert.NoError(t, ns.DeleteNode(tx, node1.Address))
// Read node.
_, err = ns.ReadNode(ctx, node1.Address)
_, err = ns.ReadNode(tx, node1.Address)
if assert.Error(t, err) {
assert.True(t, errors.Is(err, dax.ErrNodeDoesNotExist))
}
assert.NoError(t, tx.Commit())
})
}

6
dax/directive.go
View file

@ -1,7 +1,5 @@
package dax
import "context"
// Directive contains the instructions, sent from MDS, which a compute node is
// to follow. A Directive is typically JSON-encoded and POSTed to a compute
// node's `/directive` endpoint.
@ -22,7 +20,7 @@ type Directive struct {
}
type DirectiveVersion interface {
Increment(ctx context.Context, delta uint64) (uint64, error)
Increment(tx Transaction, delta uint64) (uint64, error)
}
// DirectiveMethod is used to tell the compute node how it should handle the
@ -177,5 +175,5 @@ func (d *Directive) IsEmpty() bool {
type Directives []*Directive
func (d Directives) Len() int { return len(d) }
func (d Directives) Less(i, j int) bool { return d[i].Address.String() < d[j].Address.String() }
func (d Directives) Less(i, j int) bool { return d[i].Version < d[j].Version }
func (d Directives) Swap(i, j int) { d[i], d[j] = d[j], d[i] }

26
dax/errors.go
View file

@ -7,6 +7,9 @@ import (
)
const (
ErrDatabaseIDExists errors.Code = "DatabaseIDExists"
ErrDatabaseIDDoesNotExist errors.Code = "DatabaseIDDoesNotExist"
ErrTableIDExists errors.Code = "TableIDExists"
ErrTableKeyExists errors.Code = "TableKeyExists"
ErrTableNameExists errors.Code = "TableNameExists"
@ -17,12 +20,28 @@ const (
ErrFieldExists errors.Code = "FieldExists"
ErrFieldDoesNotExist errors.Code = "FieldDoesNotExist"
ErrInvalidTransaction errors.Code = "InvalidTransaction"
ErrUnimplemented errors.Code = "Unimplemented"
)
// The following are helper functions for constructing coded errors containing
// relevant information about the specific error.
func NewErrDatabaseIDExists(qdbid QualifiedDatabaseID) error {
return errors.New(
ErrDatabaseIDExists,
fmt.Sprintf("database ID '%s' already exists", qdbid),
)
}
func NewErrDatabaseIDDoesNotExist(qdbid QualifiedDatabaseID) error {
return errors.New(
ErrDatabaseIDDoesNotExist,
fmt.Sprintf("database ID '%s' does not exist", qdbid),
)
}
func NewErrTableIDDoesNotExist(qtid QualifiedTableID) error {
return errors.New(
ErrTableIDDoesNotExist,
@ -78,3 +97,10 @@ func NewErrFieldExists(fieldName FieldName) error {
fmt.Sprintf("field '%s' already exists", fieldName),
)
}
func NewErrInvalidTransaction() error {
return errors.New(
ErrInvalidTransaction,
"tx is not a *boltdb.Tx",
)
}

72
dax/mds/client/client.go
View file

@ -54,8 +54,8 @@ func (c *Client) TableByID(ctx context.Context, qtid dax.QualifiedTableID) (*dax
}
// TODO(tlt): collapse TableID into this
func (c *Client) TableByName(ctx context.Context, qual dax.TableQualifier, tname dax.TableName) (*dax.QualifiedTable, error) {
qtid, err := c.TableID(ctx, qual, tname)
func (c *Client) TableByName(ctx context.Context, qdbid dax.QualifiedDatabaseID, tname dax.TableName) (*dax.QualifiedTable, error) {
qtid, err := c.TableID(ctx, qdbid, tname)
if err != nil {
return nil, errors.Wrap(err, "getting table id")
}
@ -93,14 +93,14 @@ func (c *Client) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.Qua
return qtable, nil
}
func (c *Client) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
func (c *Client) TableID(ctx context.Context, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
url := fmt.Sprintf("%s/table-id", c.address.WithScheme(defaultScheme))
dflt := dax.QualifiedTableID{}
req := dax.QualifiedTableID{
TableQualifier: qual,
Name: name,
QualifiedDatabaseID: qdbid,
Name: name,
}
// Encode the request.
@ -130,12 +130,12 @@ func (c *Client) TableID(ctx context.Context, qual dax.TableQualifier, name dax.
return qtid, nil
}
func (c *Client) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
func (c *Client) Tables(ctx context.Context, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
url := fmt.Sprintf("%s/tables", c.address.WithScheme(defaultScheme))
req := mdshttp.TablesRequest{
OrganizationID: qual.OrganizationID,
DatabaseID: qual.DatabaseID,
OrganizationID: qdbid.OrganizationID,
DatabaseID: qdbid.DatabaseID,
TableIDs: ids,
}
@ -166,6 +166,62 @@ func (c *Client) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax
return qtables, nil
}
func (c *Client) CreateDatabase(ctx context.Context, qdb *dax.QualifiedDatabase) error {
url := fmt.Sprintf("%s/create-database", c.address.WithScheme(defaultScheme))
// Encode the request.
postBody, err := json.Marshal(qdb)
if err != nil {
return errors.Wrap(err, "marshalling post request")
}
responseBody := bytes.NewBuffer(postBody)
// Post the request.
resp, err := http.Post(url, "application/json", responseBody)
if err != nil {
return errors.Wrap(err, "posting create database request")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
b, _ := io.ReadAll(resp.Body)
return errors.Errorf("status code: %d: %s", resp.StatusCode, b)
}
return nil
}
func (c *Client) DatabaseByID(ctx context.Context, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
url := fmt.Sprintf("%s/database-by-id", c.address.WithScheme(defaultScheme))
// Encode the request.
postBody, err := json.Marshal(qdbid)
if err != nil {
return nil, errors.Wrap(err, "marshalling post request")
}
responseBody := bytes.NewBuffer(postBody)
// Post the request.
c.logger.Debugf("POST database request: url: %s", url)
resp, err := http.Post(url, "application/json", responseBody)
if err != nil {
return nil, errors.Wrap(err, "posting table request")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
b, _ := io.ReadAll(resp.Body)
return nil, errors.Errorf("status code: %d: %s", resp.StatusCode, b)
}
var qdb *dax.QualifiedDatabase
if err := json.NewDecoder(resp.Body).Decode(&qdb); err != nil {
return nil, errors.Wrap(err, "reading response body")
}
return qdb, nil
}
func (c *Client) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
url := fmt.Sprintf("%s/create-table", c.address.WithScheme(defaultScheme))

78
dax/mds/controller/balancer.go
View file

@ -1,38 +1,42 @@
package controller
import (
"context"
"fmt"
"github.com/featurebasedb/featurebase/v3/dax"
)
type Balancer interface {
AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error)
RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error)
AddJobs(ctx context.Context, job ...fmt.Stringer) ([]dax.WorkerDiff, error)
RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error)
Balance(ctx context.Context) ([]dax.WorkerDiff, error)
CurrentState(ctx context.Context) ([]dax.WorkerInfo, error)
WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error)
WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error)
// AddWorker adds a worker to the global pool of available workers.
AddWorker(tx dax.Transaction, node *dax.Node) ([]dax.WorkerDiff, error)
// WorkersForJobPrefix returns all workers and their job
// assignments which start with `prefix` for all jobs that start
// with `prefix`. If there are free jobs that start with `prefix`
// an error is returned.
//
// The motivating use case is getting all workers for a particular
// table so we can execute a query that will hit every shard in a
// table. If there are jobs representing shards in that table
// which are not assigned to any worker, that means the query
// would return incomplete data, so we want to error.
WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error)
// RemoveWorker removes a worker from the system. If the worker is currently
// assigned to a database and has jobs, it will be removed and its jobs will
// be either transferred to other workers or placed on the free job list.
RemoveWorker(tx dax.Transaction, addr dax.Address) ([]dax.WorkerDiff, error)
// RemoveJobs is for e.g. when dropping a table remove all jobs
// associated with that table without needing to look up in
// advance which shards or partitions are actually present.
RemoveJobs(ctx context.Context, prefix string) ([]dax.WorkerDiff, error)
// AddJobs adds new jobs for the given database.
AddJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error)
// RemoveJobs removes jobs for the given database.
RemoveJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error)
// BalanceDatabase forces a database balance. TODO(tlt): currently this is
// only used in tests, so perhaps we can get rid of it.
BalanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerDiff, error)
// CurrentState returns the workers and jobs currently active for the given
// database.
CurrentState(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error)
// WorkerState returns the jobs currently active for the given worker.
WorkerState(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) (dax.WorkerInfo, error)
// WorkersForJobs returns the workers and jobs currently responsible for the
// given jobs.
WorkersForJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) ([]dax.WorkerInfo, error)
// WorkersForTable returns the workers responsible for any job related to
// the given table.
WorkersForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.WorkerInfo, error)
}
// Ensure type implements interface.
@ -45,34 +49,30 @@ func NewNopBalancer() *NopBalancer {
return &NopBalancer{}
}
func (b *NopBalancer) AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
func (b *NopBalancer) AddWorker(tx dax.Transaction, node *dax.Node) ([]dax.WorkerDiff, error) {
return []dax.WorkerDiff{}, nil
}
func (b *NopBalancer) RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
func (b *NopBalancer) RemoveWorker(tx dax.Transaction, addr dax.Address) ([]dax.WorkerDiff, error) {
return []dax.WorkerDiff{}, nil
}
func (b *NopBalancer) AddJobs(ctx context.Context, job ...fmt.Stringer) ([]dax.WorkerDiff, error) {
func (b *NopBalancer) AddJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
return []dax.WorkerDiff{}, nil
}
func (b *NopBalancer) RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
func (b *NopBalancer) RemoveJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
return []dax.WorkerDiff{}, nil
}
func (b *NopBalancer) Balance(ctx context.Context) ([]dax.WorkerDiff, error) {
func (b *NopBalancer) BalanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerDiff, error) {
return []dax.WorkerDiff{}, nil
}
func (b *NopBalancer) CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) {
func (b *NopBalancer) CurrentState(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error) {
return []dax.WorkerInfo{}, nil
}
func (b *NopBalancer) WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) {
func (b *NopBalancer) WorkerState(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) (dax.WorkerInfo, error) {
return dax.WorkerInfo{}, nil
}
func (b *NopBalancer) WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) {
func (b *NopBalancer) WorkersForJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) ([]dax.WorkerInfo, error) {
return []dax.WorkerInfo{}, nil
}
func (b *NopBalancer) WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) {
func (b *NopBalancer) WorkersForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.WorkerInfo, error) {
return []dax.WorkerInfo{}, nil
}
func (b *NopBalancer) RemoveJobs(ctx context.Context, prefix string) ([]dax.WorkerDiff, error) {
return nil, nil
}

821
dax/mds/controller/balancer/balancer.go Normal file
View file

@ -0,0 +1,821 @@
// Package balancer is an implementation of the controller's Balancer interface.
package balancer
import (
"log"
"math"
"sort"
"strings"
"time"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
"github.com/featurebasedb/featurebase/v3/errors"
"github.com/featurebasedb/featurebase/v3/logger"
)
// Ensure type implements interface.
var _ controller.Balancer = (*Balancer)(nil)
// Balancer is an implementation of the controller.Balancer interface which
// isolates workers and jobs by database. It helps manage the relationships
// between workers and jobs. The logic it uses to balance jobs across workers is
// very simple; it bases everything off the number of workers and number of
// jobs. It does not take anything else (such as job size, worker capabilities,
// etc) into consideration.
type Balancer struct {
// current represents the current state of worker/job assigments.
current WorkerJobService
// freeJobs is the set of jobs which have yet to be assigned to a worker.
// This could be because there are no available workers, or because a worker
// has been removed and the jobs for which it was responsible have yet to be
// reassigned.
freeJobs FreeJobService
freeWorkers FreeWorkerService
schemar schemar.Schemar
logger logger.Logger
}
// New returns a new instance of Balancer.
func New(fjs FreeJobService, wjs WorkerJobService, fws FreeWorkerService, schemar schemar.Schemar, logger logger.Logger) *Balancer {
return &Balancer{
current: wjs,
freeJobs: fjs,
freeWorkers: fws,
schemar: schemar,
logger: logger,
}
}
// AddWorker adds the given Node to the Balancer's available worker pool.
// TODO(tlt): this method takes a Node (as opposed to a Worker) because in the
// future we may want to maintain separate worker pools based on RoleType
// (compute, translate, etc.).
func (b *Balancer) AddWorker(tx dax.Transaction, node *dax.Node) ([]dax.WorkerDiff, error) {
addr := node.Address
b.logger.Debugf("AddWorker(%s)", addr)
diffs := NewInternalDiffs()
// This logic means that a node is used for ALL of the role types specified.
// In other words, specifying roleTypes = {compute, translate}, does not
// mean that the node can be used as either a compute worker or a translate
// worker. It means that it will be used as both.
for _, rt := range node.RoleTypes {
if err := b.addWorker(tx, rt, addr); err != nil {
return nil, errors.Wrapf(err, "adding worker: (%s) %s", rt, addr)
}
}
// Process the freeWorkers.
// TODO(tlt): this is a little heavy-handed. I'm sure we'll need to be more
// intentional about knowing which databases needs workers, as opposed to
// this brute force loop over all databases every time.
if diff, err := b.balance(tx); err != nil {
return nil, errors.Wrapf(err, "balancing new worker: %s", addr)
} else {
diffs.Merge(diff)
}
return diffs.Output(), nil
}
// addWorker adds a worker to the free worker list. From there, it can be used
// by any database which needs a worker.
func (b *Balancer) addWorker(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) error {
// If this worker already exists, don't do anything.
if dbkey := b.current.DatabaseForWorker(tx, addr); dbkey != "" {
return nil
}
if err := b.freeWorkers.AddWorkers(tx, roleType, addr); err != nil {
return errors.Wrap(err, "adding free worker")
}
return nil
}
func (b *Balancer) assignMinWorkers(tx dax.Transaction, roleType dax.RoleType) (InternalDiffs, error) {
// Find out how many free workers we have.
freeWorkers, err := b.freeWorkers.ListWorkers(tx, roleType)
if err != nil {
return nil, errors.Wrap(err, "getting free worker list")
}
freeWorkerCount := len(freeWorkers)
// If there are no free workers, return early.
if freeWorkerCount == 0 {
return InternalDiffs{}, nil
}
// Get all database and their minWorkerCount (Database.Options.WorkersMin).
qdbs, err := b.schemar.Databases(tx, "")
if err != nil {
return nil, errors.Wrap(err, "getting all database")
}
// Create a map[database]int where int is the number of workers required to
// reach that database's minWorkerCount. This map will only contain database
// which need more workers in order to reach their minimum.
m := make(map[dax.QualifiedDatabaseID]int)
for _, qdb := range qdbs {
qdbid := qdb.QualifiedID()
minWorkers := qdb.Options.WorkersMin
if minWorkers == 0 {
continue
}
// If the database doesn't have any jobs, there's no need to go any
// further. In other words, we don't want to assign a worker to a
// database until it has at least one job.
if hasJobs, err := b.databaseHasJobs(tx, roleType, qdbid); err != nil {
return nil, errors.Wrapf(err, "checking has jobs: (%s) %s", roleType, qdbid)
} else if !hasJobs {
continue
}
// Get the number of workers assigned to this database.
workerCount, err := b.current.WorkerCount(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting worker count: (%s) %s", roleType, qdbid)
}
diff := minWorkers - workerCount
// If we have more workers than the min required, or if we have the
// exact number of workers , don't do anything for that database.
if diff <= 0 {
continue
}
m[qdbid] = diff
}
diffs := NewInternalDiffs()
// Create an ordered slice of map keys so that tests are predicatable.
qdbids := make([]dax.QualifiedDatabaseID, 0, len(m))
for qdbid := range m {
qdbids = append(qdbids, qdbid)
}
sort.Sort(dax.QualifiedDatabaseIDs(qdbids))
// For each database, if there are enough free workers to
// satisfy its min, then pop that number of workers from the free list. If
// not, contine to the next database until either reaching the end of the
// database list or until there are no more free workers in the list,
// whichever comes first.
for _, qdbid := range qdbids {
need := m[qdbid]
if freeWorkerCount == 0 {
break
}
if freeWorkerCount >= need {
addrs, err := b.freeWorkers.PopWorkers(tx, roleType, need)
if err != nil {
return nil, errors.Wrapf(err, "popping free worker: (%s)", roleType)
}
if diff, err := b.addDatabaseWorkers(tx, roleType, qdbid, addrs...); err != nil {
return nil, errors.Wrapf(err, "adding database workers: (%s) %s, %v", roleType, qdbid, addrs)
} else {
diffs.Merge(diff)
}
}
}
return diffs, nil
}
// addDatabaseWorkers adds workers from the free worker list to the pool of
// workers for a specific database.
func (b *Balancer) addDatabaseWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addrs ...dax.Address) (InternalDiffs, error) {
for _, addr := range addrs {
if err := b.current.CreateWorker(tx, roleType, qdbid, addr); err != nil {
return nil, errors.Wrap(err, "creating worker")
}
}
// Process the freeJobs.
return b.processFreeJobs(tx, roleType, qdbid)
}
// databaseHasJobs returns true if the database has at least one job.
func (b *Balancer) databaseHasJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (bool, error) {
// Free jobs.
if freeJobs, err := b.freeJobs.ListJobs(tx, roleType, qdbid); err != nil {
return false, errors.Wrapf(err, "getting free jobs: (%s) %s", roleType, qdbid)
} else if len(freeJobs) > 0 {
return true, nil
}
// Assigned jobs.
if wis, err := b.current.WorkersJobs(tx, roleType, qdbid); err != nil {
return false, errors.Wrapf(err, "getting free jobs: (%s) %s", roleType, qdbid)
} else {
for _, wi := range wis {
if len(wi.Jobs) > 0 {
return true, nil
}
}
}
return false, nil
}
func (b *Balancer) RemoveWorker(tx dax.Transaction, addr dax.Address) ([]dax.WorkerDiff, error) {
diffs := NewInternalDiffs()
// See if the worker is assigned to a database.
dbkey := b.current.DatabaseForWorker(tx, addr)
if dbkey == "" {
return diffs.Output(), nil
}
qdbid := dbkey.QualifiedDatabaseID()
for _, rt := range []dax.RoleType{dax.RoleTypeCompute, dax.RoleTypeTranslate} {
// Remove the worker form the free worker list (if it's there).
if err := b.freeWorkers.RemoveWorker(tx, rt, addr); err != nil {
return nil, errors.Wrapf(err, "removing worker from free list: (%s) %s", rt, addr)
}
if diff, err := b.removeDatabaseWorker(tx, rt, qdbid, addr); err != nil {
return nil, errors.Wrapf(err, "removing worker: (%s) %s", rt, addr)
} else {
diffs.Merge(diff)
}
}
// Balance the affected database.
if diff, err := b.balanceDatabase(tx, qdbid); err != nil {
return nil, errors.Wrapf(err, "balancing database: %s", qdbid)
} else {
diffs.Merge(diff)
}
return diffs.Output(), nil
}
func (b *Balancer) removeDatabaseWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) (InternalDiffs, error) {
jobs, err := b.current.ListJobs(tx, roleType, qdbid, addr)
if err != nil {
return nil, errors.Wrap(err, "listing jobs")
}
// Before removing the worker, mark its jobs as free.
if err := b.freeJobs.MergeJobs(tx, roleType, qdbid, jobs); err != nil {
return nil, errors.Wrap(err, "merging free jobs")
}
// Remove the worker.
if err := b.current.DeleteWorker(tx, roleType, qdbid, addr); err != nil {
return nil, errors.Wrap(err, "deleting worker")
}
// Even though this may not be useful to the caller (for example, in the
// case where the worker has died and no longer exists), return the diffs
// which represent the removal of jobs from the worker.
diff := NewInternalDiffs()
for _, job := range jobs {
diff.Removed(addr, job)
}
return diff, nil
}
func (b *Balancer) AddJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
start := time.Now()
defer func() {
log.Printf("ELAPSED: Balancer.AddJob: %v", time.Since(start))
}()
switch len(jobs) {
case 0:
// No jobs so return early.
b.logger.Debugf("%s: AddJobs (no jobs provided)", roleType)
return []dax.WorkerDiff{}, nil
case 1:
b.logger.Debugf("%s: AddJobs (%s)", roleType, jobs[0])
default:
b.logger.Debugf("%s: AddJobs (%d)", roleType, len(jobs))
}
// TODO(tlt): we don't currently use "table" in this method; even though we
// pass a table, we're still encoding the tableKey in the job. In theory, we
// could exclude tableKey from the job coming into this method, and add it
// here.
qdbid := qtid.QualifiedDatabaseID
diff, err := b.addJobs(tx, roleType, qdbid, jobs...)
if err != nil {
return nil, errors.Wrap(err, "adding job")
}
return diff.Output(), nil
}
func (b *Balancer) addJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) (InternalDiffs, error) {
diffs := NewInternalDiffs()
if len(jobs) == 0 {
return diffs, nil
}
if cnt, err := b.current.WorkerCount(tx, roleType, qdbid); err != nil {
return nil, errors.Wrap(err, "getting worker count")
} else if cnt == 0 {
if err := b.freeJobs.CreateJobs(tx, roleType, qdbid, jobs...); err != nil {
return nil, errors.Wrap(err, "creating free job")
}
// Since we've just added free jobs to the database, try to add a worker
// for the database. This case would happen because when a database is
// first created, it is not assigned any workers. A database is not
// assigned workers until it has at least one job (which this database
// now has).
if diff, err := b.balanceDatabaseForRole(tx, roleType, qdbid); err != nil {
return nil, errors.Wrapf(err, "assigning min workers: (%s)", roleType)
} else {
diffs.Merge(diff)
}
// Now check, again, to see if the database has a worker.
if cnt2, err := b.current.WorkerCount(tx, roleType, qdbid); err != nil {
return nil, errors.Wrap(err, "getting worker count, again")
} else if cnt2 == 0 {
// TODO: we might want to inform the user that a job is in the free list
// because there are no workers.
return InternalDiffs{}, nil
}
}
diff, err := b.addDatabaseJobs(tx, roleType, qdbid, jobs...)
if err != nil {
return nil, errors.Wrapf(err, "adding database jobs: (%s) %s", roleType, qdbid)
}
diffs.Merge(diff)
return diffs, nil
}
// addDatabaseJobs adds the job for the provided database.
func (b *Balancer) addDatabaseJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) (InternalDiffs, error) {
workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting workers jobs: %s", roleType)
}
jset := dax.NewSet[dax.Job]()
for _, workerInfo := range workerJobs {
jset.Merge(dax.NewSet(workerInfo.Jobs...))
}
addrs := make(dax.Addresses, 0, len(workerJobs))
jobCounts := make(map[dax.Address]int, 0)
for _, v := range workerJobs {
addrs = append(addrs, v.Address)
jobCounts[v.Address] = len(v.Jobs)
}
diffs := NewInternalDiffs()
jobsToCreate := make(map[dax.Address][]dax.Job)
for _, job := range jobs {
// Skip any job that already exists.
if jset.Contains(job) {
continue
}
// Find the worker with the fewest number of jobs and assign it this job.
var lowCount int = math.MaxInt
var lowWorker dax.Address
// We loop over addrs here instead of jobCounts because jobCounts is a
// map and it can return results in an unexpected order, which is a
// problem for testing.
for _, addr := range addrs {
jobCount := jobCounts[addr]
if jobCount < lowCount {
lowCount = jobCount
lowWorker = addr
}
}
jobsToCreate[lowWorker] = append(jobsToCreate[lowWorker], job)
jobCounts[lowWorker]++
}
for worker, jobs := range jobsToCreate {
if err := b.current.CreateJobs(tx, roleType, qdbid, worker, jobs...); err != nil {
return nil, errors.Wrap(err, "creating job")
}
for _, job := range jobs {
diffs.Added(worker, job)
}
}
return diffs, nil
}
func (b *Balancer) RemoveJobs(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID, jobs ...dax.Job) ([]dax.WorkerDiff, error) {
qdbid := qtid.QualifiedDatabaseID
// If no jobs are provided, remove all jobs for table.
if len(jobs) == 0 {
diffs, err := b.removeJobsForTable(tx, roleType, qtid)
if err != nil {
return nil, errors.Wrapf(err, "removing jobs for table: (%s) %s", roleType, qtid)
}
return diffs.Output(), nil
}
diffs := NewInternalDiffs()
for _, job := range jobs {
if diff, err := b.removeJob(tx, roleType, qdbid, job); err != nil {
return nil, errors.Wrapf(err, "removing job: (%s) %s, %s", roleType, qdbid, job)
} else {
diffs.Merge(diff)
}
}
return diffs.Output(), nil
}
func (b *Balancer) removeJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) (InternalDiffs, error) {
idiffs, err := b.current.DeleteJobsForTable(tx, roleType, qtid)
if err != nil {
return nil, errors.Wrapf(err, "deleting jobs for table: (%s) %s", roleType, qtid)
}
if err := b.freeJobs.DeleteJobsForTable(tx, roleType, qtid); err != nil {
return nil, errors.Wrapf(err, "deleting free jobs for table: (%s) %s", roleType, qtid)
}
return idiffs, nil
}
// Balance calls balanceDatabase on every database in the schemar.
func (b *Balancer) Balance(tx dax.Transaction) ([]dax.WorkerDiff, error) {
diffs, err := b.balance(tx)
if err != nil {
return nil, errors.Wrapf(err, "balancing all")
}
return diffs.Output(), nil
}
func (b *Balancer) balance(tx dax.Transaction) (InternalDiffs, error) {
qdbs, err := b.schemar.Databases(tx, "")
if err != nil {
return nil, errors.Wrapf(err, "getting all databases")
}
diffs := NewInternalDiffs()
for _, qdb := range qdbs {
qdbid := qdb.QualifiedID()
if diff, err := b.balanceDatabase(tx, qdbid); err != nil {
return nil, errors.Wrapf(err, "balancing database: %s", qdbid)
} else {
diffs.Merge(diff)
}
}
return diffs, nil
}
func (b *Balancer) BalanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerDiff, error) {
diffs, err := b.balanceDatabase(tx, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "balancing database: %s", qdbid)
}
return diffs.Output(), nil
}
func (b *Balancer) balanceDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) (InternalDiffs, error) {
diffs := NewInternalDiffs()
for _, role := range []dax.RoleType{dax.RoleTypeCompute, dax.RoleTypeTranslate} {
diff, err := b.balanceDatabaseForRole(tx, role, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting worker count: (%s) %s", role, qdbid)
}
diffs.Merge(diff)
}
return diffs, nil
}
func (b *Balancer) balanceDatabaseForRole(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (InternalDiffs, error) {
diffs := NewInternalDiffs()
// Before balancing, make sure the database has its minimum number of
// workers satisfied.
// TODO(tlt): make assignMinWorkers database specific.
if diff, err := b.assignMinWorkers(tx, roleType); err != nil {
return nil, errors.Wrapf(err, "assigning min workers: (%s) %s", roleType, qdbid)
} else {
diffs.Merge(diff)
}
// If there are no workers, we can't properly balance.
if cnt, err := b.current.WorkerCount(tx, roleType, qdbid); err != nil {
return nil, errors.Wrapf(err, "getting worker count: (%s) %s", roleType, qdbid)
} else if cnt == 0 {
return InternalDiffs{}, nil
}
// Process the freeJobs.
if diff, err := b.processFreeJobs(tx, roleType, qdbid); err != nil {
return nil, errors.Wrapf(err, "processing free jobs: (%s) %s", roleType, qdbid)
} else {
diffs.Merge(diff)
}
// Balance the jobs among workers.
diff, err := b.balanceDatabaseJobs(tx, roleType, qdbid, diffs)
if err != nil {
return nil, errors.Wrap(err, "balancing jobs")
}
return diff, nil
}
func (b *Balancer) CurrentState(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error) {
return b.current.WorkersJobs(tx, roleType, qdbid)
}
func (b *Balancer) WorkerState(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) (dax.WorkerInfo, error) {
info := dax.WorkerInfo{
Address: addr,
}
dbkey := b.current.DatabaseForWorker(tx, addr)
if dbkey == "" {
return info, nil
}
qdbid := dbkey.QualifiedDatabaseID()
jobs, err := b.current.ListJobs(tx, roleType, qdbid, addr)
if err != nil {
return dax.WorkerInfo{}, errors.Wrapf(err, "listing jobs: (%s) %s, %s", roleType, qdbid, addr)
}
info.Jobs = jobs
return info, nil
}
func (b *Balancer) WorkersForJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) ([]dax.WorkerInfo, error) {
out := make(map[dax.Address]dax.Set[dax.Job])
workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting worker jobs: (%s) %s", roleType, qdbid)
}
for _, workerInfo := range workerJobs {
jset := dax.NewSet(workerInfo.Jobs...)
matches := dax.NewSet[dax.Job]()
for _, job := range jobs {
if jset.Contains(job) {
matches.Add(job)
}
}
if len(matches) > 0 {
out[workerInfo.Address] = matches
}
}
workers := make([]dax.WorkerInfo, 0, len(out))
for addr, jset := range out {
workers = append(workers, dax.WorkerInfo{
Address: addr,
Jobs: jset.Sorted(),
})
}
sort.Sort(dax.WorkerInfos(workers))
return workers, nil
}
func (b *Balancer) WorkersForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.WorkerInfo, error) {
out := make(map[dax.Address]dax.Set[dax.Job])
qdbid := qtid.QualifiedDatabaseID
prefix := string(qtid.Key())
workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting worker jobs: (%s) %s", roleType, qdbid)
}
for _, workerInfo := range workerJobs {
matches := dax.NewSet[dax.Job]()
for _, job := range workerInfo.Jobs {
if strings.HasPrefix(string(job), prefix) {
matches.Add(job)
}
}
if len(matches) > 0 {
out[workerInfo.Address] = matches
}
}
workers := make([]dax.WorkerInfo, 0, len(out))
for addr, jset := range out {
workers = append(workers, dax.WorkerInfo{
Address: addr,
Jobs: jset.Sorted(),
})
}
sort.Sort(dax.WorkerInfos(workers))
return workers, nil
}
func (b *Balancer) removeJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) (InternalDiffs, error) {
if addr, ok, err := b.workerForJob(tx, roleType, qdbid, job); err != nil {
return nil, errors.Wrapf(err, "getting worker for job: %s", job)
} else if ok {
if err := b.current.DeleteJob(tx, roleType, qdbid, addr, job); err != nil {
return nil, errors.Wrapf(err, "deleting job: (%s) %s, %s, %s", roleType, qdbid, addr, job)
}
diffs := NewInternalDiffs()
diffs.Removed(addr, job)
return diffs, nil
}
// Just in case the job is in the free list (and wasn't assigned to a
// worker), remove it; there's no need to provide a diff. There should never
// be a case where the same job is both in the free list and assigned to a
// worker.
if err := b.freeJobs.DeleteJob(tx, roleType, qdbid, job); err != nil {
return nil, errors.Wrapf(err, "deleting free job: (%s) %s, %s", roleType, qdbid, job)
}
return InternalDiffs{}, nil
}
// balanceDatabaseJobs moves jobs among workers with the goal of having an equal
// number of jobs per worker. This method takes an `internalDiffs` as input for
// cases where some action has preceeded this call which also resulted in
// `internalDiffs`. Instead of having this method take a value, we could rely on
// the internalDiffs.merge() method, but we would need to modify that method to
// be smarter about the order in which it applies the add/remove operations.
// Until that's in place, we'll pass in a value here.
func (b *Balancer) balanceDatabaseJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, diffs InternalDiffs) (InternalDiffs, error) {
numWorkers, err := b.current.WorkerCount(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting worker count: (%s) %s", roleType, qdbid)
}
numJobs := 0
if addrs, err := b.current.ListWorkers(tx, roleType, qdbid); err != nil {
return nil, errors.Wrapf(err, "listing workers: (%s) %s", roleType, qdbid)
} else {
for _, addr := range addrs {
jobCounts, err := b.current.JobCounts(tx, roleType, qdbid, addr)
if err != nil {
return nil, errors.Wrapf(err, "getting job count: (%s) %s, %s", roleType, qdbid, addr)
}
numJobs += jobCounts[addr]
}
}
minJobsPerWorker := numJobs / numWorkers
numWorkersAboveMin := numJobs % numWorkers
// workerInfos is used now in order to guarantee a sort order.
workerInfos, err := b.CurrentState(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting current state: (%s) %s", roleType, qdbid)
}
// Loop through each worker, and if the number of jobs for the worker
// exceeds the target, then remove the job and add it back (which is
// effectively how we rebalance a job).
for i, workerInfo := range workerInfos {
numTargetJobs := minJobsPerWorker
if i < numWorkersAboveMin {
numTargetJobs += 1
}
jobCounts, err := b.current.JobCounts(tx, roleType, qdbid, workerInfo.Address)
if err != nil {
return nil, errors.Wrapf(err, "getting job count: (%s) %s, %s", roleType, qdbid, workerInfo.Address)
}
numCurrentJobs := jobCounts[workerInfo.Address]
// If we don't need to remove jobs from this worker, then just continue
// on to the next worker.
if numCurrentJobs <= numTargetJobs {
continue
}
sortedJobs, err := b.current.ListJobs(tx, roleType, qdbid, workerInfo.Address)
if err != nil {
return nil, errors.Wrapf(err, "listing jobs: (%s) %s, %s", roleType, qdbid, workerInfo.Address)
}
// Remove the extra jobs from the end of the list, and add them back
// again (which should place them on a worker with fewer jobs).
for i := numCurrentJobs - 1; i >= numTargetJobs; i-- {
if rj, err := b.removeJob(tx, roleType, qdbid, sortedJobs[i]); err != nil {
return nil, errors.Wrapf(err, "removing job: %s", sortedJobs[i])
} else {
diffs.Merge(rj)
}
if aj, err := b.addJobs(tx, roleType, qdbid, sortedJobs[i]); err != nil {
return nil, errors.Wrapf(err, "adding job: %s", sortedJobs[i])
} else {
diffs.Merge(aj)
}
}
}
return diffs, nil
}
// processFreeJobs assigns all jobs in the free list to a worker.
func (b *Balancer) processFreeJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (InternalDiffs, error) {
diffs := NewInternalDiffs()
jobs, err := b.freeJobs.ListJobs(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "listing free jobs: %s", roleType)
}
for _, job := range jobs {
if aj, err := b.addDatabaseJobs(tx, roleType, qdbid, job); err != nil {
return nil, errors.Wrapf(err, "adding job: %s", job)
} else {
diffs.Merge(aj)
}
if err := b.freeJobs.DeleteJob(tx, roleType, qdbid, job); err != nil {
return nil, errors.Wrapf(err, "deleting free job: %s", job)
}
}
return diffs, nil
}
// workerForJob returns the worker currently assigned to the given job.
func (b *Balancer) workerForJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) (dax.Address, bool, error) {
workerJobs, err := b.current.WorkersJobs(tx, roleType, qdbid)
if err != nil {
return "", false, errors.Wrapf(err, "getting workers jobs: (%s) %s", roleType, qdbid)
}
for _, workerInfo := range workerJobs {
jset := dax.NewSet(workerInfo.Jobs...)
if jset.Contains(job) {
return workerInfo.Address, true, nil
}
}
return "", false, nil
}
type WorkerJobService interface {
WorkersJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error)
WorkerCount(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (int, error)
ListWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Addresses, error)
CreateWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error
DeleteWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error
CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, job ...dax.Job) error
DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, job dax.Job) error
DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) (InternalDiffs, error)
JobCounts(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr ...dax.Address) (map[dax.Address]int, error)
ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) (dax.Jobs, error)
DatabaseForWorker(tx dax.Transaction, addr dax.Address) dax.DatabaseKey
}
type FreeJobService interface {
CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job ...dax.Job) error
DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) error
DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) error
ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Jobs, error)
MergeJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs dax.Jobs) error
}
type FreeWorkerService interface {
AddWorkers(tx dax.Transaction, roleType dax.RoleType, addrs ...dax.Address) error
RemoveWorker(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) error
PopWorkers(tx dax.Transaction, roleType dax.RoleType, num int) ([]dax.Address, error)
ListWorkers(tx dax.Transaction, roleType dax.RoleType) (dax.Addresses, error)
}

1104
dax/mds/controller/balancer/balancer_test.go Normal file
View file

File diff suppressed because it is too large Load diff

813
dax/mds/controller/balancer/boltdb/balancer.go Normal file
View file

@ -0,0 +1,813 @@
// Package boltdb contains the boltdb implementation of the Balancer interface.
package boltdb
import (
"bytes"
"encoding/json"
"fmt"
"strings"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/dax/boltdb"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller/balancer"
"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
"github.com/featurebasedb/featurebase/v3/errors"
"github.com/featurebasedb/featurebase/v3/logger"
)
var (
bucketBalancer = boltdb.Bucket("balancer")
)
// BalancerBuckets defines the buckets used by this package. It can be
// called during setup to create the buckets ahead of time.
var BalancerBuckets []boltdb.Bucket = []boltdb.Bucket{
bucketBalancer,
}
// NewBalancer returns a new instance of controller.Balancer.
func NewBalancer(db *boltdb.DB, schemar schemar.Schemar, logger logger.Logger) controller.Balancer {
fjs := newFreeJobService(db)
wjs := newWorkerJobService(db, logger)
fws := newFreeWorkerService(db)
return balancer.New(fjs, wjs, fws, schemar, logger)
}
// Ensure type implements interface.
var _ balancer.WorkerJobService = (*workerJobService)(nil)
type workerJobService struct {
db *boltdb.DB
logger logger.Logger
}
func newWorkerJobService(db *boltdb.DB, logger logger.Logger) *workerJobService {
return &workerJobService{
db: db,
logger: logger,
}
}
func (w *workerJobService) WorkersJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) ([]dax.WorkerInfo, error) {
workerInfos, err := w.getWorkerInfos(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrapf(err, "getting worker infos: %s", roleType)
}
return workerInfos, nil
}
func (w *workerJobService) WorkerCount(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (int, error) {
workers, err := w.getWorkers(tx, roleType, qdbid)
if err != nil {
return 0, errors.Wrapf(err, "getting workers: %s", roleType)
}
return len(workers), nil
}
func (w *workerJobService) ListWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Addresses, error) {
return w.getWorkers(tx, roleType, qdbid)
}
func (w *workerJobService) getWorkers(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Addresses, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
c := txx.Bucket(bucketBalancer).Cursor()
// Deserialize rows into Worker objects.
addrs := make(dax.Addresses, 0)
prefix := []byte(fmt.Sprintf(prefixFmtWorkersDB, roleType, qdbid.Key()))
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
if v == nil {
w.logger.Printf("nil value for key: %s", k)
continue
}
addr, err := keyWorker(k)
if err != nil {
return nil, errors.Wrapf(err, "getting worker from key: %s", k)
}
addrs = append(addrs, addr)
}
return addrs, nil
}
func (w *workerJobService) getWorkerInfos(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.WorkerInfos, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
c := txx.Bucket(bucketBalancer).Cursor()
// Deserialize rows into WorkerInfo objects.
workerInfos := make(dax.WorkerInfos, 0)
prefix := []byte(fmt.Sprintf(prefixFmtWorkersDB, roleType, qdbid.Key()))
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
addr, err := keyWorker(k)
if err != nil {
return nil, errors.Wrapf(err, "getting worker from key: %s", k)
}
jobs := dax.NewSet[dax.Job]()
if v != nil {
jobs, err = decodeJobSet(v)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
workerInfo := dax.WorkerInfo{
Address: addr,
Jobs: jobs.Sorted(),
}
workerInfos = append(workerInfos, workerInfo)
}
return workerInfos, nil
}
func (w *workerJobService) CreateWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
// If this worker already exists, don't do anything.
wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
if wrkr != nil {
return nil
}
val := []byte("[]")
if err := bkt.Put(workerDBKey(roleType, qdbid, addr), val); err != nil {
return errors.Wrapf(err, "putting db worker: %s, %s", qdbid, addr)
}
if err := bkt.Put(workerAssignedKey(addr), []byte(qdbid.Key())); err != nil {
return errors.Wrapf(err, "putting assigned worker: %s, %s", qdbid, addr)
}
return nil
}
func (w *workerJobService) DeleteWorker(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
if err := bkt.Delete(workerDBKey(roleType, qdbid, addr)); err != nil {
return errors.Wrapf(err, "deleting node key: %s", workerDBKey(roleType, qdbid, addr))
}
if err := bkt.Delete(workerAssignedKey(addr)); err != nil {
return errors.Wrapf(err, "deleting assigned worker: %s", workerAssignedKey(addr))
}
return nil
}
func (w *workerJobService) CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, jobs ...dax.Job) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
jobset := dax.NewSet[dax.Job]()
var err error
// get worker
wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
if wrkr != nil {
jobset, err = decodeJobSet(wrkr)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
}
for _, job := range jobs {
jobset.Add(job)
}
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(workerDBKey(roleType, qdbid, addr), val); err != nil {
return errors.Wrap(err, "putting worker")
}
return nil
}
func (w *workerJobService) DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address, job dax.Job) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
// get worker
wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
if wrkr == nil {
return nil
}
jobset, err := decodeJobSet(wrkr)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
if !jobset.Contains(job) {
return nil
}
jobset.Remove(job)
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(workerDBKey(roleType, qdbid, addr), val); err != nil {
return errors.Wrap(err, "putting worker")
}
return nil
}
func (w *workerJobService) DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) (balancer.InternalDiffs, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
qdbid := qtid.QualifiedDatabaseID
prefix := string(qtid.Key())
workers, err := w.getWorkers(tx, roleType, qdbid)
if err != nil {
return nil, errors.Wrap(err, "getting workers")
}
idiffs := balancer.NewInternalDiffs()
for _, worker := range workers {
// get worker
wrkr := bkt.Get(workerDBKey(roleType, qdbid, worker))
if wrkr == nil {
panic("didn't find worker that should... definitely exist")
}
jobset, err := decodeJobSet(wrkr)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
jobs := jobset.RemoveByPrefix(prefix)
for _, job := range jobs {
idiffs.Removed(worker, job)
}
val, err := encodeJobSet(jobset)
if err != nil {
return nil, errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(workerDBKey(roleType, qdbid, worker), val); err != nil {
return nil, errors.Wrap(err, "putting worker")
}
}
return idiffs, nil
}
func (w *workerJobService) ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) (dax.Jobs, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
jobset := dax.NewSet[dax.Job]()
var err error
// get worker
wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
if wrkr != nil {
jobset, err = decodeJobSet(wrkr)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
return jobset.Sorted(), nil
}
func (w *workerJobService) JobCounts(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addrs ...dax.Address) (map[dax.Address]int, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
m := make(map[dax.Address]int)
for _, addr := range addrs {
jobset := dax.NewSet[dax.Job]()
var err error
// get worker
wrkr := bkt.Get(workerDBKey(roleType, qdbid, addr))
if wrkr != nil {
jobset, err = decodeJobSet(wrkr)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
m[addr] = len(jobset)
}
return m, nil
}
func (w *workerJobService) DatabaseForWorker(tx dax.Transaction, addr dax.Address) dax.DatabaseKey {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return "" // TODO(tlt): return error here?
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return ""
}
wrkr := bkt.Get(workerAssignedKey(addr))
return dax.DatabaseKey(wrkr)
}
// encodeJobSet encode the jobSet into a JSON array of strings.
func encodeJobSet(jobSet dax.Set[dax.Job]) ([]byte, error) {
arr := jobSet.Sorted()
b, err := json.Marshal(arr)
if err != nil {
return nil, errors.Wrap(err, "marshalling json")
}
return b, nil
}
// decodeJobSet decode the string (a JSON array of strings) into jobSet.
func decodeJobSet(v []byte) (dax.Set[dax.Job], error) {
var arr []string
err := json.Unmarshal(v, &arr)
if err != nil {
return nil, errors.Wrap(err, "unmarshalling json")
}
js := dax.NewSet[dax.Job]()
for _, s := range arr {
js.Add(dax.Job(s))
}
return js, nil
}
// encodeWorkerSet encode the workerSet into a JSON array of strings.
func encodeWorkerSet(workerSet dax.Set[dax.Address]) ([]byte, error) {
arr := workerSet.Sorted()
b, err := json.Marshal(arr)
if err != nil {
return nil, errors.Wrap(err, "marshalling json")
}
return b, nil
}
// decodeWorkerSet decode the string (a JSON array of strings) into workerSet.
func decodeWorkerSet(v []byte) (dax.Set[dax.Address], error) {
var arr []string
err := json.Unmarshal(v, &arr)
if err != nil {
return nil, errors.Wrap(err, "unmarshalling json")
}
ws := dax.NewSet[dax.Address]()
for _, s := range arr {
ws.Add(dax.Address(s))
}
return ws, nil
}
// Ensure type implements interface.
var _ balancer.FreeJobService = (*freeJobService)(nil)
type freeJobService struct {
db *boltdb.DB
}
func newFreeJobService(db *boltdb.DB) *freeJobService {
return &freeJobService{
db: db,
}
}
func (f *freeJobService) CreateJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs ...dax.Job) error {
return f.MergeJobs(tx, roleType, qdbid, jobs)
}
func (f *freeJobService) DeleteJob(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, job dax.Job) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
// get free jobs
fjs := bkt.Get(freeJobKey(roleType, qdbid))
if fjs == nil {
return nil
}
jobset, err := decodeJobSet(fjs)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
if !jobset.Contains(job) {
return nil
}
jobset.Remove(job)
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(freeJobKey(roleType, qdbid), val); err != nil {
return errors.Wrap(err, "putting free job")
}
return nil
}
func (f *freeJobService) DeleteJobsForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
qdbid := qtid.QualifiedDatabaseID
prefix := string(qtid.Key())
// get free jobs
fjs := bkt.Get(freeJobKey(roleType, qdbid))
if fjs == nil {
return nil
}
jobset, err := decodeJobSet(fjs)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
jobset.RemoveByPrefix(prefix)
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(freeJobKey(roleType, qdbid), val); err != nil {
return errors.Wrap(err, "putting free job")
}
return nil
}
func (f *freeJobService) ListJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) (dax.Jobs, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
jobset := dax.NewSet[dax.Job]()
var err error
// get free jobs
fjs := bkt.Get(freeJobKey(roleType, qdbid))
if fjs != nil {
jobset, err = decodeJobSet(fjs)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
return jobset.Sorted(), nil
}
func (f *freeJobService) MergeJobs(tx dax.Transaction, roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, jobs dax.Jobs) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
jobset := dax.NewSet[dax.Job]()
var err error
// get free jobs
fjs := bkt.Get(freeJobKey(roleType, qdbid))
if fjs != nil {
jobset, err = decodeJobSet(fjs)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
}
for _, j := range jobs {
jobset.Add(j)
}
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(freeJobKey(roleType, qdbid), val); err != nil {
return errors.Wrap(err, "putting free job")
}
return nil
}
//////////////////////////////////////////////////////
// Ensure type implements interface.
var _ balancer.FreeWorkerService = (*freeWorkerService)(nil)
type freeWorkerService struct {
db *boltdb.DB
}
func newFreeWorkerService(db *boltdb.DB) *freeWorkerService {
return &freeWorkerService{
db: db,
}
}
func (f *freeWorkerService) AddWorkers(tx dax.Transaction, roleType dax.RoleType, addres ...dax.Address) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
workerset := dax.NewSet[dax.Address]()
var err error
// get free workers
fws := bkt.Get(freeWorkerKey(roleType))
if fws != nil {
workerset, err = decodeWorkerSet(fws)
if err != nil {
return errors.Wrap(err, "decoding worker set")
}
}
for _, w := range addres {
workerset.Add(w)
}
val, err := encodeWorkerSet(workerset)
if err != nil {
return errors.Wrap(err, "encoding worker set")
}
if err := bkt.Put(freeWorkerKey(roleType), val); err != nil {
return errors.Wrap(err, "putting free worker")
}
return nil
}
func (f *freeWorkerService) RemoveWorker(tx dax.Transaction, roleType dax.RoleType, addr dax.Address) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
workers, err := f.ListWorkers(tx, roleType)
if err != nil {
return errors.Wrap(err, "listing free workers")
}
// Create a workerset containing the free workers which remain after
// removing num workers.
workerset := dax.NewSet[dax.Address]()
for _, w := range workers {
workerset.Add(w)
}
if !workerset.Contains(addr) {
return nil
}
// Remove the worker.
workerset.Remove(addr)
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
val, err := encodeWorkerSet(workerset)
if err != nil {
return errors.Wrap(err, "encoding worker set")
}
if err := bkt.Put(freeWorkerKey(roleType), val); err != nil {
return errors.Wrap(err, "putting free worker")
}
return nil
}
func (f *freeWorkerService) PopWorkers(tx dax.Transaction, roleType dax.RoleType, num int) ([]dax.Address, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
workers, err := f.ListWorkers(tx, roleType)
if err != nil {
return nil, errors.Wrap(err, "listing free workers")
}
if len(workers) < num {
return nil, errors.Errorf("not enough free workers to pop: wanted %d, have: %d", num, len(workers))
}
// Get num workers from the list.
workersToAssign := workers[0:num]
// Create a workerset containing the free workers which remain after
// removing num workers.
workerset := dax.NewSet[dax.Address]()
for _, worker := range workers[num:] {
workerset.Add(worker)
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
val, err := encodeWorkerSet(workerset)
if err != nil {
return nil, errors.Wrap(err, "encoding worker set")
}
if err := bkt.Put(freeWorkerKey(roleType), val); err != nil {
return nil, errors.Wrap(err, "putting free worker")
}
return workersToAssign, nil
}
func (f *freeWorkerService) ListWorkers(tx dax.Transaction, roleType dax.RoleType) (dax.Addresses, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
bkt := txx.Bucket(bucketBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketBalancer)
}
workerset := dax.NewSet[dax.Address]()
var err error
// get free workers
fws := bkt.Get(freeWorkerKey(roleType))
if fws != nil {
workerset, err = decodeWorkerSet(fws)
if err != nil {
return nil, errors.Wrap(err, "decoding worker set")
}
}
return workerset.Sorted(), nil
}
//////////////////////////////////////////////////////
const (
prefixFmtWorkersDB = "workers/role/%s/db/%s/" // %s - role, dbKey
prefixFmtWorkersAssigned = "workers/assigned/"
prefixFmtFreeJobs = "freejobs/role/%s/db/%s" // %s - role, dbKey
prefixFmtFreeWorkers = "freeworkers/role/%s" // %s - role
)
// workerDBKey returns a key based on worker.
//
// Format: workers/role/[role]/db/[dbKey]/[worker] = [job1, job2, ...]
func workerDBKey(roleType dax.RoleType, qdbid dax.QualifiedDatabaseID, addr dax.Address) []byte {
key := fmt.Sprintf(prefixFmtWorkersDB+"%s", roleType, qdbid.Key(), addr)
return []byte(key)
}
// workerAssignedKey returns a key based on worker.
//
// Format: workers/assigned/[worker] = dbKey
func workerAssignedKey(addr dax.Address) []byte {
key := fmt.Sprintf(prefixFmtWorkersAssigned+"%s", addr)
return []byte(key)
}
// keyWorker gets the worker out of the key.
func keyWorker(key []byte) (dax.Address, error) {
parts := strings.SplitN(string(key), "/", 6)
if len(parts) != 6 {
return "", errors.New(errors.ErrUncoded, "worker key format expected: `workers/role/[role]/db/[db]/worker`")
}
return dax.Address(parts[5]), nil
}
// freeJobKey returns a key for all freeJobs.
//
// Format: freejobs/role/[role]/db/[dbKey] = [job1, job2, ...]
func freeJobKey(roleType dax.RoleType, qdbid dax.QualifiedDatabaseID) []byte {
key := fmt.Sprintf(prefixFmtFreeJobs, roleType, qdbid.Key())
return []byte(key)
}
// freeWorkerKey returns a key for all freeWorkers.
//
// Format: freeworkers/role/[role] = [worker1, worker2, ...]
func freeWorkerKey(roleType dax.RoleType) []byte {
key := fmt.Sprintf(prefixFmtFreeWorkers, roleType)
return []byte(key)
}

26
dax/mds/controller/naive/types.go → dax/mds/controller/balancer/types.go
View file

@ -1,4 +1,4 @@
package naive
package balancer
import (
"sort"
@ -20,36 +20,36 @@ func newJobSetDiffs() jobSetDiffs {
}
}
type InternalDiffs map[dax.Worker]jobSetDiffs
type InternalDiffs map[dax.Address]jobSetDiffs
func NewInternalDiffs() InternalDiffs {
return make(InternalDiffs)
}
func (d InternalDiffs) Added(worker dax.Worker, job dax.Job) {
if _, ok := d[worker]; !ok {
d[worker] = newJobSetDiffs()
func (d InternalDiffs) Added(address dax.Address, job dax.Job) {
if _, ok := d[address]; !ok {
d[address] = newJobSetDiffs()
}
// Before adding the job, make sure we haven't indicated that it has been
// removed prior to this. If it has, we need to invalidate that "remove"
// instruction.
d[worker].removed.Remove(job)
d[address].removed.Remove(job)
d[worker].added.Add(job)
d[address].added.Add(job)
}
func (d InternalDiffs) Removed(worker dax.Worker, job dax.Job) {
if _, ok := d[worker]; !ok {
d[worker] = newJobSetDiffs()
func (d InternalDiffs) Removed(address dax.Address, job dax.Job) {
if _, ok := d[address]; !ok {
d[address] = newJobSetDiffs()
}
// Before removing the job, make sure we haven't indicated that it has been
// added prior to this. If it has, we need to invalidate that "add"
// instruction.
d[worker].added.Remove(job)
d[address].added.Remove(job)
d[worker].removed.Add(job)
d[address].removed.Add(job)
}
func (d InternalDiffs) Merge(d2 InternalDiffs) {
@ -69,7 +69,7 @@ func (d InternalDiffs) Output() []dax.WorkerDiff {
i := 0
for k, v := range d {
out[i].WorkerID = k
out[i].Address = k
out[i].AddedJobs = v.added.Sorted()
out[i].RemovedJobs = v.removed.Sorted()
i++

8
dax/mds/controller/config.go
View file

@ -11,10 +11,10 @@ import (
type NewBalancerFn func(string, logger.Logger) Balancer
type Config struct {
Director Director
Schemar schemar.Schemar
ComputeBalancer Balancer
TranslateBalancer Balancer
Director Director
Schemar schemar.Schemar
Balancer Balancer
StorageMethod string
BoltDB *boltdb.DB

1125
dax/mds/controller/controller.go
View file

File diff suppressed because it is too large Load diff

584
dax/mds/controller/controller_test.go
View file

File diff suppressed because it is too large Load diff

2
dax/mds/controller/errors.go
View file

@ -96,6 +96,6 @@ func NewErrInvalidRequest(msg string) error {
func NewErrUnassignedJobs(jobs []dax.Job) error {
return errors.New(
ErrCodeUnassignedJobs,
fmt.Sprintf("found %d unassigned jobs", len(jobs)),
fmt.Sprintf("found %d unassigned jobs: %+v", len(jobs), jobs),
)
}

599
dax/mds/controller/naive/balancer.go
View file

@ -1,599 +0,0 @@
// Package naive contains a naive implementation of the Balancer interface.
package naive
import (
"context"
"fmt"
"log"
"math"
"sort"
"strings"
"sync"
"time"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
"github.com/featurebasedb/featurebase/v3/errors"
"github.com/featurebasedb/featurebase/v3/logger"
)
// Ensure type implements interface.
var _ controller.Balancer = (*Balancer)(nil)
// Balancer is a naive implementation of the controller.Balancer interface. It
// helps manage the relationships between workers and jobs. The logic it uses to
// balance jobs across workers is very simple; it bases everything off the
// number of workers and number of jobs. It does not take anything else (such as
// job size, worker capabilities, etc) into consideration.
type Balancer struct {
mu sync.RWMutex
// name is used in logging to help identify the balancer responsible for the
// log.
name string
// current represents the current state of worker/job assigments.
current WorkerJobService
// freeJobs is the set of jobs which have yet to be assigned to a worker.
// This could be because there are no available workers, or because a worker
// has been removed and the jobs for which it was responsible have yet to be
// reassigned.
freeJobs FreeJobService
logger logger.Logger
}
type WorkerJobService interface {
WorkersJobs(ctx context.Context, balancerName string) ([]dax.WorkerInfo, error)
WorkerCount(ctx context.Context, balancerName string) (int, error)
ListWorkers(ctx context.Context, balancerName string) (dax.Workers, error)
WorkerExists(ctx context.Context, balancerName string, worker dax.Worker) (bool, error)
CreateWorker(ctx context.Context, balancerName string, worker dax.Worker) error
DeleteWorker(ctx context.Context, balancerName string, worker dax.Worker) error
CreateJobs(ctx context.Context, balancerName string, worker dax.Worker, job ...dax.Job) error
DeleteJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error
DeleteJobs(ctx context.Context, balancerName, prefix string) (InternalDiffs, error)
JobCounts(ctx context.Context, balancerName string, worker ...dax.Worker) (map[dax.Worker]int, error)
ListJobs(ctx context.Context, balancerName string, worker dax.Worker) (dax.Jobs, error)
}
// TODO: I don't think all these method names need "Free" in them.
type FreeJobService interface {
CreateFreeJobs(ctx context.Context, balancerName string, job ...dax.Job) error
DeleteFreeJob(ctx context.Context, balancerName string, job dax.Job) error
ListFreeJobs(ctx context.Context, balancerName string) (dax.Jobs, error)
MergeFreeJobs(ctx context.Context, balancerName string, jobs dax.Jobs) error
DeleteFreeJobs(ctx context.Context, balancerName, prefix string) error
}
// New returns a new instance of Balancer.
func New(name string, fjs FreeJobService, wjs WorkerJobService, logger logger.Logger) *Balancer {
return &Balancer{
name: name,
current: wjs,
freeJobs: fjs,
logger: logger,
}
}
// AddWorker adds a worker to the Balancer's worker pool. This may cause the
// Balancer to assign existing jobs that are currently in the free list to the
// worker. Also, the worker will immediately be available for assignments of new
// jobs.
func (b *Balancer) AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
b.logger.Debugf("%s: AddWorker(%s)", b.name, worker.String())
b.mu.Lock()
defer b.mu.Unlock()
diff, err := b.addWorker(ctx, dax.Worker(worker.String()))
if err != nil {
return nil, errors.Wrap(err, "adding worker")
}
return diff.Output(), nil
}
func (b *Balancer) addWorker(ctx context.Context, worker dax.Worker) (InternalDiffs, error) {
// If this worker already exists, don't do anything.
if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil {
return nil, errors.Wrap(err, "checking if worker exists")
} else if exists {
return InternalDiffs{}, nil
}
if err := b.current.CreateWorker(ctx, b.name, worker); err != nil {
return nil, errors.Wrap(err, "creating worker")
}
// Process the freeJobs.
return b.processFreeJobs(ctx)
}
// ReplaceWorker is meant to avoid the job re-assignment caused by performing a
// RemoveWorker followed by an AddWorker. In this case, it does both in one step
// so that it's more likely that the jobs will just get transferred directly
// over. NOT IMPLEMENTED YET.
// func (b *Balancer) ReplaceWorker(fromWorker string, toWorker string) []WorkerDiff {
// b.mu.Lock()
// defer b.mu.Unlock()
// return []WorkerDiff{}
// }
// RemoveWorker removes a worker from the worker pool and moves any of its
// currently assigned jobs to the free list. If the intention is to remove a
// worker and reassign its jobs to other workers, then RemoveWorker() should be
// followed by Balance().
func (b *Balancer) RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) {
b.mu.Lock()
defer b.mu.Unlock()
diff, err := b.removeWorker(ctx, dax.Worker(worker.String()))
if err != nil {
return nil, errors.Wrap(err, "removing worker")
}
return diff.Output(), nil
}
func (b *Balancer) removeWorker(ctx context.Context, worker dax.Worker) (InternalDiffs, error) {
// If this worker doesn't exist, don't do anything else.
if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil {
return nil, errors.Wrap(err, "checking if worker exists")
} else if !exists {
return InternalDiffs{}, nil
}
jobs, err := b.current.ListJobs(ctx, b.name, worker)
if err != nil {
return nil, errors.Wrap(err, "listing jobs")
}
// Before removing the worker, mark its jobs as free.
if err := b.freeJobs.MergeFreeJobs(ctx, b.name, jobs); err != nil {
return nil, errors.Wrap(err, "merging free jobs")
}
// Remove the worker.
if err := b.current.DeleteWorker(ctx, b.name, worker); err != nil {
return nil, errors.Wrap(err, "deleting worker")
}
// Even though this may not be useful to the caller (for example, in the
// case where the worker has died and no longer exists), return the diffs
// which represent the removal of jobs from the worker.
diff := NewInternalDiffs()
for _, job := range jobs {
diff.Removed(worker, job)
}
return diff, nil
}
// AddJobs adds one or more jobs to an existing worker. If there are no existing
// workers, the jobs are placed into the free list and will be assigned to a
// worker once one becomes available.
func (b *Balancer) AddJobs(ctx context.Context, jobs ...fmt.Stringer) ([]dax.WorkerDiff, error) {
start := time.Now()
defer func() {
log.Printf("ELAPSED: Balancer.AddJob: %v", time.Since(start))
}()
jobsToAdd := make([]dax.Job, 0, len(jobs))
for _, job := range jobs {
jobsToAdd = append(jobsToAdd, dax.Job(job.String()))
}
if len(jobsToAdd) == 1 {
b.logger.Debugf("%s: AddJobs (%s)", b.name, jobsToAdd[0])
} else {
b.logger.Debugf("%s: AddJobs (%d)", b.name, len(jobsToAdd))
}
b.mu.Lock()
defer b.mu.Unlock()
diff, err := b.addJobs(ctx, jobsToAdd...)
if err != nil {
return nil, errors.Wrap(err, "adding job")
}
return diff.Output(), nil
}
func (b *Balancer) addJobs(ctx context.Context, jobs ...dax.Job) (InternalDiffs, error) {
if cnt, err := b.current.WorkerCount(ctx, b.name); err != nil {
return nil, errors.Wrap(err, "getting worker count")
} else if cnt == 0 {
if err := b.freeJobs.CreateFreeJobs(ctx, b.name, jobs...); err != nil {
return nil, errors.Wrap(err, "creating free job")
}
// TODO: we might want to inform the user that a job is in the free list
// because there are no workers.
return InternalDiffs{}, nil
}
workerJobs, err := b.current.WorkersJobs(ctx, b.name)
if err != nil {
return nil, errors.Wrapf(err, "getting workers jobs: %s", b.name)
}
jset := dax.NewSet[dax.Job]()
for _, workerInfo := range workerJobs {
jset.Merge(dax.NewSet(workerInfo.Jobs...))
}
workerIDs := make(dax.Workers, 0, len(workerJobs))
jobCounts := make(map[dax.Worker]int, 0)
for _, v := range workerJobs {
workerIDs = append(workerIDs, v.ID)
jobCounts[v.ID] = len(v.Jobs)
}
diffs := NewInternalDiffs()
jobsToCreate := make(map[dax.Worker][]dax.Job)
for _, job := range jobs {
// Skip any job that already exists.
if jset.Contains(job) {
continue
}
// Find the worker with the fewest number of jobs and assign it this job.
var lowCount int = math.MaxInt
var lowWorker dax.Worker
// We loop over workerIDs here instead of jobCounts because jobCounts is
// a map and it can return results in an unexpected order, which is a
// problem for testing.
for _, worker := range workerIDs {
jobCount := jobCounts[worker]
if jobCount < lowCount {
lowCount = jobCount
lowWorker = worker
}
}
jobsToCreate[lowWorker] = append(jobsToCreate[lowWorker], job)
jobCounts[lowWorker]++
}
for worker, jobs := range jobsToCreate {
if err := b.current.CreateJobs(ctx, b.name, worker, jobs...); err != nil {
return nil, errors.Wrap(err, "creating job")
}
for _, job := range jobs {
diffs.Added(worker, job)
}
}
return diffs, nil
}
// RemoveJob removes a job from the worker to which is was assigned. If the job
// is not currently assigned to a worker, but it is in the free list, then it
// will be removed from the free list.
func (b *Balancer) RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
b.mu.Lock()
defer b.mu.Unlock()
diff, err := b.removeJob(ctx, dax.Job(job.String()))
if err != nil {
return nil, errors.Wrapf(err, "removing job: %s", job)
}
return diff.Output(), nil
}
func (b *Balancer) RemoveJobs(ctx context.Context, prefix string) ([]dax.WorkerDiff, error) {
b.mu.Lock()
defer b.mu.Unlock()
idiffs, err := b.current.DeleteJobs(ctx, b.name, prefix)
if err != nil {
return nil, errors.Wrap(err, "deleting worker jobs")
}
if err := b.freeJobs.DeleteFreeJobs(ctx, b.name, prefix); err != nil {
return nil, errors.Wrap(err, "deleting free jobs")
}
return idiffs.Output(), nil
}
func (b *Balancer) removeJob(ctx context.Context, job dax.Job) (InternalDiffs, error) {
if worker, ok, err := b.workerForJob(ctx, job); err != nil {
return nil, errors.Wrapf(err, "getting worker for job: %s", job)
} else if ok {
if err := b.current.DeleteJob(ctx, b.name, worker, job); err != nil {
return nil, errors.Wrapf(err, "deleting job: %s", job)
}
diffs := NewInternalDiffs()
diffs.Removed(worker, job)
return diffs, nil
}
// Just in case the job is in the free list (and wasn't assigned to a
// worker), remove it; there's no need to provide a diff. There should never
// be a case where the same job is both in the free list and assigned to a
// worker.
if err := b.freeJobs.DeleteFreeJob(ctx, b.name, job); err != nil {
return nil, errors.Wrapf(err, "deleting free job: %s", job)
}
return InternalDiffs{}, nil
}
// Balance ensures that all jobs are being handled by a worker by assigning jobs
// in the free list to workers, and by moving job assignments around in order to
// balance the load on workers.
func (b *Balancer) Balance(ctx context.Context) ([]dax.WorkerDiff, error) {
b.mu.Lock()
defer b.mu.Unlock()
// If there are no workers, we can't properly balance.
if cnt, err := b.current.WorkerCount(ctx, b.name); err != nil {
return nil, errors.Wrapf(err, "getting worker count: %s", b.name)
} else if cnt == 0 {
return []dax.WorkerDiff{}, nil
}
// Process the freeJobs.
diffs, err := b.processFreeJobs(ctx)
if err != nil {
return nil, errors.Wrapf(err, "processing free jobs: %s", b.name)
}
// Balance the jobs among workers.
diff, err := b.balance(ctx, diffs)
if err != nil {
return nil, errors.Wrap(err, "balancing jobs")
}
return diff.Output(), nil
}
// balance moves jobs among workers with the goal of having an equal number of
// jobs per worker. This method takes an `internalDiffs` as input for cases
// where some action has preceeded this call which also resulted in
// `internalDiffs`. Instead of having this method take a value, we could rely on
// the internalDiffs.merge() method, but we would need to modify that method to
// be smarter about the order in which it applies the add/remove operations.
// Until that's in place, we'll pass in a value here.
func (b *Balancer) balance(ctx context.Context, diffs InternalDiffs) (InternalDiffs, error) {
numWorkers, err := b.current.WorkerCount(ctx, b.name)
if err != nil {
return nil, errors.Wrapf(err, "getting worker count: %s", b.name)
}
numJobs := 0
if workers, err := b.current.ListWorkers(ctx, b.name); err != nil {
return nil, errors.Wrapf(err, "listing workers: %s", b.name)
} else {
for _, worker := range workers {
jobCounts, err := b.current.JobCounts(ctx, b.name, worker)
if err != nil {
return nil, errors.Wrapf(err, "getting job count: %s", worker)
}
numJobs += jobCounts[worker]
}
}
minJobsPerWorker := numJobs / numWorkers
numWorkersAboveMin := numJobs % numWorkers
// sortedWorkerInfos is used now in order to guarantee a sort order.
sortedWorkerInfos, err := b.currentState(ctx, true)
if err != nil {
return nil, errors.Wrapf(err, "getting current state: %s", b.name)
}
// Loop through each worker, and if the number of jobs for the worker
// exceeds the target, then remove the job and add it back (which is
// effectively how we rebalance a job).
for i, workerInfo := range sortedWorkerInfos {
numTargetJobs := minJobsPerWorker
if i < numWorkersAboveMin {
numTargetJobs += 1
}
jobCounts, err := b.current.JobCounts(ctx, b.name, workerInfo.ID)
if err != nil {
return nil, errors.Wrapf(err, "getting job count: %s", workerInfo.ID)
}
numCurrentJobs := jobCounts[workerInfo.ID]
// If we don't need to remove jobs from this worker, then just continue
// on to the next worker.
if numCurrentJobs <= numTargetJobs {
continue
}
sortedJobs, err := b.current.ListJobs(ctx, b.name, workerInfo.ID)
if err != nil {
return nil, errors.Wrapf(err, "listing jobs: %s", workerInfo.ID)
}
// Remove the extra jobs from the end of the list, and add them back
// again (which should place them on a worker with fewer jobs).
for i := numCurrentJobs - 1; i >= numTargetJobs; i-- {
if rj, err := b.removeJob(ctx, sortedJobs[i]); err != nil {
return nil, errors.Wrapf(err, "removing job: %s", sortedJobs[i])
} else {
diffs.Merge(rj)
}
if aj, err := b.addJobs(ctx, sortedJobs[i]); err != nil {
return nil, errors.Wrapf(err, "adding job: %s", sortedJobs[i])
} else {
diffs.Merge(aj)
}
}
}
return diffs, nil
}
// CurrentState returns the current state of worker and job assignments. Note
// that there could be unassigned jobs which are not captured in this output.
// Calling Balance() would force any unassigned jobs to be assigned (assuming
// there is at least one worker), and the output would then reflect that.
func (b *Balancer) CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
return b.currentState(ctx, true)
}
func (b *Balancer) currentState(ctx context.Context, sorted bool) ([]dax.WorkerInfo, error) {
return b.current.WorkersJobs(ctx, b.name)
}
// WorkerState returns the current state of job assignments for a given worker.
func (b *Balancer) WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
return b.workerState(ctx, worker)
}
func (b *Balancer) workerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) {
if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil {
return dax.WorkerInfo{}, errors.Wrapf(err, "checking worker exists: %s", worker)
} else if !exists {
return dax.WorkerInfo{
ID: dax.Worker(worker),
}, nil
}
jobs, err := b.current.ListJobs(ctx, b.name, worker)
if err != nil {
return dax.WorkerInfo{}, errors.Wrapf(err, "listing jobs: %s", worker)
}
return dax.WorkerInfo{
ID: dax.Worker(worker),
Jobs: jobs,
}, nil
}
// WorkersForJobs returns the list of workers for the given jobs. If a given job
// is not currently assigned to a worker, it will be ignored.
func (b *Balancer) WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
return b.workersForJobs(ctx, jobs)
}
func (b *Balancer) workersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) {
out := make(map[dax.Worker]dax.Set[dax.Job])
workerJobs, err := b.current.WorkersJobs(ctx, b.name)
if err != nil {
return nil, errors.Wrapf(err, "getting worker jobs: %s", b.name)
}
for _, workerInfo := range workerJobs {
jset := dax.NewSet(workerInfo.Jobs...)
matches := dax.NewSet[dax.Job]()
for _, job := range jobs {
if jset.Contains(job) {
matches.Add(job)
}
}
if len(matches) > 0 {
out[workerInfo.ID] = matches
}
}
workers := make([]dax.WorkerInfo, len(out))
i := 0
for w, jset := range out {
workers[i] = dax.WorkerInfo{
ID: dax.Worker(w),
Jobs: jset.Sorted(),
}
i++
}
sort.Sort(dax.WorkerInfos(workers))
return workers, nil
}
func (b *Balancer) WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) {
b.mu.RLock()
defer b.mu.RUnlock()
jobs, err := b.freeJobs.ListFreeJobs(ctx, b.name)
if err != nil {
return nil, errors.Wrap(err, "listing free jobs")
}
for _, job := range jobs {
if strings.HasPrefix(string(job), prefix) {
return nil, errors.Errorf("found free job '%s' matching prefix '%s'", job, prefix)
}
}
workerJobs, err := b.current.WorkersJobs(ctx, b.name)
if err != nil {
return nil, errors.Wrapf(err, "getting worker jobs: %s", b.name)
}
result := make([]dax.WorkerInfo, 0)
for _, workerInfo := range workerJobs {
matchedJobs := make([]dax.Job, 0)
for _, job := range workerInfo.Jobs {
if strings.HasPrefix(string(job), prefix) {
matchedJobs = append(matchedJobs, job)
}
}
if len(matchedJobs) > 0 {
result = append(result, dax.WorkerInfo{
ID: workerInfo.ID,
Jobs: matchedJobs,
})
}
}
return result, nil
}
// processFreeJobs assigns all jobs in the free list to a worker.
func (b *Balancer) processFreeJobs(ctx context.Context) (InternalDiffs, error) {
diffs := NewInternalDiffs()
jobs, err := b.freeJobs.ListFreeJobs(ctx, b.name)
if err != nil {
return nil, errors.Wrapf(err, "listing free jobs: %s", b.name)
}
for _, job := range jobs {
if aj, err := b.addJobs(ctx, job); err != nil {
return nil, errors.Wrapf(err, "adding job: %s", job)
} else {
diffs.Merge(aj)
}
if err := b.freeJobs.DeleteFreeJob(ctx, b.name, job); err != nil {
return nil, errors.Wrapf(err, "deleting free job: %s", job)
}
}
return diffs, nil
}
// workerForJob returns the worker currently assigned to the given job.
func (b *Balancer) workerForJob(ctx context.Context, job dax.Job) (dax.Worker, bool, error) {
workerJobs, err := b.current.WorkersJobs(ctx, b.name)
if err != nil {
return "", false, errors.Wrapf(err, "getting workers jobs: %s", b.name)
}
for _, workerInfo := range workerJobs {
jset := dax.NewSet(workerInfo.Jobs...)
if jset.Contains(job) {
return workerInfo.ID, true, nil
}
}
return "", false, nil
}

751
dax/mds/controller/naive/balancer_test.go
View file

@ -1,751 +0,0 @@
package naive_test
import (
"context"
"fmt"
"os"
"testing"
"github.com/featurebasedb/featurebase/v3/dax"
daxbolt "github.com/featurebasedb/featurebase/v3/dax/boltdb"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive/boltdb"
testbolt "github.com/featurebasedb/featurebase/v3/dax/test/boltdb"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/stretchr/testify/assert"
)
func newBoltBalancer(t *testing.T) (*daxbolt.DB, func()) {
db := testbolt.MustOpenDB(t)
assert.NoError(t, db.InitializeBuckets(boltdb.NaiveBalancerBuckets...))
return db, func() {
testbolt.MustCloseDB(t, db)
testbolt.CleanupDB(t, db.Path())
}
}
func TestBalancer(t *testing.T) {
ctx := context.Background()
t.Run("SingleWorker", func(t *testing.T) {
db, cleanup := newBoltBalancer(t)
defer cleanup()
bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
// addJob is a wrapper around bal.AddJobs() which we added when the
// function signature of bal.AddJobs changed to take multiple jobs (and
// it therefore no longer satisfied the fn type in this test).
addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
return bal.AddJobs(ctx, job)
}
tests := []struct {
fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
input string
expDiff []dax.WorkerDiff
expState []dax.WorkerInfo
}{
{
// Add job.
fn: addJob,
input: "p2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{},
},
{
// Add worker.
fn: bal.AddWorker,
input: "n1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p2"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p2"},
},
},
},
{
// Add another job out of order.
fn: addJob,
input: "p1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p1"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
},
},
{
// Add another job.
fn: addJob,
input: "p3",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p3"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p3"},
},
},
},
{
// Add a duplicate job.
fn: addJob,
input: "p2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p3"},
},
},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
diff, err := test.fn(ctx, newStringWrapper(test.input))
assert.NoError(t, err)
assert.Equal(t, test.expDiff, diff)
cs, err := bal.CurrentState(ctx)
assert.NoError(t, err)
assert.Equal(t, test.expState, cs)
})
}
})
t.Run("MultipleWorkers", func(t *testing.T) {
db, cleanup := newBoltBalancer(t)
defer cleanup()
bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
// addJob is a wrapper around bal.AddJobs() which we added when the
// function signature of bal.AddJobs changed to take multiple jobs (and
// it therefore no longer satisfied the fn type in this test).
addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
return bal.AddJobs(ctx, job)
}
tests := []struct {
fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
input string
balance bool
expDiff []dax.WorkerDiff
expState []dax.WorkerInfo
}{
{
// Balance when empty.
balance: true,
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{},
},
{
// Add worker.
fn: bal.AddWorker,
input: "n2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add worker again.
fn: bal.AddWorker,
input: "n2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add a second worker.
fn: bal.AddWorker,
input: "n1",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{},
},
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add job.
fn: addJob,
input: "p2",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p2"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p2"},
},
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add job.
fn: addJob,
input: "p3",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n2",
AddedJobs: []dax.Job{"p3"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p1"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add a third worker.
fn: bal.AddWorker,
input: "n0",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p4",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p4"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p5",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p5"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p0",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n2",
AddedJobs: []dax.Job{"p0"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p6",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p6"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p7",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p7"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
//////////////////// Remove /////////////////////////
{
// Remove nonexistent worker.
fn: bal.RemoveWorker,
input: "nonexistent",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Remove worker.
fn: bal.RemoveWorker,
input: "n1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{},
RemovedJobs: []dax.Job{"p1", "p2", "p7"},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Remove job (from free list).
fn: bal.RemoveJob,
input: "p2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Balance after remove.
balance: true,
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p7"},
RemovedJobs: []dax.Job{},
},
{
WorkerID: "n2",
AddedJobs: []dax.Job{"p1"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p1", "p3"},
},
},
},
{
// Remove job.
fn: bal.RemoveJob,
input: "p1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n2",
AddedJobs: []dax.Job{},
RemovedJobs: []dax.Job{"p1"},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
var diff []dax.WorkerDiff
var err error
if test.balance {
diff, err = bal.Balance(ctx)
} else {
diff, err = test.fn(ctx, newStringWrapper(test.input))
}
assert.NoError(t, err)
assert.Equal(t, test.expDiff, diff)
cs, err := bal.CurrentState(ctx)
assert.NoError(t, err)
assert.Equal(t, test.expState, cs)
})
}
})
t.Run("WorkerState", func(t *testing.T) {
db, cleanup := newBoltBalancer(t)
defer cleanup()
bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
assert.NoError(t, err)
_, err = bal.AddJobs(ctx, newStringWrapper("p1"))
assert.NoError(t, err)
exp := dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p1"},
}
ws, err := bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
// Worker doesn't exist.
exp = dax.WorkerInfo{
ID: "x1",
}
ws, err = bal.WorkerState(ctx, "x1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
})
t.Run("WorkersForJobs", func(t *testing.T) {
db, cleanup := newBoltBalancer(t)
defer cleanup()
bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
assert.NoError(t, err)
_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
assert.NoError(t, err)
for i := 0; i < 12; i++ {
_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
assert.NoError(t, err)
}
exp := dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p0", "p10", "p2", "p4", "p6", "p8"},
}
ws, err := bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n2",
Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
}
ws, err = bal.WorkerState(ctx, "n2")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
tests := []struct {
jobs []dax.Job
exp []dax.WorkerInfo
}{
{
jobs: []dax.Job{"p0"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0"}},
},
},
{
jobs: []dax.Job{"p0", "p4"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
},
},
{
jobs: []dax.Job{"p0", "p4", "p999"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
},
},
{
jobs: []dax.Job{"p0", "p1"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0"}},
{ID: "n2", Jobs: []dax.Job{"p1"}},
},
},
{
jobs: []dax.Job{"p5", "p0", "p1", "p8"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0", "p8"}},
{ID: "n2", Jobs: []dax.Job{"p1", "p5"}},
},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
workers, err := bal.WorkersForJobs(ctx, test.jobs)
assert.NoError(t, err)
assert.Equal(t, test.exp, workers)
})
}
// Some tests for WorkersForJobPrefix
workers, err := bal.WorkersForJobPrefix(ctx, "p1")
assert.NoError(t, err)
assert.ElementsMatch(t, []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p10"}},
{ID: "n2", Jobs: []dax.Job{"p1", "p11"}},
}, workers)
workers, err = bal.WorkersForJobPrefix(ctx, "p2")
assert.NoError(t, err)
assert.ElementsMatch(t, []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p2"}},
}, workers)
workers, err = bal.WorkersForJobPrefix(ctx, "pp")
assert.NoError(t, err)
assert.ElementsMatch(t, []dax.WorkerInfo{}, workers)
})
t.Run("Balance", func(t *testing.T) {
db, cleanup := newBoltBalancer(t)
defer cleanup()
bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr))
// Add two workers with some jobs evenly spread across them.
_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
assert.NoError(t, err)
_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
assert.NoError(t, err)
for i := 0; i < 13; i++ {
_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
assert.NoError(t, err)
}
exp := dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4", "p6", "p8"},
}
ws, err := bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n2",
Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
}
ws, err = bal.WorkerState(ctx, "n2")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
// Now, add a worker and confirm that it currently has no jobs assigned
// to it.
_, err = bal.AddWorker(ctx, newStringWrapper("n3"))
assert.NoError(t, err)
exp = dax.WorkerInfo{
ID: "n3",
Jobs: []dax.Job{},
}
ws, err = bal.WorkerState(ctx, "n3")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
// Finally, call Balance() and confirm that the appropriate jobs got
// reassigned.
_, err = bal.Balance(ctx)
assert.NoError(t, err)
exp = dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4"},
}
ws, err = bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n2",
Jobs: []dax.Job{"p1", "p11", "p3", "p5"},
}
ws, err = bal.WorkerState(ctx, "n2")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n3",
Jobs: []dax.Job{"p6", "p7", "p8", "p9"},
}
ws, err = bal.WorkerState(ctx, "n3")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
})
}
type stringWrapper struct {
s string
}
func newStringWrapper(s string) *stringWrapper {
return &stringWrapper{
s: s,
}
}
func (s *stringWrapper) String() string {
return s.s
}

606
dax/mds/controller/naive/boltdb/balancer.go
View file

@ -1,606 +0,0 @@
// Package boltdb contains the boltdb implementation of the Balancer interface.
package boltdb
import (
"bytes"
"context"
"encoding/json"
"fmt"
"strings"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/dax/boltdb"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive"
"github.com/featurebasedb/featurebase/v3/errors"
"github.com/featurebasedb/featurebase/v3/logger"
)
var (
bucketNaiveBalancer = boltdb.Bucket("naiveBalancer")
)
// NaiveBalancerBuckets defines the buckets used by this package. It can be
// called during setup to create the buckets ahead of time.
var NaiveBalancerBuckets []boltdb.Bucket = []boltdb.Bucket{
bucketNaiveBalancer,
}
// NewBalancer returns a new instance of controller.Balancer.
func NewBalancer(name string, db *boltdb.DB, logger logger.Logger) controller.Balancer {
fjs := newFreeJobService(db)
wjs := newWorkerJobService(db, logger)
return naive.New(name, fjs, wjs, logger)
}
// Ensure type implements interface.
var _ naive.WorkerJobService = (*workerJobService)(nil)
type workerJobService struct {
db *boltdb.DB
logger logger.Logger
}
func newWorkerJobService(db *boltdb.DB, logger logger.Logger) *workerJobService {
return &workerJobService{
db: db,
logger: logger,
}
}
func (w *workerJobService) WorkersJobs(ctx context.Context, balancerName string) ([]dax.WorkerInfo, error) {
tx, err := w.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "getting tx")
}
defer tx.Rollback()
workerInfos, err := getWorkerInfos(ctx, tx, balancerName)
if err != nil {
return nil, errors.Wrapf(err, "getting worker infos: %s", balancerName)
}
return workerInfos, nil
}
func (w *workerJobService) WorkerCount(ctx context.Context, balancerName string) (int, error) {
tx, err := w.db.BeginTx(ctx, false)
if err != nil {
return 0, errors.Wrap(err, "getting tx")
}
defer tx.Rollback()
workers, err := w.getWorkers(ctx, tx, balancerName)
if err != nil {
return 0, errors.Wrapf(err, "getting workers: %s", balancerName)
}
return len(workers), nil
}
func (w *workerJobService) ListWorkers(ctx context.Context, balancerName string) (dax.Workers, error) {
tx, err := w.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
workers, err := w.getWorkers(ctx, tx, balancerName)
if err != nil {
return nil, errors.Wrapf(err, "getting workers: %s", balancerName)
}
return workers, nil
}
func (w *workerJobService) getWorkers(ctx context.Context, tx *boltdb.Tx, balancerName string) (dax.Workers, error) {
c := tx.Bucket(bucketNaiveBalancer).Cursor()
// Deserialize rows into Worker objects.
workers := make(dax.Workers, 0)
prefix := []byte(fmt.Sprintf(prefixFmtWorkers, balancerName))
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
if v == nil {
w.logger.Printf("nil value for key: %s", k)
continue
}
worker, err := keyWorker(k)
if err != nil {
return nil, errors.Wrapf(err, "getting worker from key: %s", k)
}
workers = append(workers, worker)
}
return workers, nil
}
func getWorkerInfos(ctx context.Context, tx *boltdb.Tx, balancerName string) (dax.WorkerInfos, error) {
c := tx.Bucket(bucketNaiveBalancer).Cursor()
// Deserialize rows into WorkerInfo objects.
workerInfos := make(dax.WorkerInfos, 0)
prefix := []byte(fmt.Sprintf(prefixFmtWorkers, balancerName))
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
worker, err := keyWorker(k)
if err != nil {
return nil, errors.Wrapf(err, "getting worker from key: %s", k)
}
jobs := dax.NewSet[dax.Job]()
if v != nil {
jobs, err = decodeJobSet(v)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
workerInfo := dax.WorkerInfo{
ID: worker,
Jobs: jobs.Sorted(),
}
workerInfos = append(workerInfos, workerInfo)
}
return workerInfos, nil
}
func (w *workerJobService) WorkerExists(ctx context.Context, balancerName string, worker dax.Worker) (bool, error) {
tx, err := w.db.BeginTx(ctx, false)
if err != nil {
return false, errors.Wrapf(err, "getting tx: %s", balancerName)
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return false, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
wrkr := bkt.Get(workerKey(balancerName, worker))
return wrkr != nil, nil
}
func (w *workerJobService) CreateWorker(ctx context.Context, balancerName string, worker dax.Worker) error {
tx, err := w.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "getting transaction")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
// If this worker already exists, don't do anything.
wrkr := bkt.Get(workerKey(balancerName, worker))
if wrkr != nil {
return nil
}
val := []byte("[]")
if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
return errors.Wrap(err, "putting worker")
}
return tx.Commit()
}
func (w *workerJobService) DeleteWorker(ctx context.Context, balancerName string, worker dax.Worker) error {
tx, err := w.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
if err := bkt.Delete(workerKey(balancerName, worker)); err != nil {
return errors.Wrapf(err, "deleting node key: %s", workerKey(balancerName, worker))
}
return tx.Commit()
}
func (w *workerJobService) CreateJobs(ctx context.Context, balancerName string, worker dax.Worker, jobs ...dax.Job) error {
tx, err := w.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
jobset := dax.NewSet[dax.Job]()
// get worker
wrkr := bkt.Get(workerKey(balancerName, worker))
if wrkr != nil {
jobset, err = decodeJobSet(wrkr)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
}
for _, job := range jobs {
jobset.Add(job)
}
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
return errors.Wrap(err, "putting worker")
}
return tx.Commit()
}
func (w *workerJobService) DeleteJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error {
tx, err := w.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
// get worker
wrkr := bkt.Get(workerKey(balancerName, worker))
if wrkr == nil {
return nil
}
jobset, err := decodeJobSet(wrkr)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
if !jobset.Contains(job) {
return nil
}
jobset.Remove(job)
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
return errors.Wrap(err, "putting worker")
}
return tx.Commit()
}
func (w *workerJobService) DeleteJobs(ctx context.Context, balancerName, prefix string) (naive.InternalDiffs, error) {
tx, err := w.db.BeginTx(ctx, true)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
workers, err := w.getWorkers(ctx, tx, balancerName)
if err != nil {
return nil, errors.Wrap(err, "getting workers")
}
idiffs := naive.NewInternalDiffs()
for _, worker := range workers {
// get worker
wrkr := bkt.Get(workerKey(balancerName, worker))
if wrkr == nil {
panic("didn't find worker that should... definitely exist")
}
jobset, err := decodeJobSet(wrkr)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
jobs := jobset.RemovePrefix(prefix)
for _, job := range jobs {
idiffs.Removed(worker, job)
}
val, err := encodeJobSet(jobset)
if err != nil {
return nil, errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(workerKey(balancerName, worker), val); err != nil {
return nil, errors.Wrap(err, "putting worker")
}
}
return idiffs, tx.Commit()
}
func (w *workerJobService) ListJobs(ctx context.Context, balancerName string, worker dax.Worker) (dax.Jobs, error) {
tx, err := w.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
jobset := dax.NewSet[dax.Job]()
// get worker
wrkr := bkt.Get(workerKey(balancerName, worker))
if wrkr != nil {
jobset, err = decodeJobSet(wrkr)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
return jobset.Sorted(), nil
}
func (w *workerJobService) JobCounts(ctx context.Context, balancerName string, workers ...dax.Worker) (map[dax.Worker]int, error) {
tx, err := w.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrapf(err, "getting tx: %s", balancerName)
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
m := make(map[dax.Worker]int)
for _, worker := range workers {
jobset := dax.NewSet[dax.Job]()
// get worker
wrkr := bkt.Get(workerKey(balancerName, worker))
if wrkr != nil {
jobset, err = decodeJobSet(wrkr)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
m[worker] = len(jobset)
}
return m, nil
}
// encodeJobSet encode the jobSet into a JSON array of strings.
func encodeJobSet(jobSet dax.Set[dax.Job]) ([]byte, error) {
arr := jobSet.Sorted()
b, err := json.Marshal(arr)
if err != nil {
return nil, errors.Wrap(err, "marshalling json")
}
return b, nil
}
// decodeJobSet decode the string (a JSON array of strings) into jobSet.
func decodeJobSet(v []byte) (dax.Set[dax.Job], error) {
var arr []string
err := json.Unmarshal(v, &arr)
if err != nil {
return nil, errors.Wrap(err, "unmarshalling json")
}
js := dax.NewSet[dax.Job]()
for _, s := range arr {
js.Add(dax.Job(s))
}
return js, nil
}
// Ensure type implements interface.
var _ naive.FreeJobService = (*freeJobService)(nil)
type freeJobService struct {
db *boltdb.DB
}
func newFreeJobService(db *boltdb.DB) *freeJobService {
return &freeJobService{
db: db,
}
}
func (f *freeJobService) CreateFreeJobs(ctx context.Context, balancerName string, jobs ...dax.Job) error {
return f.MergeFreeJobs(ctx, balancerName, jobs)
}
func (f *freeJobService) DeleteFreeJob(ctx context.Context, balancerName string, job dax.Job) error {
tx, err := f.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
// get free jobs
fjs := bkt.Get(freeJobKey(balancerName))
if fjs == nil {
return nil
}
jobset, err := decodeJobSet(fjs)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
if !jobset.Contains(job) {
return nil
}
jobset.Remove(job)
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(freeJobKey(balancerName), val); err != nil {
return errors.Wrap(err, "putting free job")
}
return tx.Commit()
}
func (f *freeJobService) DeleteFreeJobs(ctx context.Context, balancerName, prefix string) error {
tx, err := f.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
// get free jobs
fjs := bkt.Get(freeJobKey(balancerName))
if fjs == nil {
return nil
}
jobset, err := decodeJobSet(fjs)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
jobset.RemovePrefix(prefix)
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(freeJobKey(balancerName), val); err != nil {
return errors.Wrap(err, "putting free job")
}
return tx.Commit()
}
func (f *freeJobService) ListFreeJobs(ctx context.Context, balancerName string) (dax.Jobs, error) {
tx, err := f.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
jobset := dax.NewSet[dax.Job]()
// get free jobs
fjs := bkt.Get(freeJobKey(balancerName))
if fjs != nil {
jobset, err = decodeJobSet(fjs)
if err != nil {
return nil, errors.Wrap(err, "decoding job set")
}
}
return jobset.Sorted(), nil
}
func (f *freeJobService) MergeFreeJobs(ctx context.Context, balancerName string, jobs dax.Jobs) error {
tx, err := f.db.BeginTx(ctx, true)
if err != nil {
return err
}
defer tx.Rollback()
bkt := tx.Bucket(bucketNaiveBalancer)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer)
}
jobset := dax.NewSet[dax.Job]()
// get free jobs
fjs := bkt.Get(freeJobKey(balancerName))
if fjs != nil {
jobset, err = decodeJobSet(fjs)
if err != nil {
return errors.Wrap(err, "decoding job set")
}
}
for _, j := range jobs {
jobset.Add(j)
}
val, err := encodeJobSet(jobset)
if err != nil {
return errors.Wrap(err, "encoding job set")
}
if err := bkt.Put(freeJobKey(balancerName), val); err != nil {
return errors.Wrap(err, "putting free job")
}
return tx.Commit()
}
//////////////////////////////////////////////////////
const (
prefixFmtWorkers = "workers/%s/" // %s - balancerName
prefixFmtFreeJobs = "freejobs/%s" // %s - balancerName
)
// workerKey returns a key based on worker.
func workerKey(bal string, worker dax.Worker) []byte {
key := fmt.Sprintf(prefixFmtWorkers+"%s", bal, worker)
return []byte(key)
}
// keyWorker gets the worker out of the key.
func keyWorker(key []byte) (dax.Worker, error) {
parts := strings.SplitN(string(key), "/", 3)
if len(parts) != 3 {
return "", errors.New(errors.ErrUncoded, "worker key format expected: `workers/balancer/worker`")
}
return dax.Worker(parts[2]), nil
}
// freeJobKey returns a key for all freeJobs.
func freeJobKey(bal string) []byte {
key := fmt.Sprintf(prefixFmtFreeJobs, bal)
return []byte(key)
}

721
dax/mds/controller/naive/boltdb/balancer_test.go
View file

@ -1,721 +0,0 @@
package boltdb_test
import (
"context"
"fmt"
"testing"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive/boltdb"
testbolt "github.com/featurebasedb/featurebase/v3/dax/test/boltdb"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/stretchr/testify/assert"
)
func TestBalancer(t *testing.T) {
db := testbolt.MustOpenDB(t)
defer testbolt.MustCloseDB(t, db)
t.Cleanup(func() {
testbolt.CleanupDB(t, db.Path())
})
ctx := context.Background()
// Initialize the buckets.
assert.NoError(t, db.InitializeBuckets(boltdb.NaiveBalancerBuckets...))
t.Run("SingleWorker", func(t *testing.T) {
bal := boltdb.NewBalancer("test-single-worker", db, logger.NopLogger)
// addJob is a wrapper around bal.AddJobs() which we added when the
// function signature of bal.AddJobs changed to take multiple jobs (and
// it therefore no longer satisfied the fn type in this test).
addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
return bal.AddJobs(ctx, job)
}
tests := []struct {
fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
input string
expDiff []dax.WorkerDiff
expState []dax.WorkerInfo
}{
{
// Add job.
fn: addJob,
input: "p2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{},
},
{
// Add worker.
fn: bal.AddWorker,
input: "n1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p2"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p2"},
},
},
},
{
// Add another job out of order.
fn: addJob,
input: "p1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p1"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
},
},
{
// Add another job.
fn: addJob,
input: "p3",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p3"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p3"},
},
},
},
{
// Add a duplicate job.
fn: addJob,
input: "p2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p3"},
},
},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
diff, err := test.fn(ctx, newStringWrapper(test.input))
assert.NoError(t, err)
assert.Equal(t, test.expDiff, diff)
cs, err := bal.CurrentState(ctx)
assert.NoError(t, err)
assert.Equal(t, test.expState, cs)
})
}
})
t.Run("MultipleWorkers", func(t *testing.T) {
bal := boltdb.NewBalancer("test-multiple-workers", db, logger.NopLogger)
// addJob is a wrapper around bal.AddJobs() which we added when the
// function signature of bal.AddJobs changed to take multiple jobs (and
// it therefore no longer satisfied the fn type in this test).
addJob := func(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) {
return bal.AddJobs(ctx, job)
}
tests := []struct {
fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error)
input string
balance bool
expDiff []dax.WorkerDiff
expState []dax.WorkerInfo
}{
{
// Balance when empty.
balance: true,
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{},
},
{
// Add worker.
fn: bal.AddWorker,
input: "n2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add worker again.
fn: bal.AddWorker,
input: "n2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add a second worker.
fn: bal.AddWorker,
input: "n1",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{},
},
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add job.
fn: addJob,
input: "p2",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p2"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p2"},
},
{
ID: "n2",
Jobs: []dax.Job{},
},
},
},
{
// Add job.
fn: addJob,
input: "p3",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n2",
AddedJobs: []dax.Job{"p3"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p1"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add a third worker.
fn: bal.AddWorker,
input: "n0",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p4",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p4"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p5",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p5"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p0",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n2",
AddedJobs: []dax.Job{"p0"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p6",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p6"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Add job.
fn: addJob,
input: "p7",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{"p7"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
//////////////////// Remove /////////////////////////
{
// Remove nonexistent worker.
fn: bal.RemoveWorker,
input: "nonexistent",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n1",
Jobs: []dax.Job{"p1", "p2", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Remove worker.
fn: bal.RemoveWorker,
input: "n1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n1",
AddedJobs: []dax.Job{},
RemovedJobs: []dax.Job{"p1", "p2", "p7"},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Remove job (from free list).
fn: bal.RemoveJob,
input: "p2",
expDiff: []dax.WorkerDiff{},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
{
// Balance after remove.
balance: true,
expDiff: []dax.WorkerDiff{
{
WorkerID: "n0",
AddedJobs: []dax.Job{"p7"},
RemovedJobs: []dax.Job{},
},
{
WorkerID: "n2",
AddedJobs: []dax.Job{"p1"},
RemovedJobs: []dax.Job{},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p1", "p3"},
},
},
},
{
// Remove job.
fn: bal.RemoveJob,
input: "p1",
expDiff: []dax.WorkerDiff{
{
WorkerID: "n2",
AddedJobs: []dax.Job{},
RemovedJobs: []dax.Job{"p1"},
},
},
expState: []dax.WorkerInfo{
{
ID: "n0",
Jobs: []dax.Job{"p4", "p5", "p6", "p7"},
},
{
ID: "n2",
Jobs: []dax.Job{"p0", "p3"},
},
},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
var diff []dax.WorkerDiff
var err error
if test.balance {
diff, err = bal.Balance(ctx)
} else {
diff, err = test.fn(ctx, newStringWrapper(test.input))
}
assert.NoError(t, err)
assert.Equal(t, test.expDiff, diff)
cs, err := bal.CurrentState(ctx)
assert.NoError(t, err)
assert.Equal(t, test.expState, cs)
})
}
})
t.Run("WorkerState", func(t *testing.T) {
bal := boltdb.NewBalancer("test-worker-state", db, logger.NopLogger)
_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
assert.NoError(t, err)
_, err = bal.AddJobs(ctx, newStringWrapper("p1"))
assert.NoError(t, err)
exp := dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p1"},
}
ws, err := bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
// Worker doesn't exist.
exp = dax.WorkerInfo{
ID: "x1",
}
ws, err = bal.WorkerState(ctx, "x1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
})
t.Run("WorkersForJobs", func(t *testing.T) {
bal := boltdb.NewBalancer("test-workers-for-jobs", db, logger.NopLogger)
_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
assert.NoError(t, err)
_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
assert.NoError(t, err)
for i := 0; i < 12; i++ {
_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
assert.NoError(t, err)
}
exp := dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p0", "p10", "p2", "p4", "p6", "p8"},
}
ws, err := bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n2",
Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
}
ws, err = bal.WorkerState(ctx, "n2")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
tests := []struct {
jobs []dax.Job
exp []dax.WorkerInfo
}{
{
jobs: []dax.Job{"p0"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0"}},
},
},
{
jobs: []dax.Job{"p0", "p4"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
},
},
{
jobs: []dax.Job{"p0", "p4", "p999"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0", "p4"}},
},
},
{
jobs: []dax.Job{"p0", "p1"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0"}},
{ID: "n2", Jobs: []dax.Job{"p1"}},
},
},
{
jobs: []dax.Job{"p5", "p0", "p1", "p8"},
exp: []dax.WorkerInfo{
{ID: "n1", Jobs: []dax.Job{"p0", "p8"}},
{ID: "n2", Jobs: []dax.Job{"p1", "p5"}},
},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
workers, err := bal.WorkersForJobs(ctx, test.jobs)
assert.NoError(t, err)
assert.Equal(t, test.exp, workers)
})
}
})
t.Run("Balance", func(t *testing.T) {
bal := boltdb.NewBalancer("test-balance", db, logger.NopLogger)
// Add two workers with some jobs evenly spread across them.
_, err := bal.AddWorker(ctx, newStringWrapper("n1"))
assert.NoError(t, err)
_, err = bal.AddWorker(ctx, newStringWrapper("n2"))
assert.NoError(t, err)
for i := 0; i < 13; i++ {
_, err = bal.AddJobs(ctx, newStringWrapper(fmt.Sprintf("p%d", i)))
assert.NoError(t, err)
}
exp := dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4", "p6", "p8"},
}
ws, err := bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n2",
Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"},
}
ws, err = bal.WorkerState(ctx, "n2")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
// Now, add a worker and confirm that it currently has no jobs assigned
// to it.
_, err = bal.AddWorker(ctx, newStringWrapper("n3"))
assert.NoError(t, err)
exp = dax.WorkerInfo{
ID: "n3",
Jobs: []dax.Job{},
}
ws, err = bal.WorkerState(ctx, "n3")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
// Finally, call Balance() and confirm that the appropriate jobs got
// reassigned.
_, err = bal.Balance(ctx)
assert.NoError(t, err)
exp = dax.WorkerInfo{
ID: "n1",
Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4"},
}
ws, err = bal.WorkerState(ctx, "n1")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n2",
Jobs: []dax.Job{"p1", "p11", "p3", "p5"},
}
ws, err = bal.WorkerState(ctx, "n2")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
exp = dax.WorkerInfo{
ID: "n3",
Jobs: []dax.Job{"p6", "p7", "p8", "p9"},
}
ws, err = bal.WorkerState(ctx, "n3")
assert.NoError(t, err)
assert.Equal(t, exp, ws)
})
}
type stringWrapper struct {
s string
}
func newStringWrapper(s string) *stringWrapper {
return &stringWrapper{
s: s,
}
}
func (s *stringWrapper) String() string {
return s.s
}

6
dax/mds/controller/sets.go
View file

@ -85,11 +85,11 @@ func (s TableSet) SortedSlice() dax.TableKeys {
return ps
}
func (s TableSet) QualifiedSortedSlice() map[dax.TableQualifier]dax.TableIDs {
m := make(map[dax.TableQualifier]dax.TableIDs)
func (s TableSet) QualifiedSortedSlice() map[dax.QualifiedDatabaseID]dax.TableIDs {
m := make(map[dax.QualifiedDatabaseID]dax.TableIDs)
for p := range s {
qtid := p.QualifiedTableID()
m[qtid.TableQualifier] = append(m[qtid.TableQualifier], qtid.ID)
m[qtid.QualifiedDatabaseID] = append(m[qtid.QualifiedDatabaseID], qtid.ID)
}
// Sort the slices in the map.

43
dax/mds/controller/snapping_turtle.go
View file

@ -30,7 +30,29 @@ func (c *Controller) snappingTurtleRoutine(period time.Duration, control chan st
func (c *Controller) snapAll() {
c.logger.Debugf("TURTLE: snapAll")
ctx := context.Background()
computeNodes, err := c.ComputeBalancer.CurrentState(ctx)
tx, err := c.boltDB.BeginTx(ctx, false)
if err != nil {
c.logger.Printf("Error getting transaction for snapping turtle: %v", err)
return
}
defer tx.Rollback()
qdbs, err := c.Schemar.Databases(tx, "")
if err != nil {
c.logger.Printf("couldn't get databases: %v", err)
}
for _, qdb := range qdbs {
c.snapAllForDatabase(tx, qdb.QualifiedID())
}
c.logger.Debugf("TURTLE: snapAll complete")
}
func (c *Controller) snapAllForDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) {
c.logger.Debugf("TURTLE: snapAllForDatabase: %s", qdbid)
computeNodes, err := c.Balancer.CurrentState(tx, dax.RoleTypeCompute, qdbid)
if err != nil {
c.logger.Printf("Error getting compute balancer state for snapping turtle: %v", err)
}
@ -52,15 +74,17 @@ func (c *Controller) snapAll() {
if err != nil {
c.logger.Printf("couldn't decode a shard out of the job: '%s', err: %v", workerInfo.Jobs[i], err)
}
c.SnapshotShardData(ctx, j.t.QualifiedTableID(), j.shardNum())
if err := c.snapshotShardData(tx, j.t.QualifiedTableID(), j.shardNum()); err != nil {
c.logger.Printf("Couldn't snapshot table: %s, shard: %d, error: %v", j.t, j.shardNum(), err)
}
}
i++
}
// Get all tables across all orgs/dbs so we can snapshot all keyed
// fields and look up whether a table is keyed to snapshot it's
// fields and look up whether a table is keyed to snapshot its
// partitions.
tables, err := c.Schemar.Tables(ctx, dax.TableQualifier{})
tables, err := c.Schemar.Tables(tx, dax.QualifiedDatabaseID{})
if err != nil {
c.logger.Printf("Couldn't get schema for snapshotting keys: %v", err)
return
@ -71,8 +95,7 @@ func (c *Controller) snapAll() {
tableMap[table.Key()] = table
for _, f := range table.Fields {
if f.StringKeys() && !f.IsPrimaryKey() {
err := c.SnapshotFieldKeys(ctx, table.QualifiedID(), f.Name)
if err != nil {
if err := c.snapshotFieldKeys(tx, table.QualifiedID(), f.Name); err != nil {
c.logger.Printf("Couldn't snapshot table: %s, field: %s, error: %v", table, f.Name, err)
}
}
@ -83,7 +106,7 @@ func (c *Controller) snapAll() {
// for any partition that goes with a keyed table. Doing the same
// weird nested loop thing to avoid doing all jobs on one node
// back to back.
translateNodes, err := c.TranslateBalancer.CurrentState(ctx)
translateNodes, err := c.Balancer.CurrentState(tx, dax.RoleTypeTranslate, qdbid)
if err != nil {
c.logger.Printf("Error getting translate balancer state for snapping turtle: %v", err)
}
@ -101,12 +124,14 @@ func (c *Controller) snapAll() {
if err != nil {
table := tableMap[j.table()]
if table.StringKeys() {
c.SnapshotTableKeys(ctx, table.QualifiedID(), j.partitionNum())
if err := c.snapshotTableKeys(tx, table.QualifiedID(), j.partitionNum()); err != nil {
c.logger.Printf("Couldn't snapshot table: %s, partition: %d, error: %v", table, j.partitionNum(), err)
}
}
c.logger.Printf("couldn't decode a partition out of the job: '%s', err: %v", workerInfo.Jobs[i], err)
}
}
i++
}
c.logger.Debugf("TURTLE: snapAll complete")
c.logger.Debugf("TURTLE: snapAllForDatabase complete: %s", qdbid)
}

8
dax/mds/controller/stringers.go
View file

@ -20,6 +20,10 @@ func (p pUnit) String() string {
return fmt.Sprintf("%s|part_%d", p.t, p.p)
}
func (p pUnit) Job() dax.Job {
return dax.Job(fmt.Sprintf("%s|part_%d", p.t, p.p))
}
func (p pUnit) table() dax.TableKey {
return p.t
}
@ -72,6 +76,10 @@ func (s sUnit) String() string {
return fmt.Sprintf("%s|shard_%s", s.t, s.s)
}
func (s sUnit) Job() dax.Job {
return dax.Job(fmt.Sprintf("%s|shard_%s", s.t, s.s))
}
func (s sUnit) table() dax.TableKey {
return s.t
}

63
dax/mds/http/handler.go
View file

@ -18,6 +18,9 @@ func Handler(mds *mds.MDS) http.Handler {
router.HandleFunc("/health", server.getHealth).Methods("GET").Name("GetHealth")
// mds endpoints.
router.HandleFunc("/create-database", server.postCreateDatabase).Methods("POST").Name("PostCreateDatabase")
router.HandleFunc("/database-by-id", server.postDatabaseByID).Methods("POST").Name("PostDatabaseByID")
router.HandleFunc("/create-table", server.postCreateTable).Methods("POST").Name("PostCreateTable")
router.HandleFunc("/drop-table", server.postDropTable).Methods("POST").Name("PostDropTable")
router.HandleFunc("/create-field", server.postCreateField).Methods("POST").Name("PostCreateField")
@ -57,6 +60,55 @@ func (s *server) getHealth(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}
// POST /create-database
func (s *server) postCreateDatabase(w http.ResponseWriter, r *http.Request) {
body := r.Body
defer body.Close()
ctx := r.Context()
req := &dax.QualifiedDatabase{}
if err := json.NewDecoder(body).Decode(req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
err := s.mds.CreateDatabase(ctx, req)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
if err := json.NewEncoder(w).Encode(req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
// POST /database
func (s *server) postDatabaseByID(w http.ResponseWriter, r *http.Request) {
body := r.Body
defer body.Close()
ctx := r.Context()
qdbid := dax.QualifiedDatabaseID{}
if err := json.NewDecoder(body).Decode(&qdbid); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
resp, err := s.mds.DatabaseByID(ctx, qdbid)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
if err := json.NewEncoder(w).Encode(resp); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
// POST /create-table
func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) {
body := r.Body
@ -76,15 +128,12 @@ func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) {
return
}
resp := CreateTableResponse(*req)
if err := json.NewEncoder(w).Encode(resp); err != nil {
if err := json.NewEncoder(w).Encode(req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
type CreateTableResponse dax.QualifiedTable
// POST /table
func (s *server) postTable(w http.ResponseWriter, r *http.Request) {
body := r.Body
@ -122,7 +171,7 @@ func (s *server) postTableID(w http.ResponseWriter, r *http.Request) {
return
}
qtid, err := s.mds.TableID(ctx, req.TableQualifier, req.Name)
qtid, err := s.mds.TableID(ctx, req.QualifiedDatabaseID, req.Name)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
@ -221,10 +270,10 @@ func (s *server) postTables(w http.ResponseWriter, r *http.Request) {
return
}
qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
qdbid := dax.NewQualifiedDatabaseID(req.OrganizationID, req.DatabaseID)
ids := req.TableIDs
resp, err := s.mds.Tables(ctx, qual, ids...)
resp, err := s.mds.Tables(ctx, qdbid, ids...)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return

247
dax/mds/mds.go
View file

@ -3,7 +3,6 @@ package mds
import (
"context"
"fmt"
"os"
"sync"
"time"
@ -12,7 +11,7 @@ import (
"github.com/featurebasedb/featurebase/v3/dax/boltdb"
"github.com/featurebasedb/featurebase/v3/dax/computer"
"github.com/featurebasedb/featurebase/v3/dax/mds/controller"
naiveboltdb "github.com/featurebasedb/featurebase/v3/dax/mds/controller/naive/boltdb"
balancerboltdb "github.com/featurebasedb/featurebase/v3/dax/mds/controller/balancer/boltdb"
"github.com/featurebasedb/featurebase/v3/dax/mds/poller"
"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
schemarboltdb "github.com/featurebasedb/featurebase/v3/dax/mds/schemar/boltdb"
@ -57,7 +56,6 @@ type MDS struct {
// we need to be sure to close the boltDBs that are created in mds.New()
// whenever mds.Close() is called. These are pointers to those DBs so we can
// close them.
schemarDB *boltdb.DB
controllerDB *boltdb.DB
logger logger.Logger
@ -88,25 +86,20 @@ func New(cfg Config) *MDS {
logr.Warnf("no DataDir given (like '/path/to/directory') using temp dir at '%s'", cfg.DataDir)
}
schemarDB, err := boltdb.NewSvcBolt(cfg.DataDir, "schemar", schemarboltdb.SchemarBuckets...)
buckets := append(schemarboltdb.SchemarBuckets, balancerboltdb.BalancerBuckets...)
controllerDB, err := boltdb.NewSvcBolt(cfg.DataDir, "controller", buckets...)
if err != nil {
logr.Printf("Error creating schemar db: %v", err)
logr.Printf(errors.Wrap(err, "creating controller bolt").Error())
os.Exit(1)
}
schemar := schemarboltdb.NewSchemar(schemarDB, logr)
controllerDB, err := boltdb.NewSvcBolt(cfg.DataDir, "balancer", naiveboltdb.NaiveBalancerBuckets...)
if err != nil {
logr.Printf(errors.Wrap(err, "creating balancer bolt").Error())
os.Exit(1)
}
schemar := schemarboltdb.NewSchemar(controllerDB, logr)
controllerCfg := controller.Config{
Director: cfg.Director,
Schemar: schemar,
ComputeBalancer: naiveboltdb.NewBalancer("compute", controllerDB, logr),
TranslateBalancer: naiveboltdb.NewBalancer("translate", controllerDB, logr),
Director: cfg.Director,
Schemar: schemar,
Balancer: balancerboltdb.NewBalancer(controllerDB, schemar, logr),
RegistrationBatchTimeout: cfg.RegistrationBatchTimeout,
SnappingTurtleTimeout: cfg.SnappingTurtleTimeout,
@ -140,7 +133,6 @@ func New(cfg Config) *MDS {
poller: poller,
schemar: schemar,
schemarDB: schemarDB,
controllerDB: controllerDB,
logger: logr,
@ -169,9 +161,6 @@ func (m *MDS) Stop() error {
m.poller.Stop()
m.controller.Stop()
if m.schemarDB != nil {
m.schemarDB.Close()
}
if m.controllerDB != nil {
m.controllerDB.Close()
}
@ -179,38 +168,17 @@ func (m *MDS) Stop() error {
return nil
}
// sanitizeQTID populates Table.ID (by looking up the table, by name, in
// schemar) for a given table having only a Name value, but no ID.
func (m *MDS) sanitizeQTID(ctx context.Context, qtid *dax.QualifiedTableID) error {
if qtid.ID == "" {
nqtid, err := m.schemar.TableID(ctx, qtid.TableQualifier, qtid.Name)
if err != nil {
return errors.Wrap(err, "getting table ID")
}
qtid.ID = nqtid.ID
}
return nil
// CreateDatabase handles a create table request.
func (m *MDS) CreateDatabase(ctx context.Context, qdb *dax.QualifiedDatabase) error {
return m.controller.CreateDatabase(ctx, qdb)
}
func (m *MDS) DatabaseByID(ctx context.Context, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
return m.controller.DatabaseByID(ctx, qdbid)
}
// CreateTable handles a create table request.
func (m *MDS) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
m.mu.Lock()
defer m.mu.Unlock()
// Create Table ID.
if _, err := qtbl.CreateID(); err != nil {
return errors.Wrap(err, "creating table ID")
}
// Create the table in schemar.
if err := m.schemar.CreateTable(ctx, qtbl); err != nil {
return errors.Wrapf(err, "creating table: %s", qtbl)
}
// TODO: if error here, we should probably roll-back the
// schemar.CreateTable() request.
// Add the table to the controller.
return m.controller.CreateTable(ctx, qtbl)
}
@ -218,216 +186,61 @@ func (m *MDS) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
// reason about consistency here? What if controller DropTable
// succeeds, but schemar fails?
func (m *MDS) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error {
m.mu.Lock()
defer m.mu.Unlock()
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
if err := m.controller.DropTable(ctx, qtid); err != nil {
return errors.Wrapf(err, "dropping table: %s", qtid)
}
return m.schemar.DropTable(ctx, qtid)
}
type CreateFieldRequest struct {
Table dax.TableName
Field *dax.Field
return m.controller.DropTable(ctx, qtid)
}
// CreateField handles a create Field request.
func (m *MDS) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error {
m.mu.Lock()
defer m.mu.Unlock()
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
// Create the field in schemar.
if err := m.schemar.CreateField(ctx, qtid, fld); err != nil {
return errors.Wrapf(err, "creating field: %s, %s", qtid, fld)
}
// Add the table to the controller.
return m.controller.CreateField(ctx, qtid, fld)
}
// DropField handles a drop Field request.
func (m *MDS) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
m.mu.Lock()
defer m.mu.Unlock()
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
// Drop the field from schemar.
if err := m.schemar.DropField(ctx, qtid, fldName); err != nil {
return errors.Wrapf(err, "dropping field: %s, %s", qtid, fldName)
}
// Drop the field from the controller.
return m.controller.DropField(ctx, qtid, fldName)
}
type DropFieldResponse struct{}
// Table handles a table request.
func (m *MDS) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
m.mu.RLock()
defer m.mu.RUnlock()
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return nil, errors.Wrap(err, "sanitizing")
}
return m.schemar.Table(ctx, qtid)
return m.controller.Table(ctx, qtid)
}
// Tables handles a tables request.
func (m *MDS) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
m.mu.RLock()
defer m.mu.RUnlock()
return m.schemar.Tables(ctx, qual, ids...)
func (m *MDS) Tables(ctx context.Context, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
return m.controller.Tables(ctx, qdbid, ids...)
}
// TableID handles a table id (i.e. by name) request.
func (m *MDS) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
m.mu.RLock()
defer m.mu.RUnlock()
return m.schemar.TableID(ctx, qual, name)
func (m *MDS) TableID(ctx context.Context, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
return m.controller.TableID(ctx, qdbid, name)
}
// IngestPartition handles an ingest partition request.
func (m *MDS) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partnNum dax.PartitionNum) (dax.Address, error) {
m.mu.RLock()
defer m.mu.RUnlock()
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return "", errors.Wrap(err, "sanitizing")
}
// Verify that the table exists.
if _, err := m.schemar.Table(ctx, qtid); err != nil {
return "", err
}
partitions := dax.PartitionNums{partnNum}
nodes, err := m.controller.TranslateNodes(ctx, qtid, partitions, true)
if err != nil {
return "", err
}
if l := len(nodes); l == 0 {
return "", controller.NewErrNoAvailableNode()
} else if l > 1 {
return "", controller.NewErrInternal(
fmt.Sprintf("unexpected number of nodes: %d", l))
}
node := nodes[0]
// Verify that the node returned is actually responsible for the partition
// requested.
if node.Table != qtid.Key() {
return "", controller.NewErrInternal(
fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid))
} else if l := len(node.Partitions); l != 1 {
return "", controller.NewErrInternal(
fmt.Sprintf("unexpected number of partitions returned: %d", l))
} else if p := node.Partitions[0]; p != partnNum {
return "", controller.NewErrInternal(
fmt.Sprintf("partition returned (%d) does not match requested (%d)", p, partnNum))
}
return node.Address, nil
return m.controller.IngestPartition(ctx, qtid, partnNum)
}
// IngestShard handles an ingest shard request.
func (m *MDS) IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shrdNum dax.ShardNum) (dax.Address, error) {
m.mu.RLock()
defer m.mu.RUnlock()
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return "", errors.Wrap(err, "sanitizing")
}
// Verify that the table exists.
if _, err := m.schemar.Table(ctx, qtid); err != nil {
return "", err
}
shards := dax.ShardNums{shrdNum}
nodes, err := m.controller.ComputeNodes(ctx, qtid, shards, true)
if err != nil {
return "", err
}
if l := len(nodes); l == 0 {
return "", controller.NewErrNoAvailableNode()
} else if l > 1 {
return "", controller.NewErrInternal(
fmt.Sprintf("unexpected number of nodes: %d", l))
}
node := nodes[0]
// Verify that the node returned is actually responsible for the shard
// requested.
if node.Table != qtid.Key() {
return "", controller.NewErrInternal(
fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid))
} else if l := len(node.Shards); l != 1 {
return "", controller.NewErrInternal(
fmt.Sprintf("unexpected number of shards returned: %d", l))
} else if s := node.Shards[0]; s != shrdNum {
return "", controller.NewErrInternal(
fmt.Sprintf("shard returned (%d) does not match requested (%d)", s, shrdNum))
}
return node.Address, nil
return m.controller.IngestShard(ctx, qtid, shrdNum)
}
// SnapshotTable handles a snapshot table request.
func (m *MDS) SnapshotTable(ctx context.Context, qtid dax.QualifiedTableID) error {
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
return m.controller.SnapshotTable(ctx, qtid)
}
// SnapshotShardData handles a snapshot shard request.
func (m *MDS) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) error {
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
return m.controller.SnapshotShardData(ctx, qtid, shardNum)
}
// SnapshotTableKeys handles a snapshot table/keys request.
func (m *MDS) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) error {
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
return m.controller.SnapshotTableKeys(ctx, qtid, partitionNum)
}
// SnapshotFieldKeys handles a snapshot field/keys request.
func (m *MDS) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return errors.Wrap(err, "sanitizing")
}
return m.controller.SnapshotFieldKeys(ctx, qtid, fldName)
}
@ -466,11 +279,7 @@ func (m *MDS) DeregisterNodes(ctx context.Context, addrs ...dax.Address) error {
// ComputeNodes gets the compute nodes responsible for the table/shards
// specified in the ComputeNodeRequest.
func (m *MDS) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shardNums ...dax.ShardNum) ([]dax.ComputeNode, error) {
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return nil, errors.Wrap(err, "sanitizing")
}
return m.controller.ComputeNodes(ctx, qtid, shardNums, false)
return m.controller.ComputeNodes(ctx, qtid, shardNums)
}
func (m *MDS) DebugNodes(ctx context.Context) ([]*dax.Node, error) {
@ -480,9 +289,5 @@ func (m *MDS) DebugNodes(ctx context.Context) ([]*dax.Node, error) {
// TranslateNodes gets the translate nodes responsible for the table/partitions
// specified in the TranslateNodeRequest.
func (m *MDS) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitionNums ...dax.PartitionNum) ([]dax.TranslateNode, error) {
if err := m.sanitizeQTID(ctx, &qtid); err != nil {
return nil, errors.Wrap(err, "sanitizing")
}
return m.controller.TranslateNodes(ctx, qtid, partitionNums, false)
return m.controller.TranslateNodes(ctx, qtid, partitionNums)
}

392
dax/mds/schemar/boltdb/schemar.go
View file

@ -4,7 +4,6 @@ package boltdb
import (
"bytes"
"context"
"encoding/json"
"fmt"
"strings"
@ -44,9 +43,224 @@ func NewSchemar(db *boltdb.DB, logger logger.Logger) *Schemar {
}
}
// CreateDatabase creates the database provided. If a database with the same
// name already exists then an error is returned. For now, we are not going to
// store the tables in the schemar Database struct.
func (s *Schemar) CreateDatabase(tx dax.Transaction, qdb *dax.QualifiedDatabase) error {
// Ensure the database id is not blank.
if qdb.ID == "" {
return schemar.NewErrDatabaseIDInvalid(qdb.ID)
}
// Ensure the database name is not blank.
if qdb.Name == "" {
return schemar.NewErrDatabaseNameInvalid(qdb.Name)
}
// Set the CreateAt value for the database.
// TODO(tlt): We may want to consider erroring here if the value is != 0.
if qdb.CreatedAt == 0 {
now := timestamp()
qdb.CreatedAt = now
}
//////////// end validation
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
// Ensure a database with that ID doesn't already exist.
if db, _ := s.databaseByID(txx, qdb.OrganizationID, qdb.ID); db != nil {
return dax.NewErrDatabaseIDExists(qdb.QualifiedID())
}
if err := s.putDatabase(txx, qdb); err != nil {
return errors.Wrap(err, "putting database")
}
// In addition to storing the database in databaseKey, we want to store a
// reverse-lookup (i.e. index) on database name to the databaseKey.
if err := s.putDatabaseName(txx, qdb); err != nil {
return errors.Wrap(err, "putting database name")
}
return nil
}
func (s *Schemar) DatabaseByID(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
return s.databaseByID(txx, qdbid.OrganizationID, qdbid.DatabaseID)
}
func (s *Schemar) databaseByID(tx *boltdb.Tx, orgID dax.OrganizationID, id dax.DatabaseID) (*dax.QualifiedDatabase, error) {
bkt := tx.Bucket(bucketSchemar)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
b := bkt.Get(databaseKey(orgID, id))
if b == nil {
return nil, dax.NewErrDatabaseIDDoesNotExist(dax.QualifiedDatabaseID{OrganizationID: orgID, DatabaseID: id})
}
database := &dax.QualifiedDatabase{}
if err := json.Unmarshal(b, database); err != nil {
return nil, errors.Wrap(err, "unmarshalling database json")
}
return database, nil
}
func (s *Schemar) putDatabase(tx *boltdb.Tx, qdb *dax.QualifiedDatabase) error {
bkt := tx.Bucket(bucketSchemar)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
val, err := json.Marshal(qdb)
if err != nil {
return errors.Wrap(err, "marshalling database to json")
}
return bkt.Put(databaseKey(qdb.OrganizationID, qdb.ID), val)
}
func (s *Schemar) putDatabaseName(tx *boltdb.Tx, qdb *dax.QualifiedDatabase) error {
bkt := tx.Bucket(bucketSchemar)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
return bkt.Put(databaseNameKey(qdb.OrganizationID, qdb.Name), databaseKey(qdb.OrganizationID, qdb.ID))
}
// DropDatabase drops the given database. If the named/IDed database does not
// exist then an error is returned.
func (s *Schemar) DropDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
// Ensure the database exists.
qdb, err := s.databaseByID(txx, qdbid.OrganizationID, qdbid.DatabaseID)
if err != nil {
return errors.Wrap(err, "getting database by id")
}
bkt := txx.Bucket(bucketSchemar)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
// Delete the database by ID.
if err := bkt.Delete(databaseKey(qdb.OrganizationID, qdb.ID)); err != nil {
return errors.Wrap(err, "deleting database by id")
}
// Delete the reverse-lookup database by Name.
if err := bkt.Delete(databaseNameKey(qdb.OrganizationID, qdb.Name)); err != nil {
return errors.Wrap(err, "deleting database by name")
}
return nil
}
// SetDatabaseOptions overwrites the existing database options with those
// provided for the given database.
func (s *Schemar) SetDatabaseOptions(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, opts dax.DatabaseOptions) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
// Get the database.
qdb, err := s.databaseByID(txx, qdbid.OrganizationID, qdbid.DatabaseID)
if err != nil {
return errors.Wrapf(err, "getting database: %s", qdbid)
}
// Set the new options.
qdb.Options = opts
// Put the database.
if err := s.putDatabase(txx, qdb); err != nil {
return errors.Wrap(err, "putting database")
}
return nil
}
func (s *Schemar) Databases(tx dax.Transaction, orgID dax.OrganizationID, ids ...dax.DatabaseID) ([]*dax.QualifiedDatabase, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
return s.getDatabases(txx, orgID, ids...)
}
func (s *Schemar) getDatabases(tx *boltdb.Tx, orgID dax.OrganizationID, ids ...dax.DatabaseID) (dax.QualifiedDatabases, error) {
c := tx.Bucket(bucketSchemar).Cursor()
// Deserialize rows into Database objects.
databases := make(dax.QualifiedDatabases, 0)
var filterByID bool
if len(ids) > 0 {
filterByID = true
}
prefix := []byte(fmt.Sprintf(prefixFmtDatabases, orgID))
if orgID == "" {
prefix = []byte(prefixDatabases)
}
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
if v == nil {
s.logger.Printf("nil value for key: %s", k)
continue
}
dbID, err := keyDatabaseID(k)
if err != nil {
return nil, errors.Wrap(err, "getting database from key")
}
// Only include databases provided in the ids filter.
if filterByID && !containsDatabaseID(ids, dbID) {
continue
}
database := &dax.QualifiedDatabase{}
if err := json.Unmarshal(v, database); err != nil {
return nil, errors.Wrap(err, "unmarshalling database json")
}
databases = append(databases, database)
}
return databases, nil
}
func containsDatabaseID(s []dax.DatabaseID, e dax.DatabaseID) bool {
for _, a := range s {
if a == e {
return true
}
}
return false
}
// CreateTable creates the table provided. If a table with the same name already
// exists then an error is returned.
func (s *Schemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error {
func (s *Schemar) CreateTable(tx dax.Transaction, qtbl *dax.QualifiedTable) error {
// Ensure the table id is not blank.
if qtbl.ID == "" {
return schemar.NewErrTableIDInvalid(qtbl.ID)
@ -76,33 +290,37 @@ func (s *Schemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) err
//////////// end validation
tx, err := s.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "getting transaction")
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
// Ensure the database, defined in the table's QualifiedDatabaseID, exists.
if _, err := s.databaseByID(txx, qtbl.OrganizationID, qtbl.DatabaseID); err != nil {
return errors.Wrap(err, "validating database")
}
defer tx.Rollback()
// Ensure a table with that ID doesn't already exist.
if t, _ := s.tableByID(tx, qtbl.TableQualifier, qtbl.ID); t != nil {
if t, _ := s.tableByID(txx, qtbl.QualifiedDatabaseID, qtbl.ID); t != nil {
return dax.NewErrTableIDExists(qtbl.QualifiedID())
}
if err := s.putTable(tx, qtbl); err != nil {
if err := s.putTable(txx, qtbl); err != nil {
return errors.Wrap(err, "putting table")
}
// In addition to storing the table in tableKey, we want to store a reverse-lookup
// (i.e. index) on table name to the tableKey.
if err := s.putTableName(tx, qtbl); err != nil {
if err := s.putTableName(txx, qtbl); err != nil {
return errors.Wrap(err, "putting table name")
}
return tx.Commit()
return nil
}
// CreateField creates the field provided in the given table. If a field with
// the same name already exists then an error is returned.
func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error {
func (s *Schemar) CreateField(tx dax.Transaction, qtid dax.QualifiedTableID, fld *dax.Field) error {
// Ensure the field name is not blank.
if fld.Name == "" {
return schemar.NewErrFieldNameInvalid(fld.Name)
@ -110,14 +328,13 @@ func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fl
//////////// end validation
tx, err := s.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "getting transaction")
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
// Get the table.
qtbl, err := s.tableByQTID(tx, qtid)
qtbl, err := s.tableByQTID(txx, qtid)
if err != nil {
return errors.Wrap(err, "getting table by id")
}
@ -130,23 +347,22 @@ func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fl
qtbl.Fields = append(qtbl.Fields, fld)
// Write table back to database.
if err := s.putTable(tx, qtbl); err != nil {
if err := s.putTable(txx, qtbl); err != nil {
return errors.Wrap(err, "putting table")
}
return tx.Commit()
return nil
}
// DropField removes the field from the table.
func (s *Schemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
tx, err := s.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "getting transaction")
func (s *Schemar) DropField(tx dax.Transaction, qtid dax.QualifiedTableID, fldName dax.FieldName) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
// Get the table.
qtbl, err := s.tableByQTID(tx, qtid)
qtbl, err := s.tableByQTID(txx, qtid)
if err != nil {
return errors.Wrap(err, "getting table by id")
}
@ -159,11 +375,11 @@ func (s *Schemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldN
_ = qtbl.RemoveField(fldName)
// Write table back to database.
if err := s.putTable(tx, qtbl); err != nil {
if err := s.putTable(txx, qtbl); err != nil {
return errors.Wrap(err, "putting table")
}
return tx.Commit()
return nil
}
func (s *Schemar) putTable(tx *boltdb.Tx, qtbl *dax.QualifiedTable) error {
@ -191,43 +407,43 @@ func (s *Schemar) putTableName(tx *boltdb.Tx, qtbl *dax.QualifiedTable) error {
// Table returns the TableInfo for the given table. An error is returned if the
// table does not exist.
func (s *Schemar) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
tx, err := s.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
func (s *Schemar) Table(tx dax.Transaction, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
return s.tableByQTID(tx, qtid)
return s.tableByQTID(txx, qtid)
}
// tableByQTID gets the full qualified table by the QualifiedTableID whether it has Name or ID set.
// tableByQTID gets the full qualified table by the QualifiedTableID whether it
// has Name or ID set.
func (s *Schemar) tableByQTID(tx *boltdb.Tx, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
if qtid.ID == "" {
return s.tableByName(tx, qtid.TableQualifier, qtid.Name)
return s.tableByName(tx, qtid.QualifiedDatabaseID, qtid.Name)
}
return s.tableByID(tx, qtid.TableQualifier, qtid.ID)
return s.tableByID(tx, qtid.QualifiedDatabaseID, qtid.ID)
}
func (s *Schemar) tableByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax.TableName) (*dax.QualifiedTable, error) {
qtid, err := s.tableIDByName(tx, qual, name)
func (s *Schemar) tableByName(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, name dax.TableName) (*dax.QualifiedTable, error) {
qtid, err := s.tableIDByName(tx, qdbid, name)
if err != nil {
return nil, errors.Wrap(err, "getting table ID")
}
return s.tableByID(tx, qtid.TableQualifier, qtid.ID) // TODO remove?
return s.tableByID(tx, qtid.QualifiedDatabaseID, qtid.ID) // TODO remove?
}
func (s *Schemar) tableByID(tx *boltdb.Tx, qual dax.TableQualifier, id dax.TableID) (*dax.QualifiedTable, error) {
func (s *Schemar) tableByID(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, id dax.TableID) (*dax.QualifiedTable, error) {
bkt := tx.Bucket(bucketSchemar)
if bkt == nil {
return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
b := bkt.Get(tableKey(qual.OrganizationID, qual.DatabaseID, id))
b := bkt.Get(tableKey(qdbid.OrganizationID, qdbid.DatabaseID, id))
if b == nil {
return nil, dax.NewErrTableIDDoesNotExist(dax.QualifiedTableID{TableQualifier: qual, ID: id})
return nil, dax.NewErrTableIDDoesNotExist(dax.QualifiedTableID{QualifiedDatabaseID: qdbid, ID: id})
}
table := &dax.QualifiedTable{}
@ -238,13 +454,13 @@ func (s *Schemar) tableByID(tx *boltdb.Tx, qual dax.TableQualifier, id dax.Table
return table, nil
}
func (s *Schemar) tableIDByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
func (s *Schemar) tableIDByName(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
bkt := tx.Bucket(bucketSchemar)
if bkt == nil {
return dax.QualifiedTableID{}, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
b := bkt.Get(tableNameKey(qual.OrganizationID, qual.DatabaseID, name))
b := bkt.Get(tableNameKey(qdbid.OrganizationID, qdbid.DatabaseID, name))
if b == nil {
return dax.QualifiedTableID{}, dax.NewErrTableNameDoesNotExist(name)
}
@ -254,17 +470,16 @@ func (s *Schemar) tableIDByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax
// Tables returns a list of Table for all existing tables. If one or more table
// IDs is provided, then only those will be included in the output.
func (s *Schemar) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
tx, err := s.db.BeginTx(ctx, false)
if err != nil {
return nil, errors.Wrap(err, "beginning tx")
func (s *Schemar) Tables(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return nil, dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
return s.getTables(ctx, tx, qual, ids...)
return s.getTables(txx, qdbid, ids...)
}
func (s *Schemar) getTables(ctx context.Context, tx *boltdb.Tx, qual dax.TableQualifier, ids ...dax.TableID) (dax.QualifiedTables, error) {
func (s *Schemar) getTables(tx *boltdb.Tx, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) (dax.QualifiedTables, error) {
c := tx.Bucket(bucketSchemar).Cursor()
// Deserialize rows into Table objects.
@ -275,11 +490,11 @@ func (s *Schemar) getTables(ctx context.Context, tx *boltdb.Tx, qual dax.TableQu
filterByID = true
}
prefix := []byte(fmt.Sprintf(prefixFmtTables, qual.OrganizationID, qual.DatabaseID))
if qual.OrganizationID == "" && qual.DatabaseID == "" {
prefix := []byte(fmt.Sprintf(prefixFmtTables, qdbid.OrganizationID, qdbid.DatabaseID))
if qdbid.OrganizationID == "" && qdbid.DatabaseID == "" {
prefix = []byte(prefixTables)
} else if qual.DatabaseID == "" {
prefix = []byte(fmt.Sprintf(prefixFmtTablesOrg, qual.OrganizationID))
} else if qdbid.DatabaseID == "" {
prefix = []byte(fmt.Sprintf(prefixFmtTablesOrg, qdbid.OrganizationID))
}
for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() {
@ -320,20 +535,19 @@ func containsTableID(s []dax.TableID, e dax.TableID) bool {
// DropTable drops the given table. If the named/IDed table does not exist
// then an error is returned.
func (s *Schemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error {
tx, err := s.db.BeginTx(ctx, true)
if err != nil {
return errors.Wrap(err, "getting transaction")
func (s *Schemar) DropTable(tx dax.Transaction, qtid dax.QualifiedTableID) error {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.NewErrInvalidTransaction()
}
defer tx.Rollback()
// Ensure the table exists.
qtbl, err := s.tableByQTID(tx, qtid)
qtbl, err := s.tableByQTID(txx, qtid)
if err != nil {
return errors.Wrap(err, "getting table by id")
}
bkt := tx.Bucket(bucketSchemar)
bkt := txx.Bucket(bucketSchemar)
if bkt == nil {
return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar)
}
@ -348,16 +562,52 @@ func (s *Schemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) erro
return errors.Wrap(err, "deleting table by name")
}
return tx.Commit()
return nil
}
func (s *Schemar) TableID(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) {
txx, ok := tx.(*boltdb.Tx)
if !ok {
return dax.QualifiedTableID{}, dax.NewErrInvalidTransaction()
}
return s.tableIDByName(txx, qdbid, name)
}
const (
prefixTables = "tables/"
prefixFmtTablesOrg = prefixTables + "%s/"
prefixFmtTables = prefixFmtTablesOrg + "%s/"
prefixFmtTableNames = "tablenames/%s/%s/"
prefixFmtTablesOrg = prefixTables + "%s/" // org-id
prefixFmtTables = prefixFmtTablesOrg + "%s/" // db-id
prefixFmtTableNames = "tablenames/%s/%s/" // org-id, db-id
prefixDatabases = "databases/"
prefixFmtDatabases = prefixDatabases + "%s/" // org-id
prefixFmtDatabase = prefixFmtDatabases + "%s" // db-id
prefixFmtDatabaseNames = "databasenames/%s/" // org-id
)
// databaseKey returns a key based on a qualified database ID.
func databaseKey(orgID dax.OrganizationID, dbID dax.DatabaseID) []byte {
key := fmt.Sprintf(prefixFmtDatabase, orgID, dbID)
return []byte(key)
}
// databaseNameKey returns a key based on a qualified database name.
func databaseNameKey(orgID dax.OrganizationID, name dax.DatabaseName) []byte {
key := fmt.Sprintf(prefixFmtDatabaseNames+"%s", orgID, name)
return []byte(key)
}
// keyDatabaseID gets the DatabaseID out of the key.
func keyDatabaseID(key []byte) (dax.DatabaseID, error) {
parts := strings.Split(string(key), "/")
if len(parts) != 3 {
return "", errors.New(errors.ErrUncoded, "database key format expected: `databases/orgID/dbID`")
}
return dax.DatabaseID(parts[2]), nil
}
// tableKey returns a key based on a qualified table ID.
func tableKey(orgID dax.OrganizationID, dbID dax.DatabaseID, tblID dax.TableID) []byte {
key := fmt.Sprintf(prefixFmtTables+"%s", orgID, dbID, tblID)
@ -388,7 +638,7 @@ func keyQualifiedTableID(key []byte) (dax.QualifiedTableID, error) {
}
return dax.NewQualifiedTableID(
dax.NewTableQualifier(
dax.NewQualifiedDatabaseID(
dax.OrganizationID(parts[1]),
dax.DatabaseID(parts[2]),
),
@ -396,16 +646,6 @@ func keyQualifiedTableID(key []byte) (dax.QualifiedTableID, error) {
), nil
}
func (s *Schemar) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) {
tx, err := s.db.BeginTx(ctx, false)
if err != nil {
return dax.QualifiedTableID{}, err
}
defer tx.Rollback()
return s.tableIDByName(tx, qual, name)
}
func timestamp() int64 {
return time.Now().UnixNano()
}

129
dax/mds/schemar/boltdb/schemar_test.go
View file

@ -16,6 +16,7 @@ import (
func TestSchemar(t *testing.T) {
orgID := dax.OrganizationID("acme")
dbID := dax.DatabaseID("db1")
dbName := dax.DatabaseName("dbname1")
invalidTableID := dax.TableID("invalidID")
tableName := dax.TableName("foo")
tableName0 := dax.TableName("foo")
@ -25,21 +26,36 @@ func TestSchemar(t *testing.T) {
partitionN := 12
ctx := context.Background()
qual := dax.NewTableQualifier(orgID, dbID)
qdbid := dax.NewQualifiedDatabaseID(orgID, dbID)
db := testbolt.MustOpenDB(t)
defer testbolt.MustCloseDB(t, db)
t.Cleanup(func() {
testbolt.CleanupDB(t, db.Path())
})
// Initialize the buckets.
assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
qdb := &dax.QualifiedDatabase{
OrganizationID: orgID,
Database: dax.Database{
ID: dbID,
Name: dbName,
},
}
t.Run("NewSchemar", func(t *testing.T) {
db := testbolt.MustOpenDB(t)
defer testbolt.MustCloseDB(t, db)
t.Cleanup(func() {
testbolt.CleanupDB(t, db.Path())
})
// Initialize the buckets.
assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
s := boltdb.NewSchemar(db, logger.NopLogger)
tx, err := db.BeginTx(ctx, true)
assert.NoError(t, err)
defer tx.Rollback()
// Create database.
assert.NoError(t, s.CreateDatabase(tx, qdb))
// Add new table.
tbl := dax.NewTable(tableName)
tbl.CreateID()
@ -53,11 +69,11 @@ func TestSchemar(t *testing.T) {
Type: dax.BaseTypeInt,
},
}
qtbl := dax.NewQualifiedTable(qual, tbl)
assert.NoError(t, s.CreateTable(ctx, qtbl))
qtbl := dax.NewQualifiedTable(qdbid, tbl)
assert.NoError(t, s.CreateTable(tx, qtbl))
// Try adding the table again.
err := s.CreateTable(ctx, qtbl)
err = s.CreateTable(tx, qtbl)
if assert.Error(t, err) {
assert.True(t, errors.Is(err, dax.ErrTableIDExists))
}
@ -66,32 +82,32 @@ func TestSchemar(t *testing.T) {
// Get the table.
{
tbl, err := s.Table(ctx, qtid)
tbl, err := s.Table(tx, qtid)
assert.NoError(t, err)
assert.Equal(t, tableName, tbl.Name)
}
// Drop the table.
assert.NoError(t, s.DropTable(ctx, qtid))
assert.NoError(t, s.DropTable(tx, qtid))
// Make sure the reverse-lookup (table by name) was dropped as well.
{
_, err := s.TableID(ctx, qual, tableName)
_, err := s.TableID(tx, qdbid, tableName)
if assert.Error(t, err) {
assert.True(t, errors.Is(err, dax.ErrTableNameDoesNotExist))
}
}
// Try adding the table (i.e. the same table name) again.
assert.NoError(t, s.CreateTable(ctx, qtbl))
assert.NoError(t, s.CreateTable(tx, qtbl))
// Drop the table again.
assert.NoError(t, s.DropTable(ctx, qtid))
assert.NoError(t, s.DropTable(tx, qtid))
// Drop invalid table.
{
iqtid := dax.NewQualifiedTableID(qual, invalidTableID)
err := s.DropTable(ctx, iqtid)
iqtid := dax.NewQualifiedTableID(qdbid, invalidTableID)
err := s.DropTable(tx, iqtid)
if assert.Error(t, err) {
assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist))
}
@ -99,19 +115,36 @@ func TestSchemar(t *testing.T) {
})
t.Run("GetTables", func(t *testing.T) {
db := testbolt.MustOpenDB(t)
defer testbolt.MustCloseDB(t, db)
t.Cleanup(func() {
testbolt.CleanupDB(t, db.Path())
})
// Initialize the buckets.
assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
s := boltdb.NewSchemar(db, logger.NopLogger)
tx, err := db.BeginTx(ctx, true)
assert.NoError(t, err)
defer tx.Rollback()
// Create database.
assert.NoError(t, s.CreateDatabase(tx, qdb))
exp := []*dax.QualifiedTable{}
tables, err := s.Tables(ctx, qual)
tables, err := s.Tables(tx, qdbid)
assert.NoError(t, err)
assert.Equal(t, exp, tables)
qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false)
qtbl1 := daxtest.TestQualifiedTableWithID(t, qual, tableID1, tableName1, partitionN, false)
qtbl0 := daxtest.TestQualifiedTableWithID(t, qdbid, tableID0, tableName0, partitionN, false)
qtbl1 := daxtest.TestQualifiedTableWithID(t, qdbid, tableID1, tableName1, partitionN, false)
// Add a couple of tables.
assert.NoError(t, s.CreateTable(ctx, qtbl0))
assert.NoError(t, s.CreateTable(ctx, qtbl1))
assert.NoError(t, s.CreateTable(tx, qtbl0))
assert.NoError(t, s.CreateTable(tx, qtbl1))
exp = []*dax.QualifiedTable{
qtbl1,
@ -119,62 +152,76 @@ func TestSchemar(t *testing.T) {
}
// All tables.
tables, err = s.Tables(ctx, qual)
tables, err = s.Tables(tx, qdbid)
assert.NoError(t, err)
assert.Equal(t, exp, tables)
// With a valid filter.
tables, err = s.Tables(ctx, qual, qtbl0.ID)
tables, err = s.Tables(tx, qdbid, qtbl0.ID)
assert.NoError(t, err)
assert.Equal(t, exp[1:], tables)
// With an invalid filter.
tables, err = s.Tables(ctx, qual, invalidTableID)
tables, err = s.Tables(tx, qdbid, invalidTableID)
assert.NoError(t, err)
assert.Equal(t, exp[0:0], tables)
// With both valid and invalid filters.
tables, err = s.Tables(ctx, qual, qtbl0.ID, invalidTableID)
tables, err = s.Tables(tx, qdbid, qtbl0.ID, invalidTableID)
assert.NoError(t, err)
assert.Equal(t, exp[1:], tables)
// With all valid filters.
tables, err = s.Tables(ctx, qual, qtbl0.ID, qtbl1.ID)
tables, err = s.Tables(tx, qdbid, qtbl0.ID, qtbl1.ID)
assert.NoError(t, err)
assert.Equal(t, exp, tables)
})
t.Run("GetTablesAll", func(t *testing.T) {
// get a fresh DB
db := testbolt.MustOpenDB(t)
defer testbolt.MustCloseDB(t, db)
t.Cleanup(func() {
testbolt.CleanupDB(t, db.Path())
})
// Initialize the buckets.
assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...))
s := boltdb.NewSchemar(db, logger.NopLogger)
qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false)
orgID2 := dax.OrganizationID("acme2")
qual2 := dax.NewTableQualifier(orgID2, dbID)
tableID2 := "3"
qtbl2 := daxtest.TestQualifiedTableWithID(t, qual2, tableID2, dax.TableName("two"), partitionN, false)
tx, err := db.BeginTx(ctx, true)
assert.NoError(t, err)
defer tx.Rollback()
assert.NoError(t, s.CreateTable(ctx, qtbl0))
assert.NoError(t, s.CreateTable(ctx, qtbl2))
qtbl0 := daxtest.TestQualifiedTableWithID(t, qdbid, tableID0, tableName0, partitionN, false)
orgID2 := dax.OrganizationID("acme2")
qdbid2 := dax.NewQualifiedDatabaseID(orgID2, dbID)
tableID2 := "3"
qtbl2 := daxtest.TestQualifiedTableWithID(t, qdbid2, tableID2, dax.TableName("two"), partitionN, false)
// Create databases.
assert.NoError(t, s.CreateDatabase(tx, qdb))
qdb2 := &dax.QualifiedDatabase{
OrganizationID: orgID2,
Database: dax.Database{
ID: dbID,
Name: dbName,
},
}
assert.NoError(t, s.CreateDatabase(tx, qdb2))
assert.NoError(t, s.CreateTable(tx, qtbl0))
assert.NoError(t, s.CreateTable(tx, qtbl2))
exp := []*dax.QualifiedTable{qtbl0, qtbl2}
tables, err := s.Tables(ctx, dax.TableQualifier{})
tables, err := s.Tables(tx, dax.QualifiedDatabaseID{})
assert.NoError(t, err)
assert.Equal(t, exp, tables)
tables, err = s.Tables(ctx, dax.TableQualifier{OrganizationID: orgID2})
tables, err = s.Tables(tx, dax.QualifiedDatabaseID{OrganizationID: orgID2})
assert.NoError(t, err)
assert.Equal(t, []*dax.QualifiedTable{qtbl2}, tables)
})
}

17
dax/mds/schemar/errors.go
View file

@ -8,6 +8,9 @@ import (
)
const (
ErrCodeDatabaseIDInvalid errors.Code = "DatabaseIDInvalid"
ErrCodeDatabaseNameInvalid errors.Code = "DatabaseNameInvalid"
ErrCodeTableIDInvalid errors.Code = "TableIDInvalid"
ErrCodeTableNameInvalid errors.Code = "TableNameInvalid"
ErrCodeInvalidPrimaryKey errors.Code = "InvalidPrimaryKey"
@ -15,6 +18,20 @@ const (
ErrCodeFieldNameInvalid errors.Code = "FieldNameInvalid"
)
func NewErrDatabaseIDInvalid(databaseID dax.DatabaseID) error {
return errors.New(
ErrCodeDatabaseIDInvalid,
fmt.Sprintf("database ID '%s' is invalid", databaseID),
)
}
func NewErrDatabaseNameInvalid(databaseName dax.DatabaseName) error {
return errors.New(
ErrCodeDatabaseNameInvalid,
fmt.Sprintf("database name '%s' is invalid", databaseName),
)
}
func NewErrTableIDInvalid(tableID dax.TableID) error {
return errors.New(
ErrCodeTableIDInvalid,

199
dax/mds/schemar/http/handler.go
View file

@ -1,199 +0,0 @@
package http
import (
"encoding/json"
"net/http"
"github.com/gorilla/mux"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/dax/mds/schemar"
)
func Handler(s schemar.Schemar) http.Handler {
svr := &server{
schemar: s,
}
router := mux.NewRouter()
router.HandleFunc("/health", svr.getHealth).Methods("GET").Name("GetHealth")
router.HandleFunc("/create-table", svr.postCreateTable).Methods("POST").Name("PostCreateTable")
router.HandleFunc("/drop-table", svr.postDropTable).Methods("POST").Name("PostDropTable")
router.HandleFunc("/table", svr.postTable).Methods("POST").Name("PostTable")
router.HandleFunc("/tables", svr.postTables).Methods("POST").Name("PostTables")
return router
}
type server struct {
schemar schemar.Schemar
}
// GET /health
func (s *server) getHealth(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}
// POST /create-table
func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) {
body := r.Body
defer body.Close()
ctx := r.Context()
req := &dax.QualifiedTable{}
if err := json.NewDecoder(body).Decode(req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
err := s.schemar.CreateTable(ctx, req)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
resp := struct{}{}
if err := json.NewEncoder(w).Encode(resp); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
// POST /drop-table
func (s *server) postDropTable(w http.ResponseWriter, r *http.Request) {
body := r.Body
defer body.Close()
ctx := r.Context()
req := DropTableRequest{}
if err := json.NewDecoder(body).Decode(&req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
qtid := req.TableKey.QualifiedTableID()
err := s.schemar.DropTable(ctx, qtid)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
resp := struct{}{}
if err := json.NewEncoder(w).Encode(resp); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
// // POST /create-field
// func (s *server) postCreateField(w http.ResponseWriter, r *http.Request) {
// body := r.Body
// defer body.Close()
// req := mds.CreateFieldRequest{}
// if err := json.NewDecoder(body).Decode(&req); err != nil {
// http.Error(w, err.Error(), http.StatusBadRequest)
// return
// }
// resp, err := s.mds.CreateField(req)
// if err != nil {
// http.Error(w, err.Error(), http.StatusBadRequest)
// return
// }
// if err := json.NewEncoder(w).Encode(resp); err != nil {
// http.Error(w, err.Error(), http.StatusBadRequest)
// return
// }
// }
// // POST /drop-field
// func (s *server) postDropField(w http.ResponseWriter, r *http.Request) {
// body := r.Body
// defer body.Close()
// req := mds.DropFieldRequest{}
// if err := json.NewDecoder(body).Decode(&req); err != nil {
// http.Error(w, err.Error(), http.StatusBadRequest)
// return
// }
// resp, err := s.mds.DropField(req)
// if err != nil {
// http.Error(w, err.Error(), http.StatusBadRequest)
// return
// }
// if err := json.NewEncoder(w).Encode(resp); err != nil {
// http.Error(w, err.Error(), http.StatusBadRequest)
// return
// }
// }
// POST /table
func (s *server) postTable(w http.ResponseWriter, r *http.Request) {
body := r.Body
defer body.Close()
ctx := r.Context()
req := TableRequest{}
if err := json.NewDecoder(body).Decode(&req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
qtid := req.TableKey.QualifiedTableID()
resp, err := s.schemar.Table(ctx, qtid)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
if err := json.NewEncoder(w).Encode(resp); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
// POST /tables
func (s *server) postTables(w http.ResponseWriter, r *http.Request) {
body := r.Body
defer body.Close()
ctx := r.Context()
req := TablesRequest{}
if err := json.NewDecoder(body).Decode(&req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
resp, err := s.schemar.Tables(ctx, qual, req.TableIDs...)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
if err := json.NewEncoder(w).Encode(resp); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
}
type DropTableRequest struct {
TableKey dax.TableKey `json:"table-key"`
}
type TableRequest struct {
TableKey dax.TableKey `json:"table-key"`
}
type TablesRequest struct {
OrganizationID dax.OrganizationID `json:"org-id"`
DatabaseID dax.DatabaseID `json:"db-id"`
TableIDs dax.TableIDs `json:"table-ids"`
}

84
dax/mds/schemar/schemar.go
View file

@ -2,29 +2,38 @@
package schemar
import (
"context"
"github.com/featurebasedb/featurebase/v3/dax"
)
type Schemar interface {
CreateTable(context.Context, *dax.QualifiedTable) error
DropTable(context.Context, dax.QualifiedTableID) error
CreateField(context.Context, dax.QualifiedTableID, *dax.Field) error
DropField(context.Context, dax.QualifiedTableID, dax.FieldName) error
Table(context.Context, dax.QualifiedTableID) (*dax.QualifiedTable, error)
CreateDatabase(dax.Transaction, *dax.QualifiedDatabase) error
DropDatabase(dax.Transaction, dax.QualifiedDatabaseID) error
DatabaseByID(dax.Transaction, dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error)
// Tables returns a list of tables. If the qualifiers DatabaseID
// is empty, all tables in the org will be returned. If the
// OrganizationID is empty, all tables will be returned. If both
// are populated, only tables in that databse will be returned. If
// greater than zero table IDs are passed in the third argument,
// only tables matching those IDs will be returned.
Tables(context.Context, dax.TableQualifier, ...dax.TableID) ([]*dax.QualifiedTable, error)
SetDatabaseOptions(dax.Transaction, dax.QualifiedDatabaseID, dax.DatabaseOptions) error
// Databases returns a list of databases. If the OrganizationID is empty,
// all databases will be returned. If greater than zero database IDs are
// passed in the second argument, only databases matching those IDs will be
// returned.
Databases(dax.Transaction, dax.OrganizationID, ...dax.DatabaseID) ([]*dax.QualifiedDatabase, error)
CreateTable(dax.Transaction, *dax.QualifiedTable) error
DropTable(dax.Transaction, dax.QualifiedTableID) error
CreateField(dax.Transaction, dax.QualifiedTableID, *dax.Field) error
DropField(dax.Transaction, dax.QualifiedTableID, dax.FieldName) error
Table(dax.Transaction, dax.QualifiedTableID) (*dax.QualifiedTable, error)
// Tables returns a list of tables. If the qualifiers DatabaseID is empty,
// all tables in the org will be returned. If the OrganizationID is empty,
// all tables will be returned. If both are populated, only tables in that
// database will be returned. If greater than zero table IDs are passed in
// the third argument, only tables matching those IDs will be returned.
Tables(dax.Transaction, dax.QualifiedDatabaseID, ...dax.TableID) ([]*dax.QualifiedTable, error)
// TableID is a reverse-lookup method to get the TableID for a given
// qualified TableName.
TableID(context.Context, dax.TableQualifier, dax.TableName) (dax.QualifiedTableID, error)
TableID(dax.Transaction, dax.QualifiedDatabaseID, dax.TableName) (dax.QualifiedTableID, error)
}
//////////////////////////////////////////////
@ -39,23 +48,50 @@ func NewNopSchemar() *NopSchemar {
return &NopSchemar{}
}
func (s *NopSchemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { return nil }
func (s *NopSchemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error {
func (s *NopSchemar) CreateDatabase(tx dax.Transaction, qtbl *dax.QualifiedDatabase) error {
return nil
}
func (s *NopSchemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error {
func (s *NopSchemar) DropDatabase(tx dax.Transaction, qdbid dax.QualifiedDatabaseID) error {
return nil
}
func (s *NopSchemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fld dax.FieldName) error {
return nil
}
func (s *NopSchemar) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
func (s *NopSchemar) DatabaseByID(dax.Transaction, dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) {
return nil, nil
}
func (s *NopSchemar) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
func (s *NopSchemar) SetDatabaseOptions(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, opts dax.DatabaseOptions) error {
return nil
}
func (s *NopSchemar) Databases(dax.Transaction, dax.OrganizationID, ...dax.DatabaseID) ([]*dax.QualifiedDatabase, error) {
return nil, nil
}
func (s *NopSchemar) CreateTable(tx dax.Transaction, qtbl *dax.QualifiedTable) error {
return nil
}
func (s *NopSchemar) DropTable(tx dax.Transaction, qtid dax.QualifiedTableID) error {
return nil
}
func (s *NopSchemar) CreateField(tx dax.Transaction, qtid dax.QualifiedTableID, fld *dax.Field) error {
return nil
}
func (s *NopSchemar) DropField(tx dax.Transaction, qtid dax.QualifiedTableID, fld dax.FieldName) error {
return nil
}
func (s *NopSchemar) Table(tx dax.Transaction, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) {
return nil, nil
}
func (s *NopSchemar) Tables(tx dax.Transaction, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) {
return []*dax.QualifiedTable{}, nil
}
func (s *NopSchemar) TableID(context.Context, dax.TableQualifier, dax.TableName) (dax.QualifiedTableID, error) {
func (s *NopSchemar) TableID(dax.Transaction, dax.QualifiedDatabaseID, dax.TableName) (dax.QualifiedTableID, error) {
return dax.QualifiedTableID{}, nil
}

124
dax/mds/schemar/schemar_test.go
View file

@ -1,124 +0,0 @@
package schemar_test
import (
"context"
"testing"
"github.com/featurebasedb/featurebase/v3/dax"
daxtest "github.com/featurebasedb/featurebase/v3/dax/test"
"github.com/featurebasedb/featurebase/v3/errors"
"github.com/stretchr/testify/assert"
)
func TestSchemar(t *testing.T) {
orgID := dax.OrganizationID("acme")
dbID := dax.DatabaseID("db1")
invalidTableID := dax.TableID("invalidID")
tableName := dax.TableName("foo")
tableName0 := dax.TableName("foo")
tableName1 := dax.TableName("bar")
tableID0 := "2"
tableID1 := "1"
partitionN := 12
ctx := context.Background()
qual := dax.NewTableQualifier(orgID, dbID)
t.Run("NewSchemar", func(t *testing.T) {
s, cleanup := daxtest.NewSchemar(t)
defer cleanup()
// Add new table.
tbl := dax.NewTable(tableName)
tbl.Fields = []*dax.Field{
{
Name: dax.PrimaryKeyFieldName,
Type: dax.BaseTypeString,
},
{
Name: "intField",
Type: dax.BaseTypeInt,
},
}
qtbl := dax.NewQualifiedTable(qual, tbl)
qtbl.CreateID()
assert.NoError(t, s.CreateTable(ctx, qtbl))
// Try adding the table again.
err := s.CreateTable(ctx, qtbl)
if assert.Error(t, err) {
assert.True(t, errors.Is(err, dax.ErrTableIDExists))
}
qtid := qtbl.QualifiedID()
// Get the table.
{
tbl, err := s.Table(ctx, qtid)
assert.NoError(t, err)
assert.Equal(t, tableName, tbl.Name)
}
// Drop the table.
{
err := s.DropTable(ctx, qtid)
assert.NoError(t, err)
}
// Drop invalid table.
{
iqtid := dax.NewQualifiedTableID(qual, invalidTableID)
err := s.DropTable(ctx, iqtid)
if assert.Error(t, err) {
assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist))
}
}
})
t.Run("GetTables", func(t *testing.T) {
s, cleanup := daxtest.NewSchemar(t)
defer cleanup()
exp := []*dax.QualifiedTable{}
tables, err := s.Tables(ctx, qual)
assert.NoError(t, err)
assert.Equal(t, exp, tables)
qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false)
qtbl1 := daxtest.TestQualifiedTableWithID(t, qual, tableID1, tableName1, partitionN, false)
// Add a couple of tables.
assert.NoError(t, s.CreateTable(ctx, qtbl0))
assert.NoError(t, s.CreateTable(ctx, qtbl1))
exp = []*dax.QualifiedTable{
qtbl1,
qtbl0,
}
// All tables.
tables, err = s.Tables(ctx, qual)
assert.NoError(t, err)
assert.Equal(t, exp, tables)
// With a valid filter.
tables, err = s.Tables(ctx, qual, qtbl0.ID)
assert.NoError(t, err)
assert.Equal(t, exp[1:], tables)
// With an invalid filter.
tables, err = s.Tables(ctx, qual, invalidTableID)
assert.NoError(t, err)
assert.Equal(t, exp[0:0], tables)
// With both valid and invalid filters.
tables, err = s.Tables(ctx, qual, qtbl0.ID, invalidTableID)
assert.NoError(t, err)
assert.Equal(t, exp[1:], tables)
// With all valid filters.
tables, err = s.Tables(ctx, qual, qtbl0.ID, qtbl1.ID)
assert.NoError(t, err)
assert.Equal(t, exp, tables)
})
}

27
dax/node.go
View file

@ -3,6 +3,7 @@ package dax
import (
"context"
"fmt"
"strings"
"github.com/featurebasedb/featurebase/v3/errors"
)
@ -15,7 +16,23 @@ type Node struct {
RoleTypes []RoleType `json:"role-types"`
}
// AssignedNode is used in API responses.
// Nodes is a slice of *Node. It's useful for printing the nodes as a list of
// node.Addresses via its String() method.
type Nodes []*Node
// String prints the slice of node addresses in Nodes.
func (n Nodes) String() string {
out := make([]string, 0, len(n))
for i := range n {
out = append(out, string(n[i].Address))
}
return "[" + strings.Join(out, ",") + "]"
}
// AssignedNode represents a Worker which has been assigned a role. Note that
// the worker which it represents might be responsible for multiple roles, but
// AssignedNode only ever represents one of those roles at a time. This is
// because it is always the response of a RoleType-specific request.
type AssignedNode struct {
Address Address `json:"address"`
Role Role `json:"role"`
@ -23,10 +40,10 @@ type AssignedNode struct {
// NodeService represents a service for managing Nodes.
type NodeService interface {
CreateNode(context.Context, Address, *Node) error
ReadNode(context.Context, Address) (*Node, error)
DeleteNode(context.Context, Address) error
Nodes(context.Context) ([]*Node, error)
CreateNode(Transaction, Address, *Node) error
ReadNode(Transaction, Address) (*Node, error)
DeleteNode(Transaction, Address) error
Nodes(Transaction) ([]*Node, error)
}
// ComputeNode represents a compute node and the table/shards for which it is

13
dax/queryer/client/client.go
View file

@ -49,12 +49,12 @@ func (c *Client) Health() bool {
return true
}
func (c *Client) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
func (c *Client) QuerySQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, sql string) (*featurebase.WireQueryResponse, error) {
url := fmt.Sprintf("%s/sql", c.address.WithScheme(defaultScheme))
req := &queryerhttp.SQLRequest{
OrganizationID: qual.OrganizationID,
DatabaseID: qual.DatabaseID,
OrganizationID: qdbid.OrganizationID,
DatabaseID: qdbid.DatabaseID,
SQL: sql,
}
@ -85,12 +85,13 @@ func (c *Client) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql stri
return wireResp, nil
}
func (c *Client) QueryPQL(ctx context.Context, qual dax.TableQualifier, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
func (c *Client) QueryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
url := fmt.Sprintf("%s/query", c.address.WithScheme(defaultScheme))
req := &queryerhttp.QueryRequest{
OrganizationID: qual.OrganizationID,
DatabaseID: qual.DatabaseID,
OrganizationID: qdbid.OrganizationID,
DatabaseID: qdbid.DatabaseID,
Table: table,
PQL: pql,
}

10
dax/queryer/http/handler.go
View file

@ -49,15 +49,15 @@ func (s *server) postQuery(w http.ResponseWriter, r *http.Request) {
var resp interface{}
var err error
qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
qdbid := dax.NewQualifiedDatabaseID(req.OrganizationID, req.DatabaseID)
if req.SQL != "" {
resp, err = s.queryer.QuerySQL(ctx, qual, req.SQL)
resp, err = s.queryer.QuerySQL(ctx, qdbid, req.SQL)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
} else {
resp, err = s.queryer.QueryPQL(ctx, qual, req.Table, req.PQL)
resp, err = s.queryer.QueryPQL(ctx, qdbid, req.Table, req.PQL)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
@ -83,8 +83,8 @@ func (s *server) postSQL(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID)
resp, err := s.queryer.QuerySQL(ctx, qual, req.SQL)
qdbid := dax.NewQualifiedDatabaseID(req.OrganizationID, req.DatabaseID)
resp, err := s.queryer.QuerySQL(ctx, qdbid, req.SQL)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return

112
dax/queryer/orchestrator.go
View file

@ -11,10 +11,10 @@ import (
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/prometheus/client_golang/prometheus"
"github.com/featurebasedb/featurebase/v3/errors"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/tracing"
"golang.org/x/sync/errgroup"
)
@ -74,7 +74,7 @@ type Translator interface {
// TODO(jaffee) the naming here is a cluster. TranslateIndexIDs takes a list, but TranslateFieldIDs takes a set, both have alternate methods that take the other thing. :facepalm:
TranslateIndexIDs(ctx context.Context, index string, ids []uint64) ([]string, error)
TranslateIndexIDSet(ctx context.Context, index string, ids map[uint64]struct{}) (map[uint64]string, error)
TranslateFieldIDs(ctx context.Context, index, field string, ids map[uint64]struct{}) (map[uint64]string, error)
TranslateFieldIDs(ctx context.Context, tableKeyer dax.TableKeyer, field string, ids map[uint64]struct{}) (map[uint64]string, error)
TranslateFieldListIDs(ctx context.Context, index, field string, ids []uint64) ([]string, error)
}
@ -87,7 +87,6 @@ type orchestrator struct {
// Client used for remote requests.
client *featurebase.InternalClient
stats stats.StatsClient
logger logger.Logger
}
@ -433,11 +432,11 @@ func (o *orchestrator) executeCall(ctx context.Context, tableKeyer dax.TableKeye
} else if err := o.validateCallArgs(c); err != nil {
return nil, errors.Wrap(err, "validating args")
}
indexTag := "index:" + string(tableKeyer.Key())
metricName := "query_" + strings.ToLower(c.Name) + "_total"
statFn := func() {
labels := prometheus.Labels{"index": string(tableKeyer.Key())}
statFn := func(ctr *prometheus.CounterVec) {
if !opt.Remote {
o.stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag})
ctr.With(labels).Inc()
}
}
@ -449,101 +448,106 @@ func (o *orchestrator) executeCall(ctx context.Context, tableKeyer dax.TableKeye
switch c.Name {
case "Sum":
statFn()
statFn(featurebase.CounterQuerySumTotal)
res, err := o.executeSum(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeSum")
case "Min":
statFn()
statFn(featurebase.CounterQueryMinTotal)
res, err := o.executeMin(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeMin")
case "Max":
statFn()
statFn(featurebase.CounterQueryMaxTotal)
res, err := o.executeMax(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeMax")
case "MinRow":
statFn()
statFn(featurebase.CounterQueryMinRowTotal)
res, err := o.executeMinRow(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeMinRow")
case "MaxRow":
statFn()
statFn(featurebase.CounterQueryMaxRowTotal)
res, err := o.executeMaxRow(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeMaxRow")
// case "Clear":
// statFn()
// statFn(featurebase.CounterQueryClearTotal)
// res, err := o.executeClearBit(ctx, index, c, opt)
// return res, errors.Wrap(err, "executeClearBit")
// case "ClearRow":
// statFn()
// statFn(featurebase.CounterQueryClearRowTotal)
// res, err := o.executeClearRow(ctx, index, c, shards, opt)
// return res, errors.Wrap(err, "executeClearRow")
case "Distinct":
statFn()
statFn(featurebase.CounterQueryDistinctTotal)
res, err := o.executeDistinct(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeDistinct")
// case "Store":
// statFn()
// statFn(featurebase.CounterQueryStoreTotal)
// res, err := o.executeSetRow(ctx, index, c, shards, opt)
// return res, errors.Wrap(err, "executeSetRow")
case "Count":
statFn()
statFn(featurebase.CounterQueryCountTotal)
res, err := o.executeCount(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeCount")
// case "Set":
// statFn()
// statFn(featurebase.CounterQuerySetTotal)
// res, err := o.executeSet(ctx, index, c, opt)
// return res, errors.Wrap(err, "executeSet")
case "TopK":
statFn()
statFn(featurebase.CounterQueryTopKTotal)
res, err := o.executeTopK(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeTopK")
case "TopN":
statFn()
statFn(featurebase.CounterQueryTopNTotal)
res, err := o.executeTopN(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeTopN")
case "Rows":
statFn()
statFn(featurebase.CounterQueryRowsTotal)
res, err := o.executeRows(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeRows")
case "Extract":
statFn()
statFn(featurebase.CounterQueryExtractTotal)
res, err := o.executeExtract(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeExtract")
case "GroupBy":
statFn()
statFn(featurebase.CounterQueryGroupByTotal)
res, err := o.executeGroupBy(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeGroupBy")
case "Options":
statFn()
statFn(featurebase.CounterQueryOptionsTotal)
res, err := o.executeOptionsCall(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeOptionsCall")
case "IncludesColumn":
statFn(featurebase.CounterQueryIncludesColumnTotal)
res, err := o.executeIncludesColumnCall(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeIncludesColumnCall")
case "FieldValue":
statFn()
statFn(featurebase.CounterQueryFieldValueTotal)
res, err := o.executeFieldValueCall(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeFieldValueCall")
case "Precomputed":
statFn(featurebase.CounterQueryPrecomputedTotal)
res, err := o.executePrecomputedCall(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executePrecomputedCall")
case "UnionRows":
statFn(featurebase.CounterQueryUnionRowsTotal)
res, err := o.executeUnionRows(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeUnionRows")
case "ConstRow":
statFn(featurebase.CounterQueryConstRowTotal)
res, err := o.executeConstRow(ctx, tableKeyer, c)
return res, errors.Wrap(err, "executeConstRow")
case "Limit":
statFn(featurebase.CounterQueryLimitTotal)
res, err := o.executeLimitCall(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeLimitCall")
case "Percentile":
statFn(featurebase.CounterQueryPercentileTotal)
res, err := o.executePercentile(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executePercentile")
// case "Delete":
// statFn() //TODO(twg) need this?
// statFn(featurebase.CounterQueryDeleteTotal)
// res, err := o.executeDeleteRecords(ctx, index, c, shards, opt)
// return res, errors.Wrap(err, "executeDelete")
default: // o.g. "Row", "Union", "Intersect" or anything that returns a bitmap.
statFn()
res, err := o.executeBitmapCall(ctx, tableKeyer, c, shards, opt)
return res, errors.Wrap(err, "executeBitmapCall")
}
@ -1116,13 +1120,42 @@ func (o *orchestrator) executeBitmapCall(ctx context.Context, tableKeyer dax.Tab
span.LogKV("pqlCallName", c.Name)
defer span.Finish()
indexTag := "index:" + string(tableKeyer.Key())
metricName := "query_" + strings.ToLower(c.Name) + "_total"
if c.Name == "Row" && c.HasConditionArg() {
metricName = "query_row_bsi_total"
labels := prometheus.Labels{"index": string(tableKeyer.Key())}
statFn := func(ctr *prometheus.CounterVec) {
if !opt.Remote {
ctr.With(labels).Inc()
}
}
if !opt.Remote {
o.stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag})
switch c.Name {
case "Row":
if c.HasConditionArg() {
statFn(featurebase.CounterQueryRowBSITotal)
} else {
statFn(featurebase.CounterQueryRowTotal)
}
case "Range":
statFn(featurebase.CounterQueryRangeTotal)
case "Difference":
statFn(featurebase.CounterQueryBitmapTotal)
case "Intersect":
statFn(featurebase.CounterQueryIntersectTotal)
case "Union":
statFn(featurebase.CounterQueryUnionTotal)
case "InnerUnionRows":
statFn(featurebase.CounterQueryInnerUnionRowsTotal)
case "Xor":
statFn(featurebase.CounterQueryXorTotal)
case "Not":
statFn(featurebase.CounterQueryNotTotal)
case "Shift":
statFn(featurebase.CounterQueryShiftTotal)
case "All":
statFn(featurebase.CounterQueryAllTotal)
default:
statFn(featurebase.CounterQueryBitmapTotal)
}
}
// Merge returned results at coordinating node.
@ -1807,6 +1840,7 @@ func (o *orchestrator) executeExtract(ctx context.Context, tableKeyer dax.TableK
}
fields[i] = fieldName
}
// TODO(tlt): is `fields` used?
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
@ -2985,9 +3019,7 @@ func (o *orchestrator) preTranslateMatrixSet(ctx context.Context, mat featurebas
}
}
index := string(tableKeyer.Key())
return o.trans.TranslateFieldIDs(ctx, index, field, ids)
return o.trans.TranslateFieldIDs(ctx, tableKeyer, field, ids)
}
func (o *orchestrator) translateResult(ctx context.Context, qtbl *dax.QualifiedTable, call *pql.Call, result interface{}, idSet map[uint64]string) (_ interface{}, err error) {
@ -3164,7 +3196,7 @@ func (o *orchestrator) translateResult(ctx context.Context, qtbl *dax.QualifiedT
fieldTranslations := make(map[string]map[uint64]string)
for field, ids := range fieldIDs {
trans, err := o.trans.TranslateFieldIDs(ctx, idx.Name, field.Name, ids)
trans, err := o.trans.TranslateFieldIDs(ctx, qtbl, field.Name, ids)
if err != nil {
return nil, errors.Wrapf(err, "translating IDs in field '%q'", field.Name)
}
@ -3480,13 +3512,13 @@ func callArgString(call *pql.Call, key string) string {
type qualifiedOrchestrator struct {
*orchestrator
qual dax.TableQualifier
qdbid dax.QualifiedDatabaseID
}
func newQualifiedOrchestrator(orch *orchestrator, qual dax.TableQualifier) *qualifiedOrchestrator {
func newQualifiedOrchestrator(orch *orchestrator, qdbid dax.QualifiedDatabaseID) *qualifiedOrchestrator {
return &qualifiedOrchestrator{
orchestrator: orch,
qual: qual,
qdbid: qdbid,
}
}
@ -3497,7 +3529,7 @@ func (o *qualifiedOrchestrator) Execute(ctx context.Context, tableKeyer dax.Tabl
switch keyer := tableKeyer.(type) {
case *dax.Table:
qtbl = dax.NewQualifiedTable(o.qual, keyer)
qtbl = dax.NewQualifiedTable(o.qdbid, keyer)
case *dax.QualifiedTable:
qtbl = keyer
default:

94
dax/queryer/queryer.go
View file

@ -22,7 +22,6 @@ import (
"github.com/featurebasedb/featurebase/v3/sql3/parser"
"github.com/featurebasedb/featurebase/v3/sql3/planner"
plannertypes "github.com/featurebasedb/featurebase/v3/sql3/planner/types"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/systemlayer"
uuid "github.com/satori/go.uuid"
)
@ -32,7 +31,7 @@ import (
// of "Queryer" nodes, which handle incoming query requests.
type Queryer struct {
mu sync.RWMutex
orchestrators map[dax.TableQualifier]*qualifiedOrchestrator
orchestrators map[dax.QualifiedDatabaseID]*qualifiedOrchestrator
fbClient *featurebase.InternalClient
@ -47,7 +46,7 @@ func New(cfg Config) *Queryer {
q := &Queryer{
noder: dax.NewNopNoder(),
schemar: dax.NewNopSchemar(),
orchestrators: make(map[dax.TableQualifier]*qualifiedOrchestrator),
orchestrators: make(map[dax.QualifiedDatabaseID]*qualifiedOrchestrator),
logger: logger.NopLogger,
}
@ -59,13 +58,13 @@ func New(cfg Config) *Queryer {
}
// Orchestrator gets (or creates) an instance of qualifiedOrchestrator based on
// the provided dax.TableQualifier.
func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
// the provided dax.QualifiedDatabaseID.
func (q *Queryer) Orchestrator(qdbid dax.QualifiedDatabaseID) *qualifiedOrchestrator {
// Try to get orchestrator under a read lock first.
if orch := func() *qualifiedOrchestrator {
q.mu.RLock()
defer q.mu.RUnlock()
if orch, ok := q.orchestrators[qual]; ok {
if orch, ok := q.orchestrators[qdbid]; ok {
return orch
}
return nil
@ -77,11 +76,11 @@ func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
// lock and try a read/write.
q.mu.Lock()
defer q.mu.Unlock()
if orch, ok := q.orchestrators[qual]; ok {
if orch, ok := q.orchestrators[qdbid]; ok {
return orch
}
sapi := newQualifiedSchemaAPI(qual, q.schemar)
sapi := newQualifiedSchemaAPI(qdbid, q.schemar)
orch := &orchestrator{
schema: sapi,
@ -89,12 +88,11 @@ func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
topology: &MDSTopology{noder: q.noder},
// TODO(jaffee) using default http.Client probably bad... need to set some timeouts.
client: q.fbClient,
stats: stats.NopStatsClient,
logger: q.logger,
}
qorch := newQualifiedOrchestrator(orch, qual)
q.orchestrators[qual] = qorch
qorch := newQualifiedOrchestrator(orch, qdbid)
q.orchestrators[qdbid] = qorch
return qorch
}
@ -133,12 +131,9 @@ func (q *Queryer) Start() error {
return nil
}
func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
func (q *Queryer) QuerySQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, sql string) (*featurebase.WireQueryResponse, error) {
start := time.Now()
if len(sql) > 0 && sql[0] == '[' {
return q.parseAndQueryPQL(ctx, qual, sql)
}
ret := &featurebase.WireQueryResponse{}
applyExecutionTime := func() {
@ -150,6 +145,19 @@ func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql str
applyExecutionTime()
}
// If PQL, run that instead.
if len(sql) > 0 && sql[0] == '[' {
if pqlResp, err := q.parseAndQueryPQL(ctx, qdbid, sql); err != nil {
applyError(errors.Wrap(err, "querying pql"))
return ret, nil
} else {
ret = pqlResp
}
applyExecutionTime()
return ret, nil
}
// Create a requestID and add it to the context.
requestID, err := uuid.NewV4()
if err != nil {
@ -166,19 +174,17 @@ func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql str
}
// SchemaAPI
sapi := newQualifiedSchemaAPI(qual, q.schemar)
sapi := newQualifiedSchemaAPI(qdbid, q.schemar)
// Importer
imp := idkmds.NewImporter(q.noder, q.schemar, qual, nil)
imp := idkmds.NewImporter(q.noder, q.schemar, qdbid, nil)
// TODO(tlt): this obviously doesn't work; we don't have an API here. We
// need a dax-compatible implementation of the SystemAPI (or at least a
// no-op implementation).
sysapi := &featurebase.FeatureBaseSystemAPI{API: nil}
// TODO(tlt): We need a dax-compatible implementation of the SystemAPI.
sysapi := &featurebase.NopSystemAPI{}
systemLayer := systemlayer.NewSystemLayer()
pl := planner.NewExecutionPlanner(q.Orchestrator(qual), sapi, sysapi, systemLayer, imp, q.logger, sql)
pl := planner.NewExecutionPlanner(q.Orchestrator(qdbid), sapi, sysapi, systemLayer, imp, q.logger, sql)
planOp, err := pl.CompilePlan(ctx, st)
if err != nil {
@ -230,7 +236,7 @@ func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql str
return ret, nil
}
func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
func (q *Queryer) parseAndQueryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, sql string) (*featurebase.WireQueryResponse, error) {
var i int
for i = 1; sql[i] != ']'; i++ {
if i == len(sql)-1 {
@ -241,7 +247,7 @@ func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier,
query := sql[i+1:]
fmt.Println("got table/query", table, query)
return q.QueryPQL(ctx, qual, dax.TableName(table), query)
return q.queryPQL(ctx, qdbid, dax.TableName(table), query)
}
// convertIndex tries to covert any "index" specified in the call.Args map to a
@ -251,9 +257,9 @@ func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier,
// modify Call.CallIndex() to be TableKeyer aware. I didn't do that along with
// these changes because I'm not sure if we want to introduce dax types into the
// pql package.
func (q *Queryer) convertIndex(ctx context.Context, qual dax.TableQualifier, call *featurebase_pql.Call) {
func (q *Queryer) convertIndex(ctx context.Context, qdbid dax.QualifiedDatabaseID, call *featurebase_pql.Call) {
if index := call.CallIndex(); index != "" {
qtbl, err := q.schemar.TableByName(ctx, qual, dax.TableName(index))
qtbl, err := q.schemar.TableByName(ctx, qdbid, dax.TableName(index))
if err != nil {
return
}
@ -262,11 +268,37 @@ func (q *Queryer) convertIndex(ctx context.Context, qual dax.TableQualifier, cal
// Apply to children.
for _, child := range call.Children {
q.convertIndex(ctx, qual, child)
q.convertIndex(ctx, qdbid, child)
}
}
func (q *Queryer) QueryPQL(ctx context.Context, qual dax.TableQualifier, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
func (q *Queryer) QueryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
start := time.Now()
ret := &featurebase.WireQueryResponse{}
applyExecutionTime := func() {
ret.ExecutionTime = time.Since(start).Microseconds()
}
applyError := func(e error) {
ret.Error = e.Error()
applyExecutionTime()
}
if pqlResp, err := q.queryPQL(ctx, qdbid, table, pql); err != nil {
applyError(errors.Wrap(err, "querying pql"))
return ret, nil
} else {
ret = pqlResp
}
applyExecutionTime()
return ret, nil
}
func (q *Queryer) queryPQL(ctx context.Context, qdbid dax.QualifiedDatabaseID, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
// Parse the pql into a pql.Query containing []pql.Call.
qry, err := featurebase_pql.NewParser(strings.NewReader(pql)).Parse()
if err != nil {
@ -277,14 +309,14 @@ func (q *Queryer) QueryPQL(ctx context.Context, qual dax.TableQualifier, table d
}
// Replace any "index" arguments within the PQL with a TableKey.
q.convertIndex(ctx, qual, qry.Calls[0])
q.convertIndex(ctx, qdbid, qry.Calls[0])
qtbl, err := q.schemar.TableByName(ctx, qual, dax.TableName(table))
qtbl, err := q.schemar.TableByName(ctx, qdbid, dax.TableName(table))
if err != nil {
return nil, errors.Wrap(err, "converting index to qualified table")
}
results, err := q.Orchestrator(qual).Execute(ctx, qtbl, qry, nil, &featurebase.ExecOptions{})
results, err := q.Orchestrator(qdbid).Execute(ctx, qtbl, qry, nil, &featurebase.ExecOptions{})
if err != nil {
return nil, errors.Wrap(err, "orchestrator.Execute")
}

34
dax/queryer/schema_api.go
View file

@ -12,31 +12,31 @@ import (
var _ pilosa.SchemaAPI = (*qualifiedSchemaAPI)(nil)
// qualifiedSchemaAPI is a wrapper around schemaAPI. It is initialized with a
// TableQualifier, and it uses this qualifer to convert between, for example,
// FeatureBase index name (a string) and TableKey. It requires a Schemar to do
// that lookup/conversion.
// QualifiedDatabaseID, and it uses this qualifer to convert between, for
// example, FeatureBase index name (a string) and TableKey. It requires a
// Schemar to do that lookup/conversion.
type qualifiedSchemaAPI struct {
qual dax.TableQualifier
qdbid dax.QualifiedDatabaseID
schemar dax.Schemar
}
func newQualifiedSchemaAPI(qual dax.TableQualifier, schema dax.Schemar) *qualifiedSchemaAPI {
func newQualifiedSchemaAPI(qdbid dax.QualifiedDatabaseID, schema dax.Schemar) *qualifiedSchemaAPI {
return &qualifiedSchemaAPI{
qual: qual,
qdbid: qdbid,
schemar: schema,
}
}
func (s *qualifiedSchemaAPI) TableByName(ctx context.Context, tname dax.TableName) (*dax.Table, error) {
qtbl, err := s.schemar.TableByName(ctx, s.qual, tname)
qtbl, err := s.schemar.TableByName(ctx, s.qdbid, tname)
if err != nil {
return nil, errors.Wrapf(err, "getting table id: (%s) %s", s.qual, tname)
return nil, errors.Wrapf(err, "getting table id: (%s) %s", s.qdbid, tname)
}
return &qtbl.Table, nil
}
func (s *qualifiedSchemaAPI) TableByID(ctx context.Context, tid dax.TableID) (*dax.Table, error) {
qtid := dax.NewQualifiedTableID(s.qual, tid)
qtid := dax.NewQualifiedTableID(s.qdbid, tid)
qtbl, err := s.schemar.TableByID(ctx, qtid)
if err != nil {
@ -47,7 +47,7 @@ func (s *qualifiedSchemaAPI) TableByID(ctx context.Context, tid dax.TableID) (*d
}
func (s *qualifiedSchemaAPI) Tables(ctx context.Context) ([]*dax.Table, error) {
qtbls, err := s.schemar.Tables(ctx, s.qual)
qtbls, err := s.schemar.Tables(ctx, s.qdbid)
if err != nil {
return nil, errors.Wrap(err, "getting tables")
}
@ -61,32 +61,32 @@ func (s *qualifiedSchemaAPI) Tables(ctx context.Context) ([]*dax.Table, error) {
}
func (s *qualifiedSchemaAPI) CreateTable(ctx context.Context, tbl *dax.Table) error {
qtbl := dax.NewQualifiedTable(s.qual, tbl)
qtbl := dax.NewQualifiedTable(s.qdbid, tbl)
return s.schemar.CreateTable(ctx, qtbl)
}
func (s *qualifiedSchemaAPI) CreateField(ctx context.Context, tname dax.TableName, fld *dax.Field) error {
qtbl, err := s.schemar.TableByName(ctx, s.qual, tname)
qtbl, err := s.schemar.TableByName(ctx, s.qdbid, tname)
if err != nil {
return errors.Wrapf(err, "getting table by name: (%s) %s", s.qual, tname)
return errors.Wrapf(err, "getting table by name: (%s) %s", s.qdbid, tname)
}
return s.schemar.CreateField(ctx, qtbl.QualifiedID(), fld)
}
func (s *qualifiedSchemaAPI) DeleteTable(ctx context.Context, tname dax.TableName) error {
qtbl, err := s.schemar.TableByName(ctx, s.qual, tname)
qtbl, err := s.schemar.TableByName(ctx, s.qdbid, tname)
if err != nil {
return errors.Wrapf(err, "getting table by name: (%s) %s", s.qual, tname)
return errors.Wrapf(err, "getting table by name: (%s) %s", s.qdbid, tname)
}
return s.schemar.DropTable(ctx, qtbl.QualifiedID())
}
func (s *qualifiedSchemaAPI) DeleteField(ctx context.Context, tname dax.TableName, fname dax.FieldName) error {
qtid, err := s.schemar.TableByName(ctx, s.qual, tname)
qtid, err := s.schemar.TableByName(ctx, s.qdbid, tname)
if err != nil {
return errors.Wrapf(err, "getting table by name: (%s) %s", s.qual, tname)
return errors.Wrapf(err, "getting table by name: (%s) %s", s.qdbid, tname)
}
return s.schemar.DropField(ctx, qtid.Key().QualifiedTableID(), fname)

10
dax/queryer/translator.go
View file

@ -229,15 +229,19 @@ func (m *mdsTranslator) TranslateIndexIDSet(ctx context.Context, table string, i
}
return ret, nil
}
func (m *mdsTranslator) TranslateFieldIDs(ctx context.Context, table, field string, ids map[uint64]struct{}) (map[uint64]string, error) {
func (m *mdsTranslator) TranslateFieldIDs(ctx context.Context, tableKeyer dax.TableKeyer, field string, ids map[uint64]struct{}) (map[uint64]string, error) {
idList := make([]uint64, 0, len(ids))
for id := range ids {
idList = append(idList, id)
}
stringList, err := m.TranslateFieldListIDs(ctx, table, field, idList)
// TODO(tlt): convert TranslateFieldListIDs (and the other Translator
// interface methods) to TableKeyer.
index := string(tableKeyer.Key())
stringList, err := m.TranslateFieldListIDs(ctx, index, field, idList)
if err != nil {
return nil, errors.Wrapf(err, "translating field ids on field: %s, %s", table, field)
return nil, errors.Wrapf(err, "translating field ids on field: %s, %s", tableKeyer, field)
}
ret := make(map[uint64]string)

8
dax/role.go
View file

@ -28,8 +28,8 @@ type Role interface {
}
// Ensure type implements interface.
var _ Role = &ComputeRole{}
var _ Role = &TranslateRole{}
var _ Role = (*ComputeRole)(nil)
var _ Role = (*TranslateRole)(nil)
// ComputeRole is a role specific to compute nodes.
type ComputeRole struct {
@ -37,7 +37,7 @@ type ComputeRole struct {
Shards ShardNums `json:"shards"`
}
// Type returns the type for ComputeRole. This is mainly to impolement the Role
// Type returns the type for ComputeRole. This is mainly to implement the Role
// interface.
func (cr *ComputeRole) Type() RoleType {
return RoleTypeCompute
@ -50,7 +50,7 @@ type TranslateRole struct {
Fields []FieldName `json:"fields"`
}
// Type returns the type for TransteRole. This is mainly to impolement the Role
// Type returns the type for TranslateRole. This is mainly to implement the Role
// interface.
func (cr *TranslateRole) Type() RoleType {
return RoleTypeTranslate

58
dax/schema.go
View file

@ -3,16 +3,52 @@ package dax
import "context"
// Schemar is similar to the pilosa.SchemaAPI interface, but it takes
// TableQualifiers into account.
// QualifiedDatabaseIDs into account. Note that it is also similar to the
// schemar.Schemar interface, but that is used internally, typically within the
// Controller, and it takes Transactions rather than a Context, because its
// methods are assumed to be used as part of a larger request.
// TODO(tlt): clean up the mds/controller/schemar Schemar interface confusion.
type Schemar interface {
TableByName(ctx context.Context, qual TableQualifier, tname TableName) (*QualifiedTable, error)
TableByID(ctx context.Context, qtid QualifiedTableID) (*QualifiedTable, error)
Tables(ctx context.Context, qual TableQualifier, tids ...TableID) ([]*QualifiedTable, error)
//////////////////////////////////////////////////////////////////////////
// Database methods
//////////////////////////////////////////////////////////////////////////
CreateDatabase(context.Context, *QualifiedDatabase) error
// DropDatabase(context.Context, QualifiedDatabaseID) error
// DatabaseByName(ctx context.Context, orgID OrganizationID, dbname DatabaseName) (*QualifiedDatabase, error)
DatabaseByID(ctx context.Context, qdbid QualifiedDatabaseID) (*QualifiedDatabase, error)
// // Databases returns a list of databases. If the OrganizationID is empty,
// // all databases will be returned. If greater than zero database IDs are
// // passed in the second argument, only databases matching those IDs will be
// // returned.
// Databases(context.Context, OrganizationID, ...DatabaseID) ([]*QualifiedDatabase, error)
// SetDatabaseOptions(context.Context, QualifiedDatabaseID, DatabaseOptions) error
//////////////////////////////////////////////////////////////////////////
// Table methods
//////////////////////////////////////////////////////////////////////////
CreateTable(ctx context.Context, qtbl *QualifiedTable) error
CreateField(ctx context.Context, qtid QualifiedTableID, fld *Field) error
DropTable(ctx context.Context, qtid QualifiedTableID) error
TableByName(ctx context.Context, qdbid QualifiedDatabaseID, tname TableName) (*QualifiedTable, error)
TableByID(ctx context.Context, qtid QualifiedTableID) (*QualifiedTable, error)
// Tables returns a list of tables. If the qualifiers DatabaseID is empty,
// all tables in the org will be returned. If the OrganizationID is empty,
// all tables will be returned. If both are populated, only tables in that
// database will be returned. If greater than zero table IDs are passed in
// the third argument, only tables matching those IDs will be returned.
Tables(ctx context.Context, qdbid QualifiedDatabaseID, tids ...TableID) ([]*QualifiedTable, error)
//////////////////////////////////////////////////////////////////////////
// Field methods
//////////////////////////////////////////////////////////////////////////
CreateField(ctx context.Context, qtid QualifiedTableID, fld *Field) error
DropField(ctx context.Context, qtid QualifiedTableID, fname FieldName) error
}
@ -28,13 +64,19 @@ func NewNopSchemar() *NopSchemar {
return &NopSchemar{}
}
func (s *NopSchemar) TableByName(context.Context, TableQualifier, TableName) (*QualifiedTable, error) {
func (s *NopSchemar) CreateDatabase(context.Context, *QualifiedDatabase) error {
return nil
}
func (s *NopSchemar) DatabaseByID(ctx context.Context, qdbid QualifiedDatabaseID) (*QualifiedDatabase, error) {
return nil, nil
}
func (s *NopSchemar) TableByName(context.Context, QualifiedDatabaseID, TableName) (*QualifiedTable, error) {
return nil, nil
}
func (s *NopSchemar) TableByID(ctx context.Context, qtid QualifiedTableID) (*QualifiedTable, error) {
return nil, nil
}
func (s *NopSchemar) Tables(ctx context.Context, qual TableQualifier, tids ...TableID) ([]*QualifiedTable, error) {
func (s *NopSchemar) Tables(ctx context.Context, qdbid QualifiedDatabaseID, tids ...TableID) ([]*QualifiedTable, error) {
return nil, nil
}
func (s *NopSchemar) CreateTable(ctx context.Context, qtbl *QualifiedTable) error {

8
dax/storage/storage_test.go
View file

@ -34,10 +34,10 @@ func TestResourceManager(t *testing.T) {
mm := NewResourceManager(sn, wl, logger.NewStandardLogger(os.Stderr))
qtid := dax.QualifiedTableID{
TableQualifier: dax.TableQualifier{
OrganizationID: dax.OrganizationID("org1"),
DatabaseID: dax.DatabaseID("db1"),
},
QualifiedDatabaseID: dax.NewQualifiedDatabaseID(
dax.OrganizationID("org1"),
dax.DatabaseID("db1"),
),
ID: dax.TableID("blah"),
Name: "blah",
}

309
dax/table.go
View file

@ -22,14 +22,17 @@ import (
//
// OrganizationID - carried over from ControlPlane; currently uuid
// DatabaseID - carried over from ControlPlane; currently uuid
// Database - base Database struct
// DatabaseKey - a string representation of OrganizationID and DatabaseID
// TableID - internally stored as a uint64; presented as a hex string.
// TableName - human-friendly string table name
// Table - base Table struct; includes a TableID and a TableName
// TableQualifier - combination of OrganizationID and DatabaseID
// QualifiedTable - TableQualifier plus a Table
// QualifiedTableID - TableQualifier plus a TableID
// QualifiedDatabase - OrganizationID plus a Database
// QualifiedDatabaseID - combination of OrganizationID and DatabaseID
// QualifiedTable - QualifiedDatabaseID plus a Table
// QualifiedTableID - QualifiedDatabaseID plus a TableID
// TableKey - a string representation of OrganizationID, DatabaseID, and
// TableID, which is safe to use as a FeatureBase index name.
// TableID, which is safe to use as a FeatureBase index name.
//
// Example:
// OrganizationID - "29-ae44-41"
@ -37,7 +40,7 @@ import (
// TableID - 123456789 (hex string: "499602d2")
// TableName - foo
// Table - {ID:"499602d2", Name: "foo", Fields: ... }
// TableQualifier - {Org: "29-ae44-41", DB: "75-d1a2-4f"}
// QualifierDatabaseID - {Org: "29-ae44-41", DB: "75-d1a2-4f"}
// QualifiedTable - {Org: "29-ae44-41", DB: "75-d1a2-4f", Table: *tbl}
// QualifiedTableID - {Org: "29-ae44-41", DB: "75-d1a2-4f", TableID: "499602d2"}
// TableKey - "tbl__29-ae44-41__75-d1a2-4f__499602d2"
@ -52,6 +55,13 @@ import (
// underscore.
const TableKeyDelimiter = "__"
// PrefixDatabase is used as a prefix to DatabaseKey strings because FeatureBase
// indexes must start with an alpha (a-z) character. Because the string
// representation of a uuid (i.e. the OrganizationID value) can start with a
// numeric value, we can't have OrganizationId (or any of the other ID values
// which make up the DatabaseKey) be at the beginning of the DatabaseKey.
const PrefixDatabase = "db"
// PrefixTable is used as a prefix to TableKey strings because FeatureBase
// indexes must start with an alpha (a-z) character. Because the string
// representation of a uuid (i.e. the OrganizationID value) can start with a
@ -103,6 +113,124 @@ type OrganizationID string
// value could be any string.
type DatabaseID string
// DatabaseKey is a globally unique identifier for a database; it is effectively the
// compound key: (org, database). This is (hopefully) the value that will
// be used when interfacing with services which are unaware of qualifiers.
type DatabaseKey string
// QualifiedDatabaseID returns the QualifiedDatabaseID based on the key. If
// DatabaseKey can't be parsed into a valid (i.e. complete) QualifiedDatabaseID,
// then blank values are used where necessary.
func (dk DatabaseKey) QualifiedDatabaseID() QualifiedDatabaseID {
qdbid, err := QualifiedDatabaseIDFromKey(string(dk))
if err != nil {
return NewQualifiedDatabaseID("", DatabaseID(dk))
}
return qdbid
}
// DatabaseName is a human-friendly string.
type DatabaseName string
// Database represents a database and its configuration.
type Database struct {
ID DatabaseID `json:"id"`
Name DatabaseName `json:"name"`
Options DatabaseOptions `json:"options"`
// Tables []*Table `json:"tables"`
Description string `json:"description,omitempty"`
Owner string `json:"owner,omitempty"`
CreatedAt int64 `json:"createdAt,omitempty"`
UpdatedAt int64 `json:"updatedAt,omitempty"`
UpdatedBy string `json:"updatedBy,omitempty"`
}
// DatabaseOptions are used to configure a database.
type DatabaseOptions struct {
WorkersMin int `json:"workers-min"`
WorkersMax int `json:"workers-max"`
}
// QualifiedDatabase is a Database along with its OrganizationID.
type QualifiedDatabase struct {
OrganizationID OrganizationID `json:"org-id"`
Database
}
type QualifiedDatabases []*QualifiedDatabase
// Key returns the string-encoded (delimited by DatabaseKeyDelimiter) globally
// unique DatabaseKey.
func (qdb QualifiedDatabase) Key() DatabaseKey {
return qdb.QualifiedID().Key()
}
// String returns a human-friendly version of the QualifiedDatabase. It is only
// used for display purposes; it is not used as any kind of key.
func (qdb QualifiedDatabase) String() string {
return fmt.Sprintf("%s (%s)", qdb.QualifiedID(), qdb.Name)
}
// QualifiedID returns the QualifiedDatabaseID for the database.
func (qdb *QualifiedDatabase) QualifiedID() QualifiedDatabaseID {
return QualifiedDatabaseID{
OrganizationID: qdb.OrganizationID,
DatabaseID: qdb.ID,
}
}
// QualifiedDatabaseID is a DatabaseID along with its OrganizationID.
type QualifiedDatabaseID struct {
OrganizationID OrganizationID `json:"org-id"`
DatabaseID DatabaseID `json:"db-id"`
}
// NewQualifiedDatabaseID is a helper function used to create a
// QualifiedDatabaseID from the provided arguments.
func NewQualifiedDatabaseID(orgID OrganizationID, dbID DatabaseID) QualifiedDatabaseID {
return QualifiedDatabaseID{
OrganizationID: orgID,
DatabaseID: dbID,
}
}
// String returns a human-friendly version of the QualifiedDatabaseID. It is only
// used for display purposes; it is not used as any kind of key. For that, see
// the QualifiedDatabaseID.Key() method and the DatabaseKey type.
func (qdbid QualifiedDatabaseID) String() string {
return fmt.Sprintf("[%s:%s]", qdbid.OrganizationID, qdbid.DatabaseID)
}
// Key returns the string-encoded (delimited by TableKeyDelimiter) globally
// unique DatabaseKey. The key has a prefix because FeatureBase index name
// restrictions require the name to start with a non-numeric value, and since a
// uuid can contain a number as its first character, we have to prefix it with
// something.
func (qdbid QualifiedDatabaseID) Key() DatabaseKey {
if qdbid.DatabaseID == "" {
panic("QualifiedDatabaseID.Key called without an ID set")
}
return DatabaseKey(fmt.Sprintf("%s%s%s%s%s",
PrefixDatabase,
TableKeyDelimiter,
qdbid.OrganizationID,
TableKeyDelimiter,
qdbid.DatabaseID))
}
// QualifiedDatabaseIDs is a list of QualifiedDatabaseID.
type QualifiedDatabaseIDs []QualifiedDatabaseID
func (s QualifiedDatabaseIDs) Len() int { return len(s) }
func (s QualifiedDatabaseIDs) Less(i, j int) bool {
if s[i].OrganizationID != s[j].OrganizationID {
return s[i].OrganizationID < s[j].OrganizationID
}
return s[i].DatabaseID < s[j].DatabaseID
}
func (s QualifiedDatabaseIDs) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
// TableKeyer is an interface implemented by any type which can produce, and be
// represented by, a TableKey. In the case of a QualifiedTable, its TableKey
// might be something like `tbl__org__db__tableid`, while a general pilosa
@ -129,7 +257,7 @@ func (s StringTableKeyer) Key() TableKey {
// TableKey as the value for index.Name.
type TableKey string
func (t TableKey) Key() TableKey { return t }
func (tk TableKey) Key() TableKey { return tk }
// QualifiedTableID returns the QualifiedTableID based on the key. If TableKey
// can't be parsed into a valid (i.e. complete) QualifiedTableID, then blank
@ -138,7 +266,7 @@ func (tk TableKey) QualifiedTableID() QualifiedTableID {
qtid, err := QualifiedTableIDFromKey(string(tk))
if err != nil {
return NewQualifiedTableID(
NewTableQualifier("", ""),
NewQualifiedDatabaseID("", ""),
TableID(tk),
)
}
@ -153,7 +281,7 @@ func (s TableKeys) Less(i, j int) bool { return s[i] < s[j] }
func (s TableKeys) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
// TableID is a table identifier. It is unique within the scope of a
// TableQualifier. Coupled with a TableQualifier, it makes up a
// QualifiedDatabaseID. Coupled with a QualifiedDatabaseID, it makes up a
// QualifiedTableID and, when encoded as a string, a TableKey.
type TableID string
@ -165,7 +293,7 @@ func (s TableIDs) Less(i, j int) bool { return s[i] < s[j] }
func (s TableIDs) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
// TableName is a human-friendly string. While it is not used as a primary key,
// uniqueness is generally enforced within the scope of a TableQualifier.
// uniqueness is generally enforced within the scope of a QualifiedDatabaseID.
type TableName string
// TableNames is a sortable slice of TableName.
@ -224,7 +352,7 @@ func (t *Table) CreateID() (TableID, error) {
}
// NewTable returns a new instance of table with a pseudo-random ID which is
// assumed to be unique within the scope of a TableQualifier.
// assumed to be unique within the scope of a QualifiedDatabaseID.
func NewTable(name TableName) *Table {
return &Table{
Name: name,
@ -315,84 +443,6 @@ func (o Tables) Len() int { return len(o) }
func (o Tables) Less(i, j int) bool { return o[i].Name < o[j].Name }
func (o Tables) Swap(i, j int) { o[i], o[j] = o[j], o[i] }
// TableQualifierKey is the unique TableQualifier values encoded as a string. The
// current encoding is delimited as `prefix|OrganizationID|DatabaseID` (where
// the pipe may be some other delimiter) by the TableQualifier.Key() method.
type TableQualifierKey string
// Qualifier returns the Qualifier based on the values encoded into the
// TableQualifierKey string.
func (tqk TableQualifierKey) Qualifier() TableQualifier {
parts := strings.Split(string(tqk), TableKeyDelimiter)
if len(parts) < 3 {
return NewTableQualifier("", "")
}
return NewTableQualifier(
OrganizationID(parts[1]),
DatabaseID(parts[2]),
)
}
// OrganizationID returns the OrganizationID value that has been encoded into
// the TableQualifierKey string.
func (tqk TableQualifierKey) OrganizationID() OrganizationID {
parts := strings.Split(string(tqk), TableKeyDelimiter)
if len(parts) < 2 {
return ""
}
return OrganizationID(parts[1])
}
// DatabaseID returns the DatabaseID value that has been encoded into the
// TableQualifierKey string.
func (tqk TableQualifierKey) DatabaseID() DatabaseID {
parts := strings.Split(string(tqk), TableKeyDelimiter)
if len(parts) < 3 {
return ""
}
return DatabaseID(parts[2])
}
// TableQualifier contains all the elements required to fully qualify a table.
type TableQualifier struct {
OrganizationID OrganizationID `json:"org-id"`
DatabaseID DatabaseID `json:"db-id"`
}
// NewTableQualifier is a helper function used to create a TableQualifier from
// the provided arguments.
func NewTableQualifier(orgID OrganizationID, dbID DatabaseID) TableQualifier {
return TableQualifier{
OrganizationID: orgID,
DatabaseID: dbID,
}
}
// String returns a human-friendly version of the TableQualifier. It is only
// used for display purposes; it is not used as any kind of key. For that, see
// the TableQualifier.Key() method and the TableQualifierKey type.
func (tq TableQualifier) String() string {
return fmt.Sprintf("[%s:%s]", tq.OrganizationID, tq.DatabaseID)
}
// Key returns the string-encoded (delimited by TableKeyDelimiter)
// TableQualifierKey.
func (tq TableQualifier) Key() TableQualifierKey {
return TableQualifierKey(fmt.Sprintf("%s%s%s%s%s",
PrefixTable,
TableKeyDelimiter,
tq.OrganizationID,
TableKeyDelimiter,
tq.DatabaseID,
))
}
////////////////////////////////////////////////
// QualifiedTableID is a globally unique table identifier. It is a
@ -400,17 +450,17 @@ func (tq TableQualifier) Key() TableQualifierKey {
// portion). Most things will take a Name or an ID and do the right
// thing™.
type QualifiedTableID struct {
TableQualifier
QualifiedDatabaseID
ID TableID `json:"id"`
Name TableName `json:"name"`
}
// NewQualifiedTableID is a helper function used to create a QualifiedTableID
// from the provided arguments.
func NewQualifiedTableID(q TableQualifier, id TableID) QualifiedTableID {
func NewQualifiedTableID(qdbid QualifiedDatabaseID, tid TableID) QualifiedTableID {
return QualifiedTableID{
TableQualifier: q,
ID: id,
QualifiedDatabaseID: qdbid,
ID: tid,
}
}
@ -422,7 +472,7 @@ func QualifiedTableIDFromKey(key string) (QualifiedTableID, error) {
case 4:
// prefix|orgID|dbID|tblID
return NewQualifiedTableID(
NewTableQualifier(
NewQualifiedDatabaseID(
OrganizationID(parts[1]),
DatabaseID(parts[2]),
),
@ -433,14 +483,31 @@ func QualifiedTableIDFromKey(key string) (QualifiedTableID, error) {
}
}
// String returns a human-friendly version of the TableQualifierID. It is only
// used for display purposes; it is not used as any kind of key. For that, see
// the TableQualifierID.Key() method.
// QualifiedDatabaseIDFromKey decodes a string key into a QualifiedDatabaseID.
// The key is assumed to have been encoded using the QualifiedDatabaseID.Key()
// method.
func QualifiedDatabaseIDFromKey(key string) (QualifiedDatabaseID, error) {
parts := strings.Split(key, TableKeyDelimiter)
switch len(parts) {
case 3:
// prefix|orgID|dbID
return NewQualifiedDatabaseID(
OrganizationID(parts[1]),
DatabaseID(parts[2]),
), nil
default:
return QualifiedDatabaseID{}, errors.Errorf("invalid key: %s", key)
}
}
// String returns a human-friendly version of the QualifiedDatabaseID. It is
// only used for display purposes; it is not used as any kind of key. For that,
// see the QualifiedDatabaseID.Key() method.
func (qtid QualifiedTableID) String() string {
if qtid.ID == "" {
return fmt.Sprintf("%s%s", qtid.TableQualifier, qtid.Name)
return fmt.Sprintf("%s%s", qtid.QualifiedDatabaseID, qtid.Name)
}
return fmt.Sprintf("%s%s", qtid.TableQualifier, qtid.ID)
return fmt.Sprintf("%s%s", qtid.QualifiedDatabaseID, qtid.ID)
}
// Key returns the string-encoded (delimited by TableKeyDelimiter) globally
@ -452,35 +519,49 @@ func (qtid QualifiedTableID) Key() TableKey {
if qtid.ID == "" {
panic("QualifiedTableID.Key called without an ID set")
}
return TableKey(fmt.Sprintf("%s%s%s",
qtid.TableQualifier.Key(),
return TableKey(fmt.Sprintf("%s%s%s%s%s%s%s",
PrefixTable,
TableKeyDelimiter,
qtid.OrganizationID,
TableKeyDelimiter,
qtid.DatabaseID,
TableKeyDelimiter,
qtid.ID))
}
// Equals returns true if `other` is the same as qtid. Note: the `Name` value is
// ignored in this comparison; only `TableQualifier` and `ID` are considered.
// ignored in this comparison; only `QualifiedDatabaseID` and `ID` are
// considered.
func (qtid QualifiedTableID) Equals(other QualifiedTableID) bool {
if qtid.TableQualifier == other.TableQualifier && qtid.ID == other.ID {
if qtid.QualifiedDatabaseID == other.QualifiedDatabaseID && qtid.ID == other.ID {
return true
}
return false
}
// Qualifier returns the QualifiedDatabaseID (qdbid) portion of the
// QualifiedTableID (qtid).
func (qtid QualifiedTableID) Qualifier() QualifiedDatabaseID {
return QualifiedDatabaseID{
OrganizationID: qtid.OrganizationID,
DatabaseID: qtid.DatabaseID,
}
}
////////////////////////////////////////////////
// QualifiedTable wraps Table and includes a TableQualifier.
// QualifiedTable wraps Table and includes a QualifiedDatabaseID.
type QualifiedTable struct {
QualifiedDatabaseID
Table
TableQualifier
}
// NewQualifiedTable returns the tbl as a QualifiedTable with the provided
// TableQualifier.
func NewQualifiedTable(qual TableQualifier, tbl *Table) *QualifiedTable {
// QualifiedDatabaseID.
func NewQualifiedTable(qdbid QualifiedDatabaseID, tbl *Table) *QualifiedTable {
return &QualifiedTable{
Table: *tbl,
TableQualifier: qual,
QualifiedDatabaseID: qdbid,
Table: *tbl,
}
}
@ -496,17 +577,17 @@ func (qt QualifiedTable) String() string {
return fmt.Sprintf("%s (%s)", qt.QualifiedID(), qt.Name)
}
// Qualifier returns the TableQualifier portion of the QualifiedTable.
func (qt *QualifiedTable) Qualifier() TableQualifier {
return qt.TableQualifier
// Qualifier returns the QualifiedDatabaseID portion of the QualifiedTable.
func (qt *QualifiedTable) Qualifier() QualifiedDatabaseID {
return qt.QualifiedDatabaseID
}
// QualifiedID returns the QualifiedTableID for the table.
func (qt *QualifiedTable) QualifiedID() QualifiedTableID {
return QualifiedTableID{
TableQualifier: qt.TableQualifier,
ID: qt.ID,
Name: qt.Name,
QualifiedDatabaseID: qt.QualifiedDatabaseID,
ID: qt.ID,
Name: qt.Name,
}
}

12
dax/table_test.go
View file

@ -313,17 +313,17 @@ func TestTable(t *testing.T) {
t.Run("New", func(t *testing.T) {
tbl := dax.NewTable(tableName)
tbl.CreateID()
qual := dax.TableQualifier{
qdbid := dax.QualifiedDatabaseID{
OrganizationID: orgID,
DatabaseID: dbID,
}
qtbl := dax.NewQualifiedTable(qual, tbl)
qtbl := dax.NewQualifiedTable(qdbid, tbl)
assert.NotEmpty(t, qtbl.ID)
assert.Equal(t, tbl.ID, qtbl.ID)
tq := qtbl.Qualifier()
assert.Equal(t, qual.OrganizationID, tq.OrganizationID)
assert.Equal(t, qual.DatabaseID, tq.DatabaseID)
assert.Equal(t, qdbid.OrganizationID, tq.OrganizationID)
assert.Equal(t, qdbid.DatabaseID, tq.DatabaseID)
wrappedTable := qtbl.Table
assert.Equal(t, tbl.Name, wrappedTable.Name)
@ -345,11 +345,11 @@ func TestTable(t *testing.T) {
t.Run("ToJSON", func(t *testing.T) {
tbl := dax.NewTable(tableName)
tbl.CreateID()
qual := dax.TableQualifier{
qdbid := dax.QualifiedDatabaseID{
OrganizationID: orgID,
DatabaseID: dbID,
}
qtbl := dax.NewQualifiedTable(qual, tbl)
qtbl := dax.NewQualifiedTable(qdbid, tbl)
id := qtbl.ID
b, err := json.Marshal(qtbl)

137
dax/test/dax/dax_test.go
View file

@ -26,7 +26,18 @@ func TestDAXIntegration(t *testing.T) {
t.Skip("skipping integration test")
}
qual := dax.NewTableQualifier("acme", "db1")
qdbid := dax.NewQualifiedDatabaseID("acme", "db1")
qdb := &dax.QualifiedDatabase{
OrganizationID: qdbid.OrganizationID,
Database: dax.Database{
ID: qdbid.DatabaseID,
Name: "dbname1",
Options: dax.DatabaseOptions{
WorkersMin: 1,
WorkersMax: 1,
},
},
}
t.Run("ServiceStart", func(t *testing.T) {
t.Run("AllServicesByDefault", func(t *testing.T) {
@ -104,6 +115,14 @@ func TestDAXIntegration(t *testing.T) {
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 1
qdb.Options.WorkersMax = 1
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
// skips is a list of tests which are currently not passing in dax. We
// need to get these passing before alpha.
skips := []string{
@ -112,7 +131,9 @@ func TestDAXIntegration(t *testing.T) {
"alterTable/alterTableBadTable", // looks like table does not exist is a different error in DAX
"top-tests/test-1", // don't know why this is failing at all
"delete_tests",
"subquerytable", // subqueries seem to be a problem
"viewtests/drop-view", // drop view does a delete
"viewtests/drop-view-if-exists-after-drop",
"viewtests/select-view-after-drop",
}
doSkip := func(name string) bool {
@ -158,7 +179,7 @@ func TestDAXIntegration(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
basicTableTestConfig(qual, tableTests...)...,
basicTableTestConfig(qdbid, tableTests...)...,
)
})
@ -166,9 +187,19 @@ func TestDAXIntegration(t *testing.T) {
mc := test.MustRunManagedCommand(t)
defer mc.Close()
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 1
qdb.Options.WorkersMax = 1
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
runTableTests(t,
mc.Manage().Queryer.Address(),
basicTableTestConfig(qual, defs.Keyed)...,
basicTableTestConfig(qdbid, defs.Keyed)...,
)
})
@ -180,6 +211,14 @@ func TestDAXIntegration(t *testing.T) {
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 2
qdb.Options.WorkersMax = 2
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
computers := svcmgr.Computers()
computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
computerKey1 := dax.ServiceKey(dax.ServicePrefixComputer + "1")
@ -187,13 +226,10 @@ func TestDAXIntegration(t *testing.T) {
// Ingest and query some data.
runTableTests(t,
svcmgr.Queryer.Address(),
basicTableTestConfig(qual, defs.Keyed)...,
basicTableTestConfig(qdbid, defs.Keyed)...,
)
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
qtid, err := mdsClient.TableID(context.Background(), qual, dax.TableName(defs.Keyed.Name(0)))
qtid, err := mdsClient.TableID(context.Background(), qdbid, dax.TableName(defs.Keyed.Name(0)))
assert.NoError(t, err)
// ensure partitions are covered
@ -236,13 +272,21 @@ func TestDAXIntegration(t *testing.T) {
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 1
qdb.Options.WorkersMax = 1
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
// Ingest and query some data.
t.Run("ingest and query some data", func(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
basicTableTestConfig(qual, defs.Keyed)...,
basicTableTestConfig(qdbid, defs.Keyed)...,
)
})
@ -266,7 +310,7 @@ func TestDAXIntegration(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
tableTestConfig{
qual: qual,
qdbid: qdbid,
test: defs.Keyed,
skipCreate: true,
skipInsert: true,
@ -282,6 +326,14 @@ func TestDAXIntegration(t *testing.T) {
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 1
qdb.Options.WorkersMax = 1
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
// Ingest and query some data.
@ -289,19 +341,16 @@ func TestDAXIntegration(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
tableTestConfig{
qual: qual,
qdbid: qdbid,
test: defs.Keyed,
insertSet: 0,
},
)
})
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Snapshot table
ctx := context.Background()
qtid, err := mdsClient.TableID(ctx, qual, dax.TableName(defs.Keyed.Name(0)))
qtid, err := mdsClient.TableID(ctx, qdbid, dax.TableName(defs.Keyed.Name(0)))
assert.NoError(t, err)
mdsClient.SnapshotTable(ctx, qtid)
@ -311,7 +360,7 @@ func TestDAXIntegration(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
tableTestConfig{
qual: qual,
qdbid: qdbid,
test: defs.Keyed,
skipCreate: true,
insertSet: 1,
@ -339,7 +388,7 @@ func TestDAXIntegration(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
tableTestConfig{
qual: qual,
qdbid: qdbid,
test: defs.Keyed,
skipCreate: true,
skipInsert: true,
@ -355,6 +404,14 @@ func TestDAXIntegration(t *testing.T) {
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 1
qdb.Options.WorkersMax = 1
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
mdsKey := dax.ServiceKey(dax.ServicePrefixMDS)
@ -363,13 +420,10 @@ func TestDAXIntegration(t *testing.T) {
// Ingest and query some data.
runTableTests(t,
svcmgr.Queryer.Address(),
basicTableTestConfig(qual, defs.Keyed)...,
basicTableTestConfig(qdbid, defs.Keyed)...,
)
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
qtid, err := mdsClient.TableID(context.Background(), qual, dax.TableName(defs.Keyed.Name(0)))
qtid, err := mdsClient.TableID(context.Background(), qdbid, dax.TableName(defs.Keyed.Name(0)))
assert.NoError(t, err)
// ensure partitions are covered
@ -408,13 +462,21 @@ func TestDAXIntegration(t *testing.T) {
svcmgr := mc.Manage()
// Set up MDS client.
mdsClient := mdsclient.New(svcmgr.MDS.Address(), svcmgr.Logger)
// Create database.
qdb.Options.WorkersMin = 1
qdb.Options.WorkersMax = 1
assert.NoError(t, mdsClient.CreateDatabase(context.Background(), qdb))
computerKey0 := dax.ServiceKey(dax.ServicePrefixComputer + "0")
// Ingest and query some data.
t.Run("ingest and query some data", func(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
basicTableTestConfig(qual, defs.Keyed)...,
basicTableTestConfig(qdbid, defs.Keyed)...,
)
})
@ -437,7 +499,7 @@ func TestDAXIntegration(t *testing.T) {
runTableTests(t,
svcmgr.Queryer.Address(),
tableTestConfig{
qual: qual,
qdbid: qdbid,
test: defs.Keyed,
skipCreate: true,
skipInsert: true,
@ -450,7 +512,7 @@ func TestDAXIntegration(t *testing.T) {
///////////////////////////////////////////////////
type tableTestConfig struct {
qual dax.TableQualifier
qdbid dax.QualifiedDatabaseID
test defs.TableTest
skipCreate bool
skipInsert bool
@ -459,20 +521,19 @@ type tableTestConfig struct {
querySet int
}
func basicTableTestConfig(qual dax.TableQualifier, tests ...defs.TableTest) []tableTestConfig {
func basicTableTestConfig(qdbid dax.QualifiedDatabaseID, tests ...defs.TableTest) []tableTestConfig {
ret := make([]tableTestConfig, len(tests))
for i := range tests {
ret[i] = tableTestConfig{
qual: qual,
test: tests[i],
qdbid: qdbid,
test: tests[i],
}
}
return ret
}
// func runTableTests(t *testing.T, queryerAddr dax.Address, qual dax.TableQualifier, doCreate bool, tests ...defs.TableTest) {
func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfig) {
emptyWireQueryResponse := &featurebase.WireQueryResponse{
Schema: featurebase.WireQuerySchema{
@ -489,7 +550,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
if !cfg.skipCreate {
// Create a table.
if cfg.test.HasTable() {
resp := runSQL(t, queryerAddr, cfg.qual, cfg.test.CreateTable())
resp := runSQL(t, queryerAddr, cfg.qdbid, cfg.test.CreateTable())
assertResponseEqual(t, emptyWireQueryResponse, resp)
}
}
@ -497,7 +558,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
if !cfg.skipInsert {
// Populate table with data.
if cfg.test.HasTable() && cfg.test.HasData() {
resp := runSQL(t, queryerAddr, cfg.qual, cfg.test.InsertInto(t, cfg.insertSet))
resp := runSQL(t, queryerAddr, cfg.qdbid, cfg.test.InsertInto(t, cfg.insertSet))
assertResponseEqual(t, emptyWireQueryResponse, resp)
}
}
@ -523,7 +584,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
expRows = sqltest.ExpRowsPlus1[cfg.querySet-1]
}
resp := runSQL(t, queryerAddr, cfg.qual, sql)
resp := runSQL(t, queryerAddr, cfg.qdbid, sql)
headers := resp.Schema.Fields
rows := resp.Data
var err error
@ -600,7 +661,7 @@ func runTableTests(t *testing.T, queryerAddr dax.Address, cfgs ...tableTestConfi
expRows = pqltest.ExpRowsPlus1[cfg.querySet-1]
}
resp := runPQL(t, queryerAddr, cfg.qual, pqltest.Table, pql)
resp := runPQL(t, queryerAddr, cfg.qdbid, pqltest.Table, pql)
headers := resp.Schema.Fields
rows := resp.Data
var err error
@ -678,23 +739,23 @@ func (c *wireResponseComparer) Equal() bool {
return assert.Equal(c.tb, c.exp, c.got)
}
func runSQL(tb testing.TB, queryerAddr dax.Address, qual dax.TableQualifier, sql string) *featurebase.WireQueryResponse {
func runSQL(tb testing.TB, queryerAddr dax.Address, qdbid dax.QualifiedDatabaseID, sql string) *featurebase.WireQueryResponse {
tb.Helper()
client := queryerclient.New(queryerAddr, logger.StderrLogger)
resp, err := client.QuerySQL(context.Background(), qual, sql)
resp, err := client.QuerySQL(context.Background(), qdbid, sql)
assert.NoError(tb, err)
return resp
}
func runPQL(tb testing.TB, queryerAddr dax.Address, qual dax.TableQualifier, table string, pql string) *featurebase.WireQueryResponse {
func runPQL(tb testing.TB, queryerAddr dax.Address, qdbid dax.QualifiedDatabaseID, table string, pql string) *featurebase.WireQueryResponse {
tb.Helper()
client := queryerclient.New(queryerAddr, logger.StderrLogger)
resp, err := client.QueryPQL(context.Background(), qual, dax.TableName(table), pql)
resp, err := client.QueryPQL(context.Background(), qdbid, dax.TableName(table), pql)
assert.NoError(tb, err)
return resp

28
dax/test/table.go
View file

@ -11,7 +11,7 @@ import (
// general configuration. This function creates a Table with a random TableID.
// If you need to specify the TableID yourself, use the TestQualifiedTableWithID
// function.
func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
func TestQualifiedTable(t *testing.T, qdbid dax.QualifiedDatabaseID, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
t.Helper()
var pkFieldType dax.BaseType
@ -22,7 +22,6 @@ func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableNam
}
tbl := dax.NewTable(name)
tbl.CreateID()
tbl.PartitionN = partitionN
tbl.Fields = []*dax.Field{
{
@ -32,14 +31,33 @@ func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableNam
}
return dax.NewQualifiedTable(
qual,
qdbid,
tbl,
)
}
// TestQualifiedDatabaseWithID is a test helper function for creating a database
// based on a general configuration, and having the specified DatabaseID.
func TestQualifiedDatabaseWithID(t *testing.T, orgID dax.OrganizationID, id dax.DatabaseID, name dax.DatabaseName, opts dax.DatabaseOptions) *dax.QualifiedDatabase {
t.Helper()
db := dax.Database{
ID: dax.DatabaseID(id),
Name: name,
Options: opts,
}
qdb := &dax.QualifiedDatabase{
OrganizationID: orgID,
Database: db,
}
return qdb
}
// TestQualifiedTableWithID is a test helper function for creating a table based
// on a general configuration, and having the specified TableID.
func TestQualifiedTableWithID(t *testing.T, qual dax.TableQualifier, id string, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
func TestQualifiedTableWithID(t *testing.T, qdbid dax.QualifiedDatabaseID, id string, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable {
t.Helper()
var pkFieldType dax.BaseType
@ -63,7 +81,7 @@ func TestQualifiedTableWithID(t *testing.T, qual dax.TableQualifier, id string,
}
return dax.NewQualifiedTable(
qual,
qdbid,
tbl,
)
}

9
dax/transaction.go Normal file
View file

@ -0,0 +1,9 @@
package dax
import "context"
type Transaction interface {
Commit() error
Context() context.Context
Rollback() error
}

38
dax/workerjob.go
View file

@ -5,44 +5,39 @@ import (
"strings"
)
// Worker is a generic identifier used to represent a service responsible for
// doing certain jobs. In the case of dax, this is typically the Address of a
// compute or translate node. Services such as the Balancer use Workers (as
// opposed to specifically using Address) in order to remain generic, and to
// keep the business logic between services slightly less coupled.
type Worker string
// Workers is a sortable slice of Worker.
type Workers []Worker
func (w Workers) Len() int { return len(w) }
func (w Workers) Less(i, j int) bool { return w[i] < w[j] }
func (w Workers) Swap(i, j int) { w[i], w[j] = w[j], w[i] }
// Job is a generic identifier used to represent a specific role assigned to a
// worker.
type Job string
// Job allows a Job to implement the Jobber interface.
func (j Job) Job() Job {
return j
}
// Jobs is a slice of Job.
type Jobs []Job
type Jobber interface {
Job() Job
}
// WorkerInfo represents a Worker and the Jobs to which it has been assigned.
type WorkerInfo struct {
ID Worker
Jobs []Job
Address Address
Jobs []Job
}
// WorkerInfos is a sortable slice of WorkerInfo.
type WorkerInfos []WorkerInfo
func (w WorkerInfos) Len() int { return len(w) }
func (w WorkerInfos) Less(i, j int) bool { return w[i].ID < w[j].ID }
func (w WorkerInfos) Less(i, j int) bool { return w[i].Address < w[j].Address }
func (w WorkerInfos) Swap(i, j int) { w[i], w[j] = w[j], w[i] }
// WorkerDiff represents the changes made to a Worker following the latest
// event.
type WorkerDiff struct {
WorkerID Worker
Address Address
AddedJobs []Job
RemovedJobs []Job
}
@ -51,7 +46,7 @@ type WorkerDiff struct {
// ID. Any job that is added and then removed or removed and then
// added cancels out and won't be present after add is called.
func (w *WorkerDiff) Add(w2 WorkerDiff) {
if w.WorkerID != w2.WorkerID {
if w.Address != w2.Address {
panic("can't add worker diffs from different workers")
}
a1 := NewSet(w.AddedJobs...)
@ -74,7 +69,7 @@ func (w *WorkerDiff) Add(w2 WorkerDiff) {
type WorkerDiffs []WorkerDiff
func (w WorkerDiffs) Len() int { return len(w) }
func (w WorkerDiffs) Less(i, j int) bool { return w[i].WorkerID < w[j].WorkerID }
func (w WorkerDiffs) Less(i, j int) bool { return w[i].Address < w[j].Address }
func (w WorkerDiffs) Swap(i, j int) { w[i], w[j] = w[j], w[i] }
// Set is a set of stringy items.
@ -109,7 +104,8 @@ func (s Set[K]) Remove(k K) {
delete(s, k)
}
func (s Set[K]) RemovePrefix(prefix string) []K {
// RemoveByPrefix removes all items from Set that have the given prefix.
func (s Set[K]) RemoveByPrefix(prefix string) []K {
ret := make([]K, 0)
for k := range s {
if strings.HasPrefix(string(k), prefix) {

104
executor.go
View file

@ -29,6 +29,7 @@ import (
"github.com/featurebasedb/featurebase/v3/testhook"
"github.com/featurebasedb/featurebase/v3/tracing"
"github.com/pkg/errors"
"github.com/prometheus/client_golang/prometheus"
"golang.org/x/sync/errgroup"
)
@ -165,16 +166,16 @@ func (e *executor) Close() error {
// PoolSize is exported to let the task pool update us
func (e *executor) PoolSize(n int) {
if e.Holder != nil {
e.Holder.Stats.Gauge("worker_total", float64(n), 0)
GaugeWorkerTotal.Set(float64(n))
}
}
// InitStats initializes stats counters. Must be called after Holder set.
func (e *executor) InitStats() {
if e.Holder != nil {
e.Holder.Stats.Count("job_total", 0, 0)
CounterJobTotal.Add(0)
l, _, _ := e.workers.Stats()
e.Holder.Stats.Gauge("worker_total", float64(l), 0)
GaugeWorkerTotal.Set(float64(l))
}
}
@ -684,11 +685,11 @@ func (e *executor) executeCall(ctx context.Context, qcx *Qcx, index string, c *p
} else if err := e.validateCallArgs(c); err != nil {
return nil, errors.Wrap(err, "validating args")
}
indexTag := "index:" + index
metricName := "query_" + strings.ToLower(c.Name) + "_total"
statFn := func() {
labels := prometheus.Labels{"index": index}
statFn := func(ctr *prometheus.CounterVec) {
if !opt.Remote {
e.Holder.Stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag})
ctr.With(labels).Inc()
}
}
@ -720,114 +721,122 @@ func (e *executor) executeCall(ctx context.Context, qcx *Qcx, index string, c *p
switch c.Name {
case "Sum":
statFn()
statFn(CounterQuerySumTotal)
res, err := e.executeSum(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeSum")
case "Min":
statFn()
statFn(CounterQueryMinTotal)
res, err := e.executeMin(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMin")
case "Max":
statFn()
statFn(CounterQueryMaxTotal)
res, err := e.executeMax(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMax")
case "MinRow":
statFn()
statFn(CounterQueryMinRowTotal)
res, err := e.executeMinRow(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMinRow")
case "MaxRow":
statFn()
statFn(CounterQueryMaxRowTotal)
res, err := e.executeMaxRow(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMaxRow")
case "Clear":
statFn()
statFn(CounterQueryClearTotal)
res, err := e.executeClearBit(ctx, qcx, index, c, opt)
return res, errors.Wrap(err, "executeClearBit")
case "ClearRow":
statFn()
statFn(CounterQueryClearRowTotal)
res, err := e.executeClearRow(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeClearRow")
case "Distinct":
statFn()
statFn(CounterQueryDistinctTotal)
res, err := e.executeDistinct(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeDistinct")
case "Store":
statFn()
statFn(CounterQueryStoreTotal)
res, err := e.executeSetRow(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeSetRow")
case "Count":
statFn()
statFn(CounterQueryCountTotal)
res, err := e.executeCount(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeCount")
case "Set":
statFn()
statFn(CounterQuerySetTotal)
res, err := e.executeSet(ctx, qcx, index, c, opt)
return res, errors.Wrap(err, "executeSet")
case "TopK":
statFn()
statFn(CounterQueryTopKTotal)
res, err := e.executeTopK(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeTopK")
case "TopN":
statFn()
statFn(CounterQueryTopNTotal)
res, err := e.executeTopN(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeTopN")
case "Rows":
statFn()
statFn(CounterQueryRowsTotal)
res, err := e.executeRows(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeRows")
case "ExternalLookup":
statFn()
statFn(CounterQueryExternalLookupTotal)
res, err := e.executeExternalLookup(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeExternalLookup")
case "Extract":
statFn()
statFn(CounterQueryExtractTotal)
res, err := e.executeExtract(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeExtract")
case "GroupBy":
statFn()
statFn(CounterQueryGroupByTotal)
res, err := e.executeGroupBy(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeGroupBy")
case "Options":
statFn()
statFn(CounterQueryOptionsTotal)
res, err := e.executeOptionsCall(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeOptionsCall")
case "IncludesColumn":
statFn(CounterQueryIncludesColumnTotal)
res, err := e.executeIncludesColumnCall(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeIncludesColumnCall")
case "FieldValue":
statFn()
statFn(CounterQueryFieldValueTotal)
res, err := e.executeFieldValueCall(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeFieldValueCall")
case "Precomputed":
statFn(CounterQueryPrecomputedTotal)
res, err := e.executePrecomputedCall(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executePrecomputedCall")
case "UnionRows":
statFn(CounterQueryUnionRowsTotal)
res, err := e.executeUnionRows(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeUnionRows")
case "ConstRow":
statFn(CounterQueryConstRowTotal)
res, err := e.executeConstRow(ctx, index, c)
return res, errors.Wrap(err, "executeConstRow")
case "Limit":
statFn(CounterQueryLimitTotal)
res, err := e.executeLimitCall(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeLimitCall")
case "Percentile":
statFn(CounterQueryPercentileTotal)
res, err := e.executePercentile(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executePercentile")
case "Delete":
statFn() // TODO(twg) need this?
statFn(CounterQueryDeleteTotal)
res, err := e.executeDeleteRecords(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeDelete")
case "Sort":
statFn(CounterQuerySortTotal)
res, err := e.executeSort(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeSort")
case "Apply":
statFn(CounterQueryApplyTotal)
res, err := e.executeApply(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeApply")
case "Arrow":
statFn(CounterQueryArrowTotal)
res, err := e.executeArrow(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeArrow")
default: // e.g. "Row", "Union", "Intersect" or anything that returns a bitmap.
statFn()
res, err := e.executeBitmapCall(ctx, qcx, index, c, shards, opt)
return res, errors.Wrap(err, "executeBitmapCall")
}
@ -1506,13 +1515,42 @@ func (e *executor) executeBitmapCall(ctx context.Context, qcx *Qcx, index string
span.LogKV("pqlCallName", c.Name)
defer span.Finish()
indexTag := "index:" + index
metricName := "query_" + strings.ToLower(c.Name) + "_total"
if c.Name == "Row" && c.HasConditionArg() {
metricName = "query_row_bsi_total"
labels := prometheus.Labels{"index": index}
statFn := func(ctr *prometheus.CounterVec) {
if !opt.Remote {
ctr.With(labels).Inc()
}
}
if !opt.Remote {
e.Holder.Stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag})
switch c.Name {
case "Row":
if c.HasConditionArg() {
statFn(CounterQueryRowBSITotal)
} else {
statFn(CounterQueryRowTotal)
}
case "Range":
statFn(CounterQueryRangeTotal)
case "Difference":
statFn(CounterQueryBitmapTotal)
case "Intersect":
statFn(CounterQueryIntersectTotal)
case "Union":
statFn(CounterQueryUnionTotal)
case "InnerUnionRows":
statFn(CounterQueryInnerUnionRowsTotal)
case "Xor":
statFn(CounterQueryXorTotal)
case "Not":
statFn(CounterQueryNotTotal)
case "Shift":
statFn(CounterQueryShiftTotal)
case "All":
statFn(CounterQueryAllTotal)
default:
statFn(CounterQueryBitmapTotal)
}
}
// Execute calls in bulk on each remote node and merge.

4
field.go
View file

@ -18,7 +18,6 @@ import (
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/testhook"
"github.com/featurebasedb/featurebase/v3/tracing"
"github.com/pkg/errors"
@ -84,7 +83,6 @@ type Field struct {
viewMap map[string]*view
broadcaster broadcaster
Stats stats.StatsClient
serializer Serializer
// Field options.
@ -396,7 +394,6 @@ func newField(holder *Holder, path, index, name string, opts FieldOption) (*Fiel
viewMap: make(map[string]*view),
broadcaster: NopBroadcaster,
Stats: stats.NopStatsClient,
serializer: NopSerializer,
options: applyDefaultOptions(&fo),
@ -1165,7 +1162,6 @@ func (f *Field) newView(path, name string) *view {
view := newView(f.holder, path, f.index, f.name, name, f.options)
view.idx = f.idx
view.fld = f
view.stats = f.Stats
view.broadcaster = f.broadcaster
return view
}

19
fragment.go
View file

@ -24,7 +24,6 @@ import (
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/featurebasedb/featurebase/v3/shardwidth"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/testhook"
"github.com/featurebasedb/featurebase/v3/tracing"
"github.com/featurebasedb/featurebase/v3/vprint"
@ -120,8 +119,6 @@ type fragment struct {
// mutexVector is used for mutex field types. It's checked for an
// existing value (to clear) prior to setting a new value.
mutexVector vector
stats stats.StatsClient
}
// newFragment returns a new instance of fragment.
@ -142,8 +139,6 @@ func newFragment(holder *Holder, idx *Index, fld *Field, vw *view, shard uint64)
CacheSize: DefaultCacheSize,
holder: holder,
stats: stats.NopStatsClient,
}
return f
}
@ -405,7 +400,7 @@ func (f *fragment) unprotectedSetBit(tx Tx, rowID, columnID uint64) (changed boo
f.cache.Add(rowID, n)
}
f.stats.Count(MetricSetBit, 1, 1.0)
CounterSetBit.Inc()
return changed, nil
}
@ -454,7 +449,7 @@ func (f *fragment) unprotectedClearBit(tx Tx, rowID, columnID uint64) (changed b
f.cache.Add(rowID, n)
}
f.stats.Count(MetricClearBit, 1, 1.0)
CounterClearBit.Inc()
return changed, nil
}
@ -510,7 +505,7 @@ func (f *fragment) unprotectedSetRow(tx Tx, row *Row, rowID uint64) (changed boo
}
}
f.stats.Count("setRow", 1, 1.0)
CounterSetRow.Inc()
return changed, nil
}
@ -1713,23 +1708,23 @@ func (p parallelSlices) Swap(i, j int) {
// operations to the op log.
func (f *fragment) importPositions(tx Tx, set, clear []uint64, rowSet map[uint64]struct{}) error {
if len(set) > 0 {
f.stats.Count(MetricImportingN, int64(len(set)), 1)
CounterImportingN.Add(float64(len(set)))
// TODO benchmark Add/RemoveN behavior with sorted/unsorted positions
changedN, err := tx.Add(f.index(), f.field(), f.view(), f.shard, set...)
if err != nil {
return errors.Wrap(err, "adding positions")
}
f.stats.Count(MetricImportedN, int64(changedN), 1)
CounterImportedN.Add(float64(changedN))
}
if len(clear) > 0 {
f.stats.Count(MetricClearingN, int64(len(clear)), 1)
CounterClearingingN.Add(float64(len(clear)))
changedN, err := tx.Remove(f.index(), f.field(), f.view(), f.shard, clear...)
if err != nil {
return errors.Wrap(err, "clearing positions")
}
f.stats.Count(MetricClearedN, int64(changedN), 1)
CounterClearedN.Add(float64(changedN))
}
return f.updateCaching(tx, rowSet)
}

12
go.mod
View file

@ -4,11 +4,11 @@ replace github.com/go-avro/avro => github.com/pilosa/avro v0.0.0-20200626214113-
replace github.com/gomem/gomem => github.com/tgruben/gomem v0.0.0-20221021111114-79fdc77dcf61
replace robpike.io/ivy => github.com/tgruben/ivy v0.0.0-20221107170120-634b546dcdac
replace robpike.io/ivy => github.com/tgruben/ivy v0.0.0-20230111144143-b80a659caeaf
require (
github.com/CAFxX/gcnotifier v0.0.0-20220409005548-0153238b886a
github.com/DataDog/datadog-go v4.8.3+incompatible
github.com/DataDog/datadog-go v4.8.3+incompatible // indirect
github.com/HdrHistogram/hdrhistogram-go v1.1.2 // indirect
github.com/Microsoft/go-winio v0.5.2 // indirect
github.com/alexbrainman/odbc v0.0.0-20211220213544-9c9a2e61c5e2
@ -43,8 +43,8 @@ require (
github.com/opentracing/opentracing-go v1.2.0
github.com/pelletier/go-toml v1.9.5
github.com/pkg/errors v0.9.1
github.com/prometheus/client_golang v1.12.2
github.com/prometheus/client_model v0.2.0
github.com/prometheus/client_golang v1.14.0
github.com/prometheus/client_model v0.3.0
github.com/prometheus/prom2json v1.3.1
github.com/rakyll/statik v0.1.7
github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
@ -158,8 +158,8 @@ require (
github.com/pierrec/lz4/v4 v4.1.15 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
github.com/prometheus/common v0.33.0 // indirect
github.com/prometheus/procfs v0.7.3 // indirect
github.com/prometheus/common v0.37.0 // indirect
github.com/prometheus/procfs v0.8.0 // indirect
github.com/sirupsen/logrus v1.7.0 // indirect
github.com/soheilhy/cmux v0.1.5 // indirect
github.com/spf13/afero v1.6.0 // indirect

12
go.sum
View file

@ -920,6 +920,8 @@ github.com/prometheus/client_golang v1.11.1/go.mod h1:Z6t4BnS23TR94PD6BsDNk8yVqr
github.com/prometheus/client_golang v1.12.1/go.mod h1:3Z9XVyYiZYEO+YQWt3RD2R3jrbd179Rt297l4aS6nDY=
github.com/prometheus/client_golang v1.12.2 h1:51L9cDoUHVrXx4zWYlcLQIZ+d+VXHgqnYKkIuq4g/34=
github.com/prometheus/client_golang v1.12.2/go.mod h1:3Z9XVyYiZYEO+YQWt3RD2R3jrbd179Rt297l4aS6nDY=
github.com/prometheus/client_golang v1.14.0 h1:nJdhIvne2eSX/XRAFV9PcvFFRbrjbcTUj0VP62TMhnw=
github.com/prometheus/client_golang v1.14.0/go.mod h1:8vpkKitgIVNcqrRBWh1C4TIUQgYNtG/XQE4E/Zae36Y=
github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
github.com/prometheus/client_model v0.0.0-20190115171406-56726106282f/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
@ -927,6 +929,8 @@ github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:
github.com/prometheus/client_model v0.1.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/prometheus/client_model v0.2.0 h1:uq5h0d+GuxiXLJLNABMgp2qUWDPiLvgCzz2dUR+/W/M=
github.com/prometheus/client_model v0.2.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4=
github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w=
github.com/prometheus/common v0.0.0-20181113130724-41aa239b4cce/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro=
github.com/prometheus/common v0.0.0-20181126121408-4724e9255275/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro=
github.com/prometheus/common v0.2.0/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
@ -939,6 +943,8 @@ github.com/prometheus/common v0.26.0/go.mod h1:M7rCNAaPfAosfx8veZJCuw84e35h3Cfd9
github.com/prometheus/common v0.32.1/go.mod h1:vu+V0TpY+O6vW9J44gczi3Ap/oXXR10b+M/gUGO4Hls=
github.com/prometheus/common v0.33.0 h1:rHgav/0a6+uYgGdNt3jwz8FNSesO/Hsang3O0T9A5SE=
github.com/prometheus/common v0.33.0/go.mod h1:gB3sOl7P0TvJabZpLY5uQMpUqRCPPCyRLCZYc7JZTNE=
github.com/prometheus/common v0.37.0 h1:ccBbHCgIiT9uSoFY0vX8H3zsNR5eLt17/RQLUvn8pXE=
github.com/prometheus/common v0.37.0/go.mod h1:phzohg0JFMnBEFGxTDbfu3QyL5GI8gTQJFhYO5B3mfA=
github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
github.com/prometheus/procfs v0.0.0-20181204211112-1dc9a6cbc91a/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
github.com/prometheus/procfs v0.0.0-20190117184657-bf6a532e95b1/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
@ -950,6 +956,8 @@ github.com/prometheus/procfs v0.3.0/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4O
github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
github.com/prometheus/procfs v0.7.3 h1:4jVXhlkAyzOScmCkXBTOLRLTz8EeU+eyjrwB/EPq0VU=
github.com/prometheus/procfs v0.7.3/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
github.com/prometheus/procfs v0.8.0 h1:ODq8ZFEaYeCaZOJlZZdJA2AbQR98dSHSM1KW/You5mo=
github.com/prometheus/procfs v0.8.0/go.mod h1:z7EfXMXOkbkqb9IINtpCn86r/to3BnA0uaxHdg830/4=
github.com/prometheus/prom2json v1.3.1 h1:OogL5hsrJpLPz3jZ4LPz4sJRTtADzViCNRQoqrzUQvk=
github.com/prometheus/prom2json v1.3.1/go.mod h1:A8Oy9aiQx4wrJY9ya1i4nHOySGmkVp5EO0aU1iSJR+g=
github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40TwIPHuXU=
@ -1058,8 +1066,8 @@ github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69
github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ=
github.com/tgruben/gomem v0.0.0-20221021111114-79fdc77dcf61 h1:3RJ3IN/m4w4B7PeYF8PPOdGLHqu045LIxtymF4lQO7g=
github.com/tgruben/gomem v0.0.0-20221021111114-79fdc77dcf61/go.mod h1:avNE+ynGJYvQNY+/5Gk6aHxtlOtt2aO6V9ShcONBvvI=
github.com/tgruben/ivy v0.0.0-20221107170120-634b546dcdac h1:xVEKycwTG+9q18T1zA+ESrkIMIMQMmlsGOFirh5VxT8=
github.com/tgruben/ivy v0.0.0-20221107170120-634b546dcdac/go.mod h1:/COPfnSdd23BhmpDXFKsYPk1kK2sJpUr21TsZ1nsgcg=
github.com/tgruben/ivy v0.0.0-20230111144143-b80a659caeaf h1:+bnHPov8gMIztSP1zxMbUXei8IiVNXbQre4+u9lpkQ0=
github.com/tgruben/ivy v0.0.0-20230111144143-b80a659caeaf/go.mod h1:/COPfnSdd23BhmpDXFKsYPk1kK2sJpUr21TsZ1nsgcg=
github.com/tidwall/btree v0.3.0/go.mod h1:huei1BkDWJ3/sLXmO+bsCNELL+Bp2Kks9OLyQFkzvA8=
github.com/tidwall/btree v1.1.0/go.mod h1:TzIRzen6yHbibdSfK6t8QimqbUnoxUSrZfeW7Uob0q4=
github.com/tidwall/buntdb v1.2.0/go.mod h1:XLza/dhlwzO6dc5o/KWor4kfZSt3BP8QV+77ZMKfI58=

15
holder.go
View file

@ -18,7 +18,6 @@ import (
"github.com/featurebasedb/featurebase/v3/logger"
rbfcfg "github.com/featurebasedb/featurebase/v3/rbf/cfg"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/storage"
"github.com/featurebasedb/featurebase/v3/testhook"
"github.com/featurebasedb/featurebase/v3/vprint"
@ -80,9 +79,6 @@ type Holder struct {
wg sync.WaitGroup
closing chan struct{}
// Stats
Stats stats.StatsClient
// Data directory path.
path string
@ -257,7 +253,6 @@ type HolderConfig struct {
Schemator disco.Schemator
Sharder disco.Sharder
CacheFlushInterval time.Duration
StatsClient stats.StatsClient
Logger logger.Logger
StorageConfig *storage.Config
@ -282,7 +277,6 @@ func DefaultHolderConfig() *HolderConfig {
Schemator: disco.NewInMemSchemator(),
Sharder: disco.InMemSharder,
CacheFlushInterval: defaultCacheFlushInterval,
StatsClient: stats.NopStatsClient,
Logger: logger.NopLogger,
StorageConfig: storage.NewDefaultConfig(),
RBFConfig: rbfcfg.NewDefaultConfig(),
@ -324,7 +318,6 @@ func NewHolder(path string, cfg *HolderConfig) *Holder {
broadcaster: NopBroadcaster,
partitionN: cfg.PartitionN,
Stats: cfg.StatsClient,
cacheFlushInterval: cfg.CacheFlushInterval,
OpenTranslateStore: cfg.OpenTranslateStore,
OpenTranslateReader: cfg.OpenTranslateReader,
@ -527,8 +520,6 @@ func (h *Holder) Open() error {
// Check if deletion was in progress when server was shutdown
h.processDeleteInflight()
h.Stats.Open()
h.opened.Close()
_ = testhook.Opened(h.Auditor, h, nil)
@ -629,8 +620,6 @@ func (h *Holder) Close() error {
fmt.Printf("%v\n", globalCallStats.report())
}
h.Stats.Close()
// Notify goroutines of closing and wait for completion.
close(h.closing)
h.wg.Wait()
@ -1157,7 +1146,6 @@ func (h *Holder) newIndex(path, name string) (*Index, error) {
if err != nil {
return nil, err
}
index.Stats = h.Stats.WithTags(fmt.Sprintf("index:%s", index.Name()))
index.broadcaster = h.broadcaster
index.serializer = h.serializer
index.OpenTranslateStore = h.OpenTranslateStore
@ -1330,9 +1318,6 @@ type holderSyncer struct {
syncers errgroup.Group
// Stats
Stats stats.StatsClient
// Signals that the sync should stop.
Closing <-chan struct{}
}

110
http_handler.go
View file

@ -9,7 +9,6 @@ import (
"encoding/gob"
"encoding/hex"
"encoding/json"
"expvar"
"fmt"
"io"
"math"
@ -45,7 +44,9 @@ import (
"github.com/featurebasedb/featurebase/v3/storage"
"github.com/featurebasedb/featurebase/v3/tracing"
"github.com/pkg/errors"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/featurebasedb/featurebase/v3/wireprotocol"
dto "github.com/prometheus/client_model/go"
"github.com/prometheus/prom2json"
uuid "github.com/satori/go.uuid"
@ -402,8 +403,7 @@ func (h *Handler) collectStats(next http.Handler) http.Handler {
next.ServeHTTP(w, r)
dur := time.Since(t)
statsTags := make([]string, 0, 5)
isSlow := "false"
longQueryTime := h.api.LongQueryTime()
if longQueryTime > 0 && dur > longQueryTime {
queryRequest := r.Context().Value(contextKeyQueryRequest)
@ -414,31 +414,27 @@ func (h *Handler) collectStats(next http.Handler) http.Handler {
}
h.logger.Printf("HTTP query duration %v exceeds %v: %s %s %s", dur, longQueryTime, r.Method, r.URL.String(), queryString)
statsTags = append(statsTags, "slow:true")
} else {
statsTags = append(statsTags, "slow:false")
isSlow = "true"
}
where := ""
pathParts := strings.Split(r.URL.Path, "/")
if externalPrefixFlag[pathParts[1]] {
statsTags = append(statsTags, "where:external")
where = "external"
} else {
statsTags = append(statsTags, "where:internal")
where = "internal"
}
statsTags = append(statsTags, "useragent:"+r.UserAgent())
path, err := mux.CurrentRoute(r).GetPathTemplate()
if err == nil {
statsTags = append(statsTags, "path:"+path)
}
statsTags = append(statsTags, "method:"+r.Method)
stats := h.api.StatsWithTags(statsTags)
if stats != nil {
stats.Timing(MetricHTTPRequest, dur, 0.1)
if err != nil {
path = ""
}
SummaryHttpRequests.With(prometheus.Labels{
"method": r.Method,
"path": path,
"slow": isSlow,
"useragent": r.UserAgent(),
"where": where,
}).Observe(dur.Seconds())
})
}
@ -506,7 +502,6 @@ func newRouter(handler *Handler) http.Handler {
// TODO: figure out how to protect these if needed
router.PathPrefix("/debug/pprof/").Handler(http.DefaultServeMux).Methods("GET")
router.PathPrefix("/debug/fgprof").Handler(fgprof.Handler()).Methods("GET")
router.Handle("/debug/vars", expvar.Handler()).Methods("GET")
router.Handle("/metrics", promhttp.Handler())
router.HandleFunc("/metrics.json", handler.chkAuthZ(handler.handleGetMetricsJSON, authz.Admin)).Methods("GET").Name("GetMetricsJSON")
@ -552,6 +547,8 @@ func newRouter(handler *Handler) http.Handler {
if handler.sqlEnabled {
router.HandleFunc("/sql", handler.chkAuthZ(handler.handlePostSQL, authz.Admin)).Methods("POST").Name("PostSQL")
}
// internal endpoint
router.HandleFunc("/sql", handler.chkAuthZ(handler.handlePostSQLPlanOperator, authz.Admin)).Headers("X-FeatureBase-Plan-Operator", "").Methods("POST").Name("PostSQLPlanOperator")
router.HandleFunc("/query-history", handler.chkAuthZ(handler.handleGetPastQueries, authz.Admin)).Methods("GET").Name("GetPastQueries")
router.HandleFunc("/version", handler.handleGetVersion).Methods("GET").Name("GetVersion")
@ -1393,10 +1390,70 @@ func (h *Handler) writeBadRequest(w http.ResponseWriter, r *http.Request, err er
}
}
// handlePostSQL handles /sql requests.
func (h *Handler) handlePostSQL(w http.ResponseWriter, r *http.Request) {
includePlan := false
// handlePostSQLOperator handles an internal sql3 plan operator execution request
// these requests come from other nodes in the cluster
// handlePostSQLOperator will 'rehydrate' a plan operator and return data in the
// featurebase wire format for effciency
// we do not track these requests as user requests
// TODO(pok) - thus is there anything we need here to align with how we do this for other nodes
func (h *Handler) handlePostSQLPlanOperator(w http.ResponseWriter, r *http.Request) {
writeError := func(err error) {
if err != nil {
w.Write(wireprotocol.WriteError(err))
}
}
// always finish with a done message
defer w.Write(wireprotocol.WriteDone())
ctx := r.Context()
rootOperator, err := h.api.RehydratePlanOperator(ctx, r.Body)
if err != nil {
writeError(err)
return
}
// get a query iterator.
iter, err := rootOperator.Iterator(ctx, nil)
if err != nil {
writeError(err)
return
}
// read schema & write to response.
columns := rootOperator.Schema()
b, err := wireprotocol.WriteSchema(columns)
if err != nil {
writeError(err)
return
}
w.Write(b)
var rowErr error
var currentRow types.Row
var nextErr error
for currentRow, nextErr = iter.Next(ctx); nextErr == nil; currentRow, nextErr = iter.Next(ctx) {
b, err := wireprotocol.WriteRow(currentRow, columns)
if err != nil {
rowErr = err
break
}
w.Write(b)
}
if nextErr != nil && nextErr != types.ErrNoMoreRows {
rowErr = nextErr
}
writeError(rowErr)
}
// handlePostSQL handles /sql requests
// supports a ?plan=true|false parameter to send back the plan in the
// query response
func (h *Handler) handlePostSQL(w http.ResponseWriter, r *http.Request) {
includePlan := false
includePlanValue := r.URL.Query().Get("plan")
if len(includePlanValue) > 0 {
var err error
@ -1407,6 +1464,7 @@ func (h *Handler) handlePostSQL(w http.ResponseWriter, r *http.Request) {
}
}
// get the body
b, err := io.ReadAll(r.Body)
if err != nil {
h.writeBadRequest(w, r, err)
@ -1420,6 +1478,9 @@ func (h *Handler) handlePostSQL(w http.ResponseWriter, r *http.Request) {
// put the requestId in the context
ctx := fbcontext.WithRequestID(r.Context(), requestID.String())
// update the counter for requests
PerfCounterSQLRequestSec.Add(1)
// Write response back to client.
w.Header().Set("Content-Type", "application/json")
@ -4146,6 +4207,7 @@ func (h *Handler) handleGetHealth(w http.ResponseWriter, r *http.Request) {
func init() {
gob.Register(arrow.PrimitiveTypes.Int64)
gob.Register(arrow.PrimitiveTypes.Float64)
gob.Register(arrow.BinaryTypes.String)
}
// EXPERIMENTAL API MAY CHANGE

2
idk/Makefile
View file

@ -231,7 +231,7 @@ TPKG ?= ./...
test-run: testenv vendor
$(DOCKER_COMPOSE) build idk-test
$(DOCKER_COMPOSE) run -T idk-test bash -c "set -o pipefail; go test -v -mod=vendor -tags=odbc,dynamic $(TPKG) -covermode=atomic -coverpkg=$(TPKG) -coverprofile=/testdata/$(PROJECT)_base_coverage.out"
$(DOCKER_COMPOSE) run -T idk-test /go/src/github.com/molecula/featurebase/idk/reingest_test.sh
$(DOCKER_COMPOSE) run -T idk-test /go/src/github.com/featurebasedb/featurebase/idk/reingest_test.sh
test-run-race: testenv vendor

29
idk/common/configMap.go
View file

@ -44,6 +44,8 @@ func LaunchKafkaEventConfirmer(producer *confluent.Producer, finished *int32, it
return doneChan
}
// For a list of confluent consumer configuraiton options, go here:
// https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md
func SetupConfluent(m *idk.ConfluentCommand) (*confluent.ConfigMap, error) {
var err error
configMap := &confluent.ConfigMap{}
@ -188,5 +190,32 @@ func SetupConfluent(m *idk.ConfluentCommand) (*confluent.ConfigMap, error) {
}
}
if m.KafkaGroupInstanceId != "" {
err = configMap.SetKey("group.instance.id", m.KafkaGroupInstanceId)
if err != nil {
return nil, err
}
}
if m.KafkaMaxPollInterval != "" {
err = configMap.SetKey("max.poll.interval.ms", m.KafkaMaxPollInterval)
if err != nil {
return nil, err
}
}
if m.KafkaSessionTimeout != "" {
err = configMap.SetKey("session.timeout.ms", m.KafkaSessionTimeout)
if err != nil {
return nil, err
}
}
if m.KafkaSocketKeepaliveEnable != "" {
err = configMap.SetKey("socket.keepalive.enable", m.KafkaSocketKeepaliveEnable)
if err != nil {
return nil, err
}
}
return configMap, nil
}

89
idk/ingest.go
View file

@ -34,9 +34,7 @@ import (
"github.com/featurebasedb/featurebase/v3/idk/mds"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/prometheus"
proto "github.com/featurebasedb/featurebase/v3/proto"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/pkg/errors"
prom "github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
@ -55,7 +53,6 @@ const (
)
// TODO Jaeger
// TODO Prometheus
// Main holds all config for general ingest
type Main struct {
@ -132,7 +129,6 @@ type Main struct {
newNexter func(c int) (IDAllocator, error)
ra RangeAllocator
stats stats.StatsClient
metricsServer *http.Server
log logger.Logger
@ -185,9 +181,13 @@ type ConfluentCommand struct {
KafkaSslEndpointIdentificationAlgorithm string `help:"The endpoint identification algorithm used by clients to validate server host name (ssl.endpoint.identification.algorithm) "`
KafkaEnableSslCertificateVerification bool `help:"(enable.ssl.certificate.verification)"`
KafkaSocketTimeoutMs int `help:"(socket.timeout.ms)"`
KafkaSocketKeepaliveEnable string `help:"The (socket.keepalive.enable) kafka consumer configuration"`
KafkaClientId string `help:"(client.id)"`
KafkaDebug string `help:"Kafka debug string (debug)"`
KafkaClientId string `help:"(client.id)"`
KafkaDebug string `help:"The (debug) kafka consumer configuration. A comma-separated list of debug contexts to enable. Detailed Consumer: consumer,cgrp,topic,fetch. Set to 'all' for most verbose option."`
KafkaMaxPollInterval string `help:"The (max.poll.interval.ms) kafka consumer configuration. The max time the consumer can go without polling the broker. Consumer exits after this timeout."`
KafkaSessionTimeout string `help:"The (session.timeout.ms) kafka consumer configuration. The max time the consumer can go without sending a heartbeat to the broker"`
KafkaGroupInstanceId string `help:"The (group.instance.id) kafka consumer configuration."`
KafkaSaslUsername string `help:"SASL authentication username (sasl.username)"`
KafkaSaslPassword string `help:"SASL authentication password (sasl.password)"`
@ -218,7 +218,7 @@ func NewMain() *Main {
Concurrency: 1,
CacheLength: 64,
PackBools: "bools",
Namespace: "ingester",
Namespace: "ingester", // this is now ignored and hardcoded in metrics.go
IDAllocKeyPrefix: "ingest",
UseShardTransactionalEndpoint: os.Getenv("IDK_DEFAULT_SHARD_TRANSACTIONAL") != "",
@ -228,8 +228,6 @@ func NewMain() *Main {
SchemaManager: NopSchemaManager,
stats: stats.NopStatsClient,
log: logger.NewStandardLogger(os.Stderr),
}
}
@ -315,7 +313,7 @@ func (m *Main) runIngester(c int, l *msgCounter) error {
err = source.Close()
if err != nil {
if m.log != nil {
m.log.Printf("error on close %v", err)
m.log.Errorf("Closing source: %v", err)
}
}
}()
@ -457,7 +455,6 @@ initialFetch:
if v, ok := source.(Metadata); ok {
m.log.Printf("new schema - subject: %#v; version: %d; schema: %#v",
v.SchemaSubject(), v.SchemaVersion(), v.SchemaSchema())
// m.log.Printf("new schema: %#v", v.SchemaMetadata())
} else {
m.log.Printf("new schema: %#v", schema)
}
@ -465,7 +462,7 @@ initialFetch:
if err != nil {
return errors.Wrap(err, "batchFromSchema")
}
m.stats.Count(MetricIngesterSchemaChanges, 1, 1)
CounterIngesterSchemaChanges.Inc()
csvSlice = make([]string, len(schema))
if m.csvWriter != nil {
for i := range schema {
@ -571,7 +568,7 @@ initialFetch:
// skip bad rows only
if !rowHasError {
err = batch.Add(*row)
m.stats.Count(MetricIngesterRowsAdded, 1, 1)
CounterIngesterRowsAdded.Inc()
}
if err == pilosabatch.ErrBatchNowFull || err == pilosabatch.ErrBatchNowStale {
@ -958,7 +955,7 @@ func (m *Main) commitRecord(ctx context.Context, rec Record, limitCounter *msgCo
return errors.Wrap(err, "committing")
}
limitCounter.Increment(numRecords)
m.stats.Count(MetricCommittedRecords, int64(numRecords), 1)
CounterCommittedRecords.Add(float64(numRecords))
return nil
}
@ -972,7 +969,7 @@ func (m *Main) NewLookupClient() (*PostgresClient, error) {
func (m *Main) setupClient() (*tls.Config, error) {
var tlsConfig *tls.Config
var err error
var opts = []pilosaclient.ClientOption{pilosaclient.OptClientStatsClient(m.stats)}
var opts = []pilosaclient.ClientOption{}
if m.TLS.CertificatePath != "" {
tlsConfig, err = GetTLSConfig(&m.TLS, m.Log())
if err != nil {
@ -1012,10 +1009,10 @@ func (m *Main) setupClient() (*tls.Config, error) {
// MDS doesn't auto-create a table based on IDK ingest; the table must
// already exist.
mdsClient := mdsclient.New(dax.Address(m.MDSAddress), m.log)
qual := dax.NewTableQualifier(m.OrganizationID, m.DatabaseID)
qtid, err := mdsClient.TableID(ctx, qual, m.TableName)
qdbid := dax.NewQualifiedDatabaseID(m.OrganizationID, m.DatabaseID)
qtid, err := mdsClient.TableID(ctx, qdbid, m.TableName)
if err != nil {
return nil, errors.Wrapf(err, "getting table id: qual: %s, table name: %s", qual, m.TableName)
return nil, errors.Wrapf(err, "getting table id: qual: %s, table name: %s", qdbid, m.TableName)
}
qtbl, err := mdsClient.Table(ctx, qtid)
if err != nil {
@ -1024,7 +1021,7 @@ func (m *Main) setupClient() (*tls.Config, error) {
m.Qtbl = qtbl
m.Index = string(qtbl.Key())
m.SchemaManager = mds.NewSchemaManager(dax.Address(m.MDSAddress), qual, m.log)
m.SchemaManager = mds.NewSchemaManager(dax.Address(m.MDSAddress), qdbid, m.log)
m.NewImporterFn = func() pilosacore.Importer {
return mds.NewImporter(mdsClient, mdsClient, qtbl.Qualifier(), &qtbl.Table)
@ -1037,34 +1034,24 @@ func (m *Main) setupClient() (*tls.Config, error) {
}
func (m *Main) setupStats() error {
if m.Stats != "" {
opts := []prometheus.ClientOption{prometheus.OptClientNamespace(m.Namespace)}
m.stats, _ = prometheus.NewPrometheusClient(opts...) // ignore error that must be nil
mux := http.NewServeMux()
// reg := prom.NewRegistry() // TODO switch to this once pilosa PrometheusClient is fixed and doesn't use the global registry internally.
// also change prom.DefaultGatherer to be "reg" at that time
reg := prom.DefaultRegisterer
promHandler := promhttp.InstrumentMetricHandler(reg, promhttp.HandlerFor(prom.DefaultGatherer, promhttp.HandlerOpts{}))
mux.Handle("/metrics", promHandler)
mux.Handle("/metrics.json", metricsJSONHandler{metricsURI: "http://" + m.Stats + "/metrics"})
m.metricsServer = &http.Server{Addr: m.Stats, Handler: mux}
ln, err := net.Listen("tcp", m.Stats)
if err != nil {
return errors.Wrapf(err, "listen for metrics on '%s'", m.Stats)
}
go func() {
m.log.Printf("Serving Prometheus metrics with namespace \"%s\" at %v/metrics\n", m.Namespace, m.Stats)
err = m.metricsServer.Serve(ln)
if err != http.ErrServerClosed {
m.log.Printf("serve metrics on '%s': %v", m.Stats, err)
}
}()
mux := http.NewServeMux()
promHandler := promhttp.InstrumentMetricHandler(prom.DefaultRegisterer, promhttp.HandlerFor(prom.DefaultGatherer, promhttp.HandlerOpts{}))
mux.Handle("/metrics", promHandler)
mux.Handle("/metrics.json", metricsJSONHandler{metricsURI: "http://" + m.Stats + "/metrics"})
m.metricsServer = &http.Server{Addr: m.Stats, Handler: mux}
ln, err := net.Listen("tcp", m.Stats)
if err != nil {
return errors.Wrapf(err, "listen for metrics on '%s'", m.Stats)
}
return nil
go func() {
m.log.Printf("Serving Prometheus metrics with namespace \"%s\" at %v/metrics\n", m.Namespace, m.Stats)
err = m.metricsServer.Serve(ln)
if err != http.ErrServerClosed {
m.log.Printf("serve metrics on '%s': %v", m.Stats, err)
}
}()
return nil
}
type metricsJSONHandler struct {
@ -1218,7 +1205,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error {
if err != nil {
return errors.Wrap(err, "clearing bools")
}
m.stats.CountWithCustomTags(MetricDeleterRowsAdded, 1, 1, []string{"type:packed-bool"})
CounterDeleterRowsAdded.With(prom.Labels{"type": "packed-bool"}).Inc()
continue
} else {
fieldName = directive
@ -1267,7 +1254,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error {
if err != nil {
return errors.Wrap(err, "clearing set")
}
m.stats.CountWithCustomTags(MetricDeleterRowsAdded, 1, 1, []string{"type:set"})
CounterDeleterRowsAdded.With(prom.Labels{"type": "set"}).Inc()
case pilosaclient.FieldTypeMutex:
if val == "" {
continue
@ -1278,7 +1265,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error {
if err != nil {
return errors.Wrap(err, "clearing mutex")
}
m.stats.CountWithCustomTags(MetricDeleterRowsAdded, 1, 1, []string{"type:mutex"})
CounterDeleterRowsAdded.With(prom.Labels{"type": "mutex"}).Inc()
case pilosaclient.FieldTypeBool:
_, err := client.Query(index.BatchQuery(
field.Clear(0, recordID),
@ -1287,19 +1274,19 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error {
if err != nil {
return errors.Wrap(err, "clearing bool")
}
m.stats.CountWithCustomTags(MetricDeleterRowsAdded, 1, 1, []string{"type:bool"})
CounterDeleterRowsAdded.With(prom.Labels{"type": "bool"}).Inc()
case pilosaclient.FieldTypeInt:
_, err := client.Query(field.Clear(0, recordID))
if err != nil {
return errors.Wrap(err, "clearing int")
}
m.stats.CountWithCustomTags(MetricDeleterRowsAdded, 1, 1, []string{"type:int"})
CounterDeleterRowsAdded.With(prom.Labels{"type": "int"}).Inc()
case pilosaclient.FieldTypeDecimal:
_, err := client.Query(field.Clear(0, recordID))
if err != nil {
return errors.Wrap(err, "clearing decimal")
}
m.stats.CountWithCustomTags(MetricDeleterRowsAdded, 1, 1, []string{"type:decimal"})
CounterDeleterRowsAdded.With(prom.Labels{"type": "decimal"}).Inc()
case pilosaclient.FieldTypeTime:
return errors.Errorf("deletion on time fields unimplemented")
default:

27
idk/ingest_test.go
View file

@ -58,7 +58,7 @@ func configureTestFlagsMDS(main *Main, address dax.Address, qtbl *dax.QualifiedT
mdsClient := mdsclient.New(dax.Address(address), logger.StderrLogger)
main.NewImporterFn = func() pilosa.Importer {
return mds.NewImporter(mdsClient, mdsClient, qtbl.TableQualifier, &qtbl.Table)
return mds.NewImporter(mdsClient, mdsClient, qtbl.QualifiedDatabaseID, &qtbl.Table)
}
}
@ -525,7 +525,7 @@ func TestIngesterServesPrometheusEndpoint(t *testing.T) {
if err != nil {
t.Errorf("read error: %v", err)
}
if strings.Contains(string(contents), MetricIngesterRowsAdded) {
if !strings.Contains(string(contents), MetricIngesterRowsAdded) {
t.Errorf("metric name missing: %v", MetricIngesterRowsAdded)
}
close(records)
@ -1756,6 +1756,24 @@ func TestBatchTargetMDS(t *testing.T) {
orgID := dax.OrganizationID("acme")
dbID := dax.DatabaseID("db1")
mdsClient := mdsclient.New(mdsAddress, logger.StderrLogger)
ctx := context.Background()
// Create the database.
qdb := &dax.QualifiedDatabase{
OrganizationID: orgID,
Database: dax.Database{
ID: dbID,
Name: "dbname1",
Options: dax.DatabaseOptions{
WorkersMin: 1,
WorkersMax: 1,
},
},
}
mdsClient.CreateDatabase(ctx, qdb)
t.Run("FieldTypes", func(t *testing.T) {
tests := []struct {
fieldType dax.BaseType
@ -1859,14 +1877,11 @@ func TestBatchTargetMDS(t *testing.T) {
}
qtbl := dax.NewQualifiedTable(
dax.NewTableQualifier(orgID, dbID),
dax.NewQualifiedDatabaseID(orgID, dbID),
tbl,
)
ctx := context.Background()
// Create the table in MDS Schemar.
mdsClient := mdsclient.New(mdsAddress, logger.StderrLogger)
if err := mdsClient.CreateTable(ctx, qtbl); err != nil {
t.Fatalf("creating table: %v", err)
}

20
idk/kafka/cmd.go
View file

@ -13,7 +13,8 @@ type Main struct {
Group string `help:"Kafka group."`
Topics []string `help:"Kafka topics to read from."`
Timeout time.Duration `help:"Time to wait for more records from Kafka before flushing a batch. 0 to disable."`
SkipOld bool `short:"" help:"Skip to the most recent Kafka message rather than starting at the beginning."`
SkipOld bool `short:"" help:"False sets kafka consumer configuration auto.offset.reset to earliest, True sets it to latest."`
ConsumerCloseTimeout int `help:"The amount of time in seconds to wait for the consumer to close properly."`
}
func NewMain() (*Main, error) {
@ -22,13 +23,16 @@ func NewMain() (*Main, error) {
ConfluentCommand: idk.ConfluentCommand{
KafkaBootstrapServers: []string{"localhost:9092"},
},
Group: "defaultgroup",
Topics: []string{"defaulttopic"},
Timeout: time.Second,
Group: "defaultgroup",
Topics: []string{"defaulttopic"},
Timeout: time.Second,
ConsumerCloseTimeout: 30,
}
m.SchemaRegistryURL = "http://" + defaultRegistryHost
m.OffsetMode = true
m.Main.Namespace = "ingester_kafka"
//m.Main.OffsetMode = m.OffsetMode
m.OffsetMode = true
m.NewSource = func() (idk.Source, error) {
source := NewSource()
source.KafkaBootstrapServers = m.KafkaBootstrapServers
@ -43,6 +47,12 @@ func NewMain() (*Main, error) {
source.SchemaRegistryUsername = m.SchemaRegistryUsername
source.SchemaRegistryPassword = m.SchemaRegistryPassword
source.Verbose = m.Verbose
source.KafkaMaxPollInterval = m.KafkaMaxPollInterval
source.KafkaSessionTimeout = m.KafkaSessionTimeout
source.KafkaGroupInstanceId = m.KafkaGroupInstanceId
source.KafkaDebug = m.KafkaDebug
source.KafkaSocketKeepaliveEnable = m.KafkaSocketKeepaliveEnable
source.consumerCloseTimeout = m.ConsumerCloseTimeout
if err := source.Open(); err != nil {
return nil, errors.Wrap(err, "opening source")

41
idk/kafka/source.go
View file

@ -33,14 +33,15 @@ import (
// achieve concurrency, create multiple Sources.
type Source struct {
idk.ConfluentCommand
Topics []string
Group string
Log logger.Logger
Timeout time.Duration
SkipOld bool
Verbose bool
schema Schema
TLS idk.TLSConfig
Topics []string
Group string
Log logger.Logger
Timeout time.Duration
SkipOld bool
Verbose bool
schema Schema
TLS idk.TLSConfig
consumerCloseTimeout int
spoolBase uint64
spool []confluent.TopicPartition
@ -257,7 +258,9 @@ func (s *Source) CommitMessages(recs []confluent.TopicPartition) ([]confluent.To
return s.client.CommitOffsets(recs)
}
// Open initializes the kafka source.
// Open initializes the kafka source. (i.e. creating and configuring a consumer)
// The configuration options for the confluentinc/confluent-kafka-go/kafka
// libarary are: https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md
func (s *Source) Open() error {
cfg, err := common.SetupConfluent(&s.ConfluentCommand)
if err != nil {
@ -439,10 +442,26 @@ func (c *Source) generator() {
func (s *Source) Close() error {
if s.client != nil {
if s.opened { // only close opened sources
var err error
closedReturned := make(chan error, 1)
// send quit message to polling routine & wait for it to exit
s.quit <- struct{}{}
s.wg.Wait()
err := s.client.Close()
s.opened = false
s.Log.Debugf("Trying to close consumer %s...", s.client.String())
go func() {
closedReturned <- s.client.Close()
}()
start := time.Now()
select {
case err = <-closedReturned:
if err == nil {
s.Log.Debugf("Successfully closed consumer %s!", s.client.String())
s.opened = false
}
case <-time.After(time.Duration(s.consumerCloseTimeout * 1000 * 1000 * 1000)):
err = fmt.Errorf("Unable to properly close consumer %s after %f seconds", s.client.String(), time.Since(start).Seconds())
}
return errors.Wrap(err, "closing kafka consumer")
}
}

16
idk/mds/importer.go
View file

@ -21,16 +21,16 @@ type importer struct {
noder dax.Noder
schemar dax.Schemar
mu sync.Mutex
qual dax.TableQualifier
tbl *dax.Table
mu sync.Mutex
qdbid dax.QualifiedDatabaseID
tbl *dax.Table
}
func NewImporter(noder dax.Noder, schemar dax.Schemar, qual dax.TableQualifier, tbl *dax.Table) *importer {
func NewImporter(noder dax.Noder, schemar dax.Schemar, qdbid dax.QualifiedDatabaseID, tbl *dax.Table) *importer {
return &importer{
noder: noder,
schemar: schemar,
qual: qual,
qdbid: qdbid,
tbl: tbl,
}
}
@ -248,8 +248,6 @@ func (m *importer) DoImport(ctx context.Context, tid dax.TableID, fld *dax.Field
return fbClient.DoImport(string(qtbl.Key()), shard, path, data)
}
func (m *importer) StatsTiming(name string, value time.Duration, rate float64) {}
// getQtbl takes a table (TableKey) and sets the local m.qtbl value. When we
// originally set up this type, it was only used by IDK, and the table was known
// at the beginning of the process, so it could be set on this import. But
@ -263,10 +261,10 @@ func (m *importer) getQtbl(ctx context.Context, tid dax.TableID) (*dax.Qualified
defer m.mu.Unlock()
if m.tbl != nil {
return dax.NewQualifiedTable(m.qual, m.tbl), nil
return dax.NewQualifiedTable(m.qdbid, m.tbl), nil
}
qtid := dax.NewQualifiedTableID(m.qual, tid)
qtid := dax.NewQualifiedTableID(m.qdbid, tid)
qtbl, err := m.schemar.TableByID(ctx, qtid)
if err != nil {

8
idk/mds/schemamanager.go
View file

@ -18,14 +18,14 @@ import (
// schemaManager
type schemaManager struct {
client *mdsclient.Client
qual dax.TableQualifier
qdbid dax.QualifiedDatabaseID
logger logger.Logger
}
func NewSchemaManager(mdsAddress dax.Address, qual dax.TableQualifier, logger logger.Logger) *schemaManager {
func NewSchemaManager(mdsAddress dax.Address, qdbid dax.QualifiedDatabaseID, logger logger.Logger) *schemaManager {
return &schemaManager{
client: mdsclient.New(mdsAddress, logger),
qual: qual,
qdbid: qdbid,
logger: logger,
}
}
@ -41,7 +41,7 @@ func (s *schemaManager) Schema() (*featurebase_client.Schema, error) {
// method returns.
schema := featurebase_client.NewSchema()
tables, err := s.client.Tables(context.Background(), s.qual)
tables, err := s.client.Tables(context.Background(), s.qdbid)
if err != nil {
return nil, err
}

44
idk/metrics.go
View file

@ -1,8 +1,52 @@
package idk
import "github.com/prometheus/client_golang/prometheus"
const (
MetricDeleterRowsAdded = "deleter_rows_added_total"
MetricIngesterRowsAdded = "ingester_rows_added_total"
MetricIngesterSchemaChanges = "ingester_schema_changes_total"
MetricCommittedRecords = "committed_records"
)
var CounterIngesterSchemaChanges = prometheus.NewCounter(
prometheus.CounterOpts{
Namespace: "ingester",
Name: MetricIngesterSchemaChanges,
Help: "TODO",
},
)
var CounterIngesterRowsAdded = prometheus.NewCounter(
prometheus.CounterOpts{
Namespace: "ingester",
Name: MetricIngesterRowsAdded,
Help: "TODO",
},
)
var CounterCommittedRecords = prometheus.NewCounter(
prometheus.CounterOpts{
Namespace: "ingester",
Name: MetricCommittedRecords,
Help: "TODO",
},
)
var CounterDeleterRowsAdded = prometheus.NewCounterVec(
prometheus.CounterOpts{
Namespace: "ingester",
Name: MetricDeleterRowsAdded,
Help: "TODO",
},
[]string{
"type",
},
)
func init() {
prometheus.MustRegister(CounterIngesterSchemaChanges)
prometheus.MustRegister(CounterIngesterRowsAdded)
prometheus.MustRegister(CounterCommittedRecords)
prometheus.MustRegister(CounterDeleterRowsAdded)
}

2
importer.go
View file

@ -20,8 +20,6 @@ type Importer interface {
EncodeImportValues(ctx context.Context, tid dax.TableID, fld *dax.Field, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error)
EncodeImport(ctx context.Context, tid dax.TableID, fld *dax.Field, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error)
DoImport(ctx context.Context, tid dax.TableID, fld *dax.Field, shard uint64, path string, data []byte) error
StatsTiming(name string, value time.Duration, rate float64)
}
// Ensure type implements interface.

7
index.go
View file

@ -17,9 +17,9 @@ import (
"github.com/featurebasedb/featurebase/v3/disco"
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/testhook"
"github.com/pkg/errors"
"github.com/prometheus/client_golang/prometheus"
"golang.org/x/sync/errgroup"
)
@ -44,7 +44,6 @@ type Index struct {
broadcaster broadcaster
serializer Serializer
Stats stats.StatsClient
// Passed to field for foreign-index lookup.
holder *Holder
@ -83,7 +82,6 @@ func NewIndex(holder *Holder, path, name string) (*Index, error) {
fields: make(map[string]*Field),
broadcaster: NopBroadcaster,
Stats: stats.NopStatsClient,
holder: holder,
trackExistence: true,
@ -511,7 +509,7 @@ func (i *Index) AvailableShards(localOnly bool) *roaring.Bitmap {
b.UnionInPlace(f.AvailableShards(localOnly))
}
i.Stats.Gauge(MetricMaxShard, float64(b.Max()), 1.0)
GaugeIndexMaxShard.With(prometheus.Labels{"index": i.name}).Set(float64(b.Max()))
return b
}
@ -924,7 +922,6 @@ func (i *Index) newField(path, name string) (*Field, error) {
return nil, err
}
f.idx = i
f.Stats = i.Stats
f.broadcaster = i.broadcaster
f.serializer = i.serializer
f.OpenTranslateStore = i.OpenTranslateStore

2
lattice/README.md
View file

@ -19,4 +19,4 @@ Open [http://localhost:3000](http://localhost:3000) to view it in the browser.
Lattice can be embedded within the Pilosa binary, so the UI is fully accessible directly from the server, reducing operational complexity.
If additional build dependencies `yarn` (`brew install yarn` and `brew upgrade yarn` perhaps) and `statik` (`make install-statik`) are available on your system, running `make generate-statik` before `make install` should produce a Pilosa binary with Lattice embedded. For up to date instructions, check the Pilosa [README](https://github.com/molecula/pilosa#getting-started).
If additional build dependencies `yarn` (`brew install yarn` and `brew upgrade yarn` perhaps) and `statik` (`make install-statik`) are available on your system, running `make generate-statik` before `make install` should produce a Pilosa binary with Lattice embedded. For up to date instructions, check the Featurebase [README](https://github.com/featurebasedb/featurebase#getting-started).

1
lattice/src/services/__mocks__/eventServices.tsx
View file

@ -10,3 +10,4 @@ const pilosa = {
};
module.exports.pilosa = pilosa;
export {};

996
metrics.go
View file

File diff suppressed because it is too large Load diff

203
performancecounters.go Normal file
View file

@ -0,0 +1,203 @@
package pilosa
import (
"sync/atomic"
"github.com/prometheus/client_golang/prometheus"
)
// PerformanceCounter holds data about a performance counter for external consumers
type PerformanceCounter struct {
NameSpace string
SubSystem string
CounterName string
Help string
Value int64
CounterType int64
}
// constants for the counter types
const (
// raw - for when you just want a count of something
CTR_TYPE_RAW = 0
// per second - for when you accumulate counts of things
// a consumer would sample this at intervals to arrive at a delta
// then divide by the time in seconds between the samples to get a
// per-second value
CTR_TYPE_PER_SECOND = 1
// ratio - for when you accumulate a count of something that you
// want to use as a numerator in a ratio calculation
// e.g. 'cache hits' could be a counter of this type and you could
// divide it by a 'cache lookups' counter to get the hit ratio (see below)
CTR_TYPE_RATIO = 2
// ratio base - for when you accumulate a count of something that you
// want to use as a denominator in a ratio calculation
// e.g. 'cache lookups' could be a counter of this type and you could
// use it as the denominator in a division with a 'cache hits' counter
// as the numerator to get the hit ratio
CTR_TYPE_RATIO_BASE = 3
)
type PerformanceCounters struct {
counterValues [5]perfCtrWrapper
}
type perfCtr struct {
nameSpace string
subSystem string
name string
help string
value int64
counterType int64
}
func (p *perfCtr) Add(increment int64) {
atomic.AddInt64(&p.value, increment)
}
type perfCtrWrapper struct {
ctr *perfCtr
fun prometheus.CounterFunc
}
var PerfCounterSQLRequestSec = perfCtr{
nameSpace: "pilosa",
subSystem: "sql_statistics",
name: "sql_requests_sec",
help: "TODO",
value: 0,
counterType: CTR_TYPE_PER_SECOND,
}
var PerfCounterSQLInsertsSec = perfCtr{
nameSpace: "pilosa",
subSystem: "sql_statistics",
name: "sql_inserts_sec",
help: "TODO",
value: 0,
counterType: CTR_TYPE_PER_SECOND,
}
var PerfCounterSQLBulkInsertsSec = perfCtr{
nameSpace: "pilosa",
subSystem: "sql_statistics",
name: "sql_bulk_inserts_sec",
help: "TODO",
value: 0,
counterType: CTR_TYPE_PER_SECOND,
}
var PerfCounterSQLBulkInsertBatchesSec = perfCtr{
nameSpace: "pilosa",
subSystem: "sql_statistics",
name: "sql_bulk_insert_batches_sec",
help: "TODO",
value: 0,
counterType: CTR_TYPE_PER_SECOND,
}
var PerfCounterSQLDeletesSec = perfCtr{
nameSpace: "pilosa",
subSystem: "sql_statistics",
name: "sql_deletes_sec",
help: "TODO",
value: 0,
counterType: CTR_TYPE_PER_SECOND,
}
var PerfCounters *PerformanceCounters = newPerformanceCounters()
func newPerformanceCounters() *PerformanceCounters {
ctrs := &PerformanceCounters{
counterValues: [5]perfCtrWrapper{
{
&PerfCounterSQLRequestSec,
prometheus.NewCounterFunc(
prometheus.CounterOpts{
Namespace: PerfCounterSQLRequestSec.nameSpace,
Subsystem: PerfCounterSQLRequestSec.subSystem,
Name: PerfCounterSQLRequestSec.name,
Help: PerfCounterSQLRequestSec.help,
},
func() float64 {
return float64(atomic.LoadInt64(&PerfCounterSQLRequestSec.value))
}),
},
{
&PerfCounterSQLInsertsSec,
prometheus.NewCounterFunc(
prometheus.CounterOpts{
Namespace: PerfCounterSQLInsertsSec.nameSpace,
Subsystem: PerfCounterSQLInsertsSec.subSystem,
Name: PerfCounterSQLInsertsSec.name,
Help: PerfCounterSQLInsertsSec.help,
},
func() float64 {
return float64(atomic.LoadInt64(&PerfCounterSQLInsertsSec.value))
}),
},
{
&PerfCounterSQLBulkInsertsSec,
prometheus.NewCounterFunc(
prometheus.CounterOpts{
Namespace: PerfCounterSQLBulkInsertsSec.nameSpace,
Subsystem: PerfCounterSQLBulkInsertsSec.subSystem,
Name: PerfCounterSQLBulkInsertsSec.name,
Help: PerfCounterSQLBulkInsertsSec.help,
},
func() float64 {
return float64(atomic.LoadInt64(&PerfCounterSQLBulkInsertsSec.value))
}),
},
{
&PerfCounterSQLBulkInsertBatchesSec,
prometheus.NewCounterFunc(
prometheus.CounterOpts{
Namespace: PerfCounterSQLBulkInsertBatchesSec.nameSpace,
Subsystem: PerfCounterSQLBulkInsertBatchesSec.subSystem,
Name: PerfCounterSQLBulkInsertBatchesSec.name,
Help: PerfCounterSQLBulkInsertBatchesSec.help,
},
func() float64 {
return float64(atomic.LoadInt64(&PerfCounterSQLBulkInsertBatchesSec.value))
}),
},
{
&PerfCounterSQLDeletesSec,
prometheus.NewCounterFunc(
prometheus.CounterOpts{
Namespace: PerfCounterSQLDeletesSec.nameSpace,
Subsystem: PerfCounterSQLDeletesSec.subSystem,
Name: PerfCounterSQLDeletesSec.name,
Help: PerfCounterSQLDeletesSec.help,
},
func() float64 {
return float64(atomic.LoadInt64(&PerfCounterSQLDeletesSec.value))
}),
},
},
}
for _, w := range ctrs.counterValues {
prometheus.MustRegister(w.fun)
}
return ctrs
}
// list all the counters
// we can just read here without locking because if the counters get changed
// midway thru the loop, absent evidence to the contrary, the world will not end
func (p *PerformanceCounters) ListCounters() ([]PerformanceCounter, error) {
result := make([]PerformanceCounter, len(p.counterValues))
for i, c := range p.counterValues {
result[i] = PerformanceCounter{
NameSpace: c.ctr.nameSpace,
SubSystem: c.ctr.subSystem,
CounterName: c.ctr.name,
Value: c.ctr.value,
CounterType: c.ctr.counterType,
}
}
return result, nil
}

308
prometheus/prometheus.go
View file

@ -1,308 +0,0 @@
// Copyright 2022 Molecula Corp. (DBA FeatureBase).
// SPDX-License-Identifier: Apache-2.0
package prometheus
import (
"sort"
"strings"
"sync"
"time"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/prometheus/client_golang/prometheus"
)
const (
// namespace is prepended to each metric event name with "_"
defaultNamespace = "general"
)
// Ensure client implements interface.
var _ stats.StatsClient = &prometheusClient{}
// Module-level mutex to avoid copying in WithTags()
var mu sync.Mutex
// prometheusClient represents a Prometheus implementation of pilosa.statsClient.
type prometheusClient struct {
tags []string
logger logger.Logger
counters map[string]prometheus.Counter
counterVecs map[string]*prometheus.CounterVec
gauges map[string]prometheus.Gauge
gaugeVecs map[string]*prometheus.GaugeVec
observers map[string]prometheus.Observer
summaryVecs map[string]*prometheus.SummaryVec
namespace string
}
// ClientOption is a functional option type for prometheusClient
type ClientOption func(c *prometheusClient)
// OptClientPrefix is a functional option on prometheusClient used to set the namespace
func OptClientNamespace(namespace string) ClientOption {
return func(c *prometheusClient) {
c.namespace = namespace
}
}
// NewPrometheusClient returns a new instance of StatsClient.
func NewPrometheusClient(opts ...ClientOption) (*prometheusClient, error) {
client := &prometheusClient{
logger: logger.NopLogger,
counters: make(map[string]prometheus.Counter),
counterVecs: make(map[string]*prometheus.CounterVec),
gauges: make(map[string]prometheus.Gauge),
gaugeVecs: make(map[string]*prometheus.GaugeVec),
observers: make(map[string]prometheus.Observer),
summaryVecs: make(map[string]*prometheus.SummaryVec),
namespace: defaultNamespace,
}
for _, opt := range opts {
opt(client)
}
return client, nil
}
// Open no-op to satisfy interface
func (c *prometheusClient) Open() {}
// Close no-op to satisfy interface
func (c *prometheusClient) Close() error {
return nil
}
// Tags returns a sorted list of tags on the client.
func (c *prometheusClient) Tags() []string {
return c.tags
}
// labels returns an instance of prometheus.Labels with the value of the set tags.
func (c *prometheusClient) labels() prometheus.Labels {
return tagsToLabels(c.tags, c.logger)
}
// WithTags returns a new client with additional tags appended.
func (c *prometheusClient) WithTags(tags ...string) stats.StatsClient {
return &prometheusClient{
tags: unionStringSlice(c.tags, tags),
logger: c.logger,
counters: c.counters,
counterVecs: c.counterVecs,
gauges: c.gauges,
gaugeVecs: c.gaugeVecs,
observers: c.observers,
summaryVecs: c.summaryVecs,
namespace: c.namespace,
}
}
// Count tracks the number of times something occurs per second.
func (c *prometheusClient) Count(name string, value int64, rate float64) {
mu.Lock()
defer mu.Unlock()
var counter prometheus.Counter
var ok bool
name = strings.Replace(name, ".", "_", -1)
labels := c.labels()
opts := prometheus.CounterOpts{
Namespace: c.namespace,
Name: name,
}
if len(labels) == 0 {
counter, ok = c.counters[name]
if !ok {
counter = prometheus.NewCounter(opts)
c.counters[name] = counter
prometheus.MustRegister(counter)
}
} else {
var counterVec *prometheus.CounterVec
counterVec, ok = c.counterVecs[name]
if !ok {
counterVec = prometheus.NewCounterVec(
opts,
labelKeys(labels),
)
c.counterVecs[name] = counterVec
prometheus.MustRegister(counterVec)
}
var err error
counter, err = counterVec.GetMetricWith(labels)
if err != nil {
c.logger.Errorf("counterVec.GetMetricWith error: %s", err)
}
}
if value == 1 {
counter.Inc()
} else {
counter.Add(float64(value))
}
}
// CountWithCustomTags tracks the number of times something occurs per second with custom tags.
func (c *prometheusClient) CountWithCustomTags(name string, value int64, rate float64, t []string) {
c.WithTags(append(c.tags, t...)...).Count(name, value, rate)
}
// Gauge sets the value of a metric.
func (c *prometheusClient) Gauge(name string, value float64, rate float64) {
mu.Lock()
defer mu.Unlock()
var gauge prometheus.Gauge
var ok bool
name = strings.Replace(name, ".", "_", -1)
labels := c.labels()
opts := prometheus.GaugeOpts{
Namespace: c.namespace,
Name: name,
}
if len(labels) == 0 {
gauge, ok = c.gauges[name]
if !ok {
gauge = prometheus.NewGauge(opts)
c.gauges[name] = gauge
prometheus.MustRegister(gauge)
}
} else {
var gaugeVec *prometheus.GaugeVec
gaugeVec, ok = c.gaugeVecs[name]
if !ok {
gaugeVec = prometheus.NewGaugeVec(
opts,
labelKeys(labels),
)
c.gaugeVecs[name] = gaugeVec
prometheus.MustRegister(gaugeVec)
}
var err error
gauge, err = gaugeVec.GetMetricWith(labels)
if err != nil {
c.logger.Errorf("gaugeVec.GetMetricWith error: %s", err)
return
}
}
gauge.Set(float64(value))
}
// Histogram tracks statistical distribution of a metric.
func (c *prometheusClient) Histogram(name string, value float64, rate float64) {
mu.Lock()
defer mu.Unlock()
var observer prometheus.Observer
var ok bool
name = strings.Replace(name, ".", "_", -1)
labels := c.labels()
opts := prometheus.SummaryOpts{
Namespace: c.namespace,
Name: name,
Objectives: map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
}
if len(labels) == 0 {
observer, ok = c.observers[name]
if !ok {
summary := prometheus.NewSummary(opts)
observer = summary
c.observers[name] = observer
prometheus.MustRegister(summary)
}
} else {
var summaryVec *prometheus.SummaryVec
summaryVec, ok = c.summaryVecs[name]
if !ok {
summaryVec = prometheus.NewSummaryVec(
opts,
labelKeys(labels),
)
c.summaryVecs[name] = summaryVec
prometheus.MustRegister(summaryVec)
}
var err error
observer, err = summaryVec.GetMetricWith(labels)
if err != nil {
c.logger.Errorf("summaryVec.GetMetricWith error: %s", err)
return
}
}
observer.Observe(value)
}
// Set tracks number of unique elements.
func (c *prometheusClient) Set(name string, value string, rate float64) {
c.logger.Infof("prometheusClient.Set unimplemented: %s=%s", name, value)
}
// Timing tracks timing information for a metric.
func (c *prometheusClient) Timing(name string, value time.Duration, rate float64) {
c.Histogram(name, value.Seconds(), rate)
}
// SetLogger sets the logger for client.
func (c *prometheusClient) SetLogger(logger logger.Logger) {
c.logger = logger
}
// unionStringSlice returns a sorted set of tags which combine a & b.
func unionStringSlice(a, b []string) []string {
// Sort both sets first.
sort.Strings(a)
sort.Strings(b)
// Find size of largest slice.
n := len(a)
if len(b) > n {
n = len(b)
}
// Exit if both sets are empty.
if n == 0 {
return nil
}
// Iterate over both in order and merge.
other := make([]string, 0, n)
for len(a) > 0 || len(b) > 0 {
if len(a) == 0 {
other, b = append(other, b[0]), b[1:]
} else if len(b) == 0 {
other, a = append(other, a[0]), a[1:]
} else if a[0] < b[0] {
other, a = append(other, a[0]), a[1:]
} else if b[0] < a[0] {
other, b = append(other, b[0]), b[1:]
} else {
other, a, b = append(other, a[0]), a[1:], b[1:]
}
}
return other
}
func tagsToLabels(tags []string, logger logger.Logger) (labels prometheus.Labels) {
labels = make(prometheus.Labels)
for _, tag := range tags {
tagParts := strings.SplitAfterN(tag, ":", 2)
if len(tagParts) != 2 {
// only process tags in "key:value" form
logger.Errorf("invalid Prometheus label: %v\n", tag)
continue
}
labels[tagParts[0][0:len(tagParts[0])-1]] = tagParts[1]
}
return labels
}
func labelKeys(labels prometheus.Labels) (keys []string) {
keys = make([]string, len(labels))
i := 0
for k := range labels {
keys[i] = k
i++
}
return keys
}

51
prometheus/prometheus_test.go
View file

@ -3,62 +3,33 @@
package prometheus_test
import (
"reflect"
"testing"
"time"
pilosaPrometheus "github.com/featurebasedb/featurebase/v3/prometheus"
"github.com/featurebasedb/featurebase/v3/test"
"github.com/prometheus/client_golang/prometheus"
io_prometheus_client "github.com/prometheus/client_model/go"
)
func TestPrometheusClient_WithTags(t *testing.T) {
// Create a new client.
c, err := pilosaPrometheus.NewPrometheusClient()
if err != nil {
t.Fatal(err)
}
defer c.Close()
// Create a new client with additional tags.
c1 := c.WithTags("foo", "bar")
if tags := c1.Tags(); !reflect.DeepEqual(tags, []string{"bar", "foo"}) {
t.Fatalf("unexpected tags: %+v", tags)
}
// Create a new client from the clone with more tags.
c2 := c1.WithTags("bar", "baz")
if tags := c2.Tags(); !reflect.DeepEqual(tags, []string{"bar", "baz", "foo"}) {
t.Fatalf("unexpected tags: %+v", tags)
}
}
func TestPrometheusClient_Methods(t *testing.T) {
// Create a new client.
c, err := pilosaPrometheus.NewPrometheusClient(
pilosaPrometheus.OptClientNamespace("testns"),
)
if err != nil {
t.Fatal(err)
}
c := test.MustRunCluster(t, 1)
defer c.Close()
dur, _ := time.ParseDuration("123us")
c.CountWithCustomTags("ct", 1, 1.0, []string{"foo:bar"})
c.Count("cc", 1, 1.0)
c.Gauge("gg", 10, 1.0)
c.Histogram("hh", 1, 1.0)
c.Timing("tt", dur, 1.0)
metricFams, err := prometheus.DefaultGatherer.Gather()
if err != nil {
t.Fatal(err)
}
for _, metricName := range []string{"testns_ct", "testns_cc", "testns_gg", "testns_hh", "testns_tt"} {
for _, metricName := range []string{
"pilosa_sql_statistics_sql_bulk_insert_batches_sec",
"pilosa_sql_statistics_sql_bulk_inserts_sec",
"pilosa_sql_statistics_sql_deletes_sec",
"pilosa_sql_statistics_sql_inserts_sec",
"pilosa_sql_statistics_sql_requests_sec",
} {
if metricExists(metricName, metricFams) {
continue
}
t.Fatalf("Metric was not recorded: %s", metricName)
t.Fatalf("metric does not exist: %s", metricName)
}
}

4
row.go
View file

@ -362,7 +362,7 @@ func (r *Row) Difference(others ...*Row) *Row {
// be incorrect.
//
// Why unsupported? For a full description, see:
// https://github.com/molecula/pilosa/issues/403.
// https://github.com/featurebasedb/pilosa/issues/403.
// In short, the current implementation will shift a bit
// at the edge of a shard out of the shard and into a
// container which is assumed to be an invalid container
@ -612,7 +612,7 @@ func (s *RowSegment) Xor(other *RowSegment) *RowSegment {
// Shift returns s shifted by 1 bit.
func (s *RowSegment) Shift() (*RowSegment, error) {
// TODO: deal with overflow
// See issue: https://github.com/molecula/pilosa/issues/403
// See issue: https://github.com/featurebasedb/pilosa/issues/403
data, err := s.data.Shift(1)
if err != nil {
return nil, errors.Wrap(err, "shifting roaring data")

3
schema.go
View file

@ -28,6 +28,9 @@ func NewOnPremSchema(api *API) *onPremSchema {
func (s *onPremSchema) TableByName(ctx context.Context, tname dax.TableName) (*dax.Table, error) {
idx, err := s.api.IndexInfo(context.Background(), string(tname))
if err != nil {
if err == ErrIndexNotFound {
return nil, dax.NewErrTableNameDoesNotExist(tname)
}
return nil, errors.Wrapf(err, "getting index info for table name: %s", tname)
}

48
server.go
View file

@ -5,6 +5,7 @@ package pilosa
import (
"context"
"fmt"
"io"
"log"
"os"
"os/exec"
@ -26,7 +27,6 @@ import (
"github.com/featurebasedb/featurebase/v3/sql3"
"github.com/featurebasedb/featurebase/v3/sql3/parser"
planner_types "github.com/featurebasedb/featurebase/v3/sql3/planner/types"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/storage"
"github.com/pkg/errors"
"golang.org/x/sync/errgroup"
@ -246,15 +246,6 @@ func OptServerPrimaryTranslateStore(store TranslateStore) ServerOption {
}
}
// OptServerStatsClient is a functional option on Server
// used to specify the stats client.
func OptServerStatsClient(sc stats.StatsClient) ServerOption {
return func(s *Server) error {
s.holderConfig.StatsClient = sc
return nil
}
}
// OptServerDiagnosticsInterval is a functional option on Server
// used to specify the duration between diagnostic checks.
func OptServerDiagnosticsInterval(dur time.Duration) ServerOption {
@ -542,7 +533,6 @@ func NewServer(opts ...ServerOption) (*Server, error) {
return nil, err
}
s.holder = NewHolder(path, s.holderConfig)
s.holder.Stats.SetLogger(s.logger)
cwd, err := os.Getwd()
if err != nil {
return nil, err
@ -558,9 +548,6 @@ func NewServer(opts ...ServerOption) (*Server, error) {
s.cluster.sharder = s.sharder
s.cluster.serverlessStorage = s.serverlessStorage
// Append the NodeID tag to stats.
s.holder.Stats = s.holder.Stats.WithTags(fmt.Sprintf("node_id:%s", s.nodeID))
s.executor.Holder = s.holder
s.holder.executor = s.executor
s.executor.Cluster = s.cluster
@ -617,6 +604,18 @@ func (s *Server) Open() error {
log.Println(errors.Wrap(err, "logging startup"))
}
// Do version check in. This is in a goroutine so that we don't block server startup if the server endpoint is down/having issues.
go func() {
s.logger.Printf("Beginning featurebase version check-in")
vc := VersionChecker{URL: "https://analytics.featurebase.com/v2/featurebase/version"}
resp, err := vc.CheckIn()
if err != nil {
s.logger.Errorf("doing version checkin. Error was %s", err)
return
}
s.logger.Printf("Version check-in complete. Latest version is %s", resp.Information.Version)
}()
// Start DisCo.
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
@ -658,7 +657,6 @@ func (s *Server) Open() error {
s.syncer.Node = node
s.syncer.Cluster = s.cluster
s.syncer.Closing = s.closing
s.syncer.Stats = s.holder.Stats.WithTags("component:HolderSyncer")
// Start background process listening for translation
// sync resets.
@ -1250,26 +1248,26 @@ func (s *Server) monitorRuntime() {
return
case <-s.gcNotifier.AfterGC():
// GC just ran.
s.holder.Stats.Count(MetricGarbageCollection, 1, 1.0)
CounterGarbageCollection.Inc()
case <-ticker.C:
}
// Record the number of go routines.
s.holder.Stats.Gauge(MetricGoroutines, float64(runtime.NumGoroutine()), 1.0)
GaugeGoroutines.Set(float64(runtime.NumGoroutine()))
openFiles, err := countOpenFiles()
// Open File handles.
if err == nil {
s.holder.Stats.Gauge(MetricOpenFiles, float64(openFiles), 1.0)
GaugeOpenFiles.Set(float64(openFiles))
}
// Runtime memory metrics.
runtime.ReadMemStats(&m)
s.holder.Stats.Gauge(MetricHeapAlloc, float64(m.HeapAlloc), 1.0)
s.holder.Stats.Gauge(MetricHeapInuse, float64(m.HeapInuse), 1.0)
s.holder.Stats.Gauge(MetricStackInuse, float64(m.StackInuse), 1.0)
s.holder.Stats.Gauge(MetricMallocs, float64(m.Mallocs), 1.0)
s.holder.Stats.Gauge(MetricFrees, float64(m.Frees), 1.0)
GaugeHeapAlloc.Set(float64(m.HeapAlloc))
GaugeHeapInUse.Set(float64(m.HeapInuse))
GaugeStackInUse.Set(float64(m.StackInuse))
GaugeMallocs.Set(float64(m.Mallocs))
GaugeFrees.Set(float64(m.Frees))
}
}
@ -1413,6 +1411,10 @@ func (s *Server) CompileExecutionPlan(ctx context.Context, q string) (planner_ty
return s.executionPlannerFn(s.executor, s.executor.client.api, q).CompilePlan(ctx, st)
}
func (s *Server) RehydratePlanOperator(ctx context.Context, reader io.Reader) (planner_types.PlanOperator, error) {
return s.executionPlannerFn(s.executor, s.executor.client.api, "").RehydratePlanOp(ctx, reader)
}
// countOpenFiles on operating systems that support lsof.
func countOpenFiles() (int, error) {
switch runtime.GOOS {

43
server/grpc.go
View file

@ -22,7 +22,6 @@ import (
pb "github.com/featurebasedb/featurebase/v3/proto"
vdsm_pb "github.com/featurebasedb/featurebase/v3/proto/vdsm"
"github.com/featurebasedb/featurebase/v3/sql"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/improbable-eng/grpc-web/go/grpcweb"
"github.com/pkg/errors"
"google.golang.org/grpc"
@ -41,12 +40,11 @@ type GRPCHandler struct {
perms *authz.GroupPermissions
logger logger.Logger
queryLogger logger.Logger
stats stats.StatsClient
inspectDeprecated sync.Once
}
func NewGRPCHandler(api *pilosa.API) *GRPCHandler {
return &GRPCHandler{api: api, logger: logger.NopLogger, stats: stats.NopStatsClient}
return &GRPCHandler{api: api, logger: logger.NopLogger}
}
func (h *GRPCHandler) WithLogger(logger logger.Logger) *GRPCHandler {
@ -54,11 +52,6 @@ func (h *GRPCHandler) WithLogger(logger logger.Logger) *GRPCHandler {
return h
}
func (h *GRPCHandler) WithStats(stats stats.StatsClient) *GRPCHandler {
h.stats = stats
return h
}
func (h *GRPCHandler) WithPerms(perms *authz.GroupPermissions) *GRPCHandler {
h.perms = perms
return h
@ -140,7 +133,7 @@ func errToStatusError(err error) error {
}
func (h *GRPCHandler) execSQL(ctx context.Context, queryStr string) (pb.ToRowser, error) {
h.stats.Count(pilosa.MetricSqlQueries, 1, 1)
pilosa.CounterSQLQueries.Inc()
return execSQL(ctx, h.api, h.logger, queryStr)
}
@ -340,9 +333,9 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ
return errToStatusError(err)
}
durFormat := time.Since(t)
h.stats.Timing(pilosa.MetricGRPCStreamQueryDurationSeconds, durQuery, 0.1)
h.stats.Timing(pilosa.MetricGRPCStreamFormatDurationSeconds, durFormat, 0.1)
h.stats.Count(pilosa.MetricPqlQueries, 1, 1)
pilosa.SummaryGRPCStreamQueryDurationSeconds.Observe(durQuery.Seconds())
pilosa.SummaryGRPCStreamFormatDurationSeconds.Observe(durFormat.Seconds())
pilosa.CounterPQLQueries.Inc()
return errToStatusError(nil)
}
@ -406,9 +399,9 @@ func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest
return nil, errors.Wrap(err, "sending header")
}
h.stats.Timing(pilosa.MetricGRPCUnaryQueryDurationSeconds, durQuery, 0.1)
h.stats.Timing(pilosa.MetricGRPCUnaryFormatDurationSeconds, durFormat, 0.1)
h.stats.Count(pilosa.MetricPqlQueries, 1, 1)
pilosa.SummaryGRPCUnaryQueryDurationSeconds.Observe(durQuery.Seconds())
pilosa.SummaryGRPCUnaryFormatDurationSeconds.Observe(durFormat.Seconds())
pilosa.CounterPQLQueries.Inc()
return table, errToStatusError(nil)
}
@ -510,11 +503,10 @@ type VDSMGRPCHandler struct {
grpcHandler *GRPCHandler
api *pilosa.API
logger logger.Logger
stats stats.StatsClient
}
func NewVDSMGRPCHandler(grpcHandler *GRPCHandler, api *pilosa.API) *VDSMGRPCHandler {
return &VDSMGRPCHandler{grpcHandler: grpcHandler, api: api, logger: logger.NopLogger, stats: stats.NopStatsClient}
return &VDSMGRPCHandler{grpcHandler: grpcHandler, api: api, logger: logger.NopLogger}
}
func (h *VDSMGRPCHandler) WithLogger(logger logger.Logger) *VDSMGRPCHandler {
@ -522,11 +514,6 @@ func (h *VDSMGRPCHandler) WithLogger(logger logger.Logger) *VDSMGRPCHandler {
return h
}
func (h *VDSMGRPCHandler) WithStats(stats stats.StatsClient) *VDSMGRPCHandler {
h.stats = stats
return h
}
// GetVDSs returns a single VDS given a name
func (h *VDSMGRPCHandler) GetVDS(ctx context.Context, req *vdsm_pb.GetVDSRequest) (*vdsm_pb.GetVDSResponse, error) {
typedIdOrName := req.GetIdOrName()
@ -1481,7 +1468,6 @@ type grpcServer struct {
logger logger.Logger
queryLogger logger.Logger
stats stats.StatsClient
}
type grpcServerOption func(s *grpcServer) error
@ -1514,13 +1500,6 @@ func OptGRPCServerLogger(logger logger.Logger) grpcServerOption {
}
}
func OptGRPCServerStats(stats stats.StatsClient) grpcServerOption {
return func(s *grpcServer) error {
s.stats = stats
return nil
}
}
func OptGRPCServerAuth(authn *authn.Auth) grpcServerOption {
return func(s *grpcServer) error {
s.auth = authn
@ -1640,7 +1619,7 @@ func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) {
// create grpc server
server.grpcServer = grpc.NewServer(gopts...)
grpcHandler := NewGRPCHandler(server.api).WithLogger(server.logger).WithStats(server.stats).WithQueryLogger(server.queryLogger)
grpcHandler := NewGRPCHandler(server.api).WithLogger(server.logger).WithQueryLogger(server.queryLogger)
// add server permissions if we've got 'em
if server.perms != nil {
@ -1648,7 +1627,7 @@ func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) {
}
pb.RegisterPilosaServer(server.grpcServer, grpcHandler)
vdsm_pb.RegisterMoleculaServer(server.grpcServer, NewVDSMGRPCHandler(grpcHandler, server.api).WithLogger(server.logger).WithStats(server.stats))
vdsm_pb.RegisterMoleculaServer(server.grpcServer, NewVDSMGRPCHandler(grpcHandler, server.api).WithLogger(server.logger))
// register the server so its services are available to grpc_cli and others
reflection.Register(server.grpcServer)

9
server/handler_test.go
View file

@ -1080,15 +1080,6 @@ func TestHandler_Endpoints(t *testing.T) {
}
})
t.Run("Expvars", func(t *testing.T) {
w := httptest.NewRecorder()
r := test.MustNewHTTPRequest("GET", "/debug/vars", nil)
h.ServeHTTP(w, r)
if w.Code != http.StatusOK {
t.Fatalf("unexpected status code: %d", w.Code)
}
})
t.Run("Recalculate Caches", func(t *testing.T) {
w := httptest.NewRecorder()
h.ServeHTTP(w, test.MustNewHTTPRequest("POST", "/recalculate-caches", nil))

28
server/server.go
View file

@ -39,12 +39,9 @@ import (
"github.com/featurebasedb/featurebase/v3/gopsutil"
"github.com/featurebasedb/featurebase/v3/logger"
pnet "github.com/featurebasedb/featurebase/v3/net"
"github.com/featurebasedb/featurebase/v3/prometheus"
"github.com/featurebasedb/featurebase/v3/sql3"
"github.com/featurebasedb/featurebase/v3/sql3/planner"
"github.com/featurebasedb/featurebase/v3/statik"
"github.com/featurebasedb/featurebase/v3/stats"
"github.com/featurebasedb/featurebase/v3/statsd"
"github.com/featurebasedb/featurebase/v3/systemlayer"
"github.com/featurebasedb/featurebase/v3/syswrap"
"github.com/featurebasedb/featurebase/v3/testhook"
@ -474,11 +471,6 @@ func (m *Command) setupServer() error {
diagnosticsInterval = defaultDiagnosticsInterval
}
statsClient, err := newStatsClient(m.Config.Metric.Service, m.Config.Metric.Host, m.Config.Namespace())
if err != nil {
return errors.Wrap(err, "new stats client")
}
if m.Config.Listener == nil {
m.ln, err = getListener(*uri, m.tlsConfig)
if err != nil {
@ -581,7 +573,6 @@ func (m *Command) setupServer() error {
pilosa.OptServerQueryLogger(m.queryLogger),
pilosa.OptServerSystemInfo(gopsutil.NewSystemInfo()),
pilosa.OptServerGCNotifier(gcnotify.NewActiveGCNotifier()),
pilosa.OptServerStatsClient(statsClient),
pilosa.OptServerURI(advertiseURI),
pilosa.OptServerGRPCURI(advertiseGRPCURI),
pilosa.OptServerClusterName(m.Config.Cluster.Name),
@ -691,7 +682,6 @@ func (m *Command) setupServer() error {
OptGRPCServerListener(m.grpcLn),
OptGRPCServerTLSConfig(m.tlsConfig),
OptGRPCServerLogger(m.logger),
OptGRPCServerStats(statsClient),
OptGRPCServerAuth(m.auth),
OptGRPCServerPerm(&p),
OptGRPCServerQueryLogger(m.queryLogger),
@ -828,24 +818,6 @@ func (m *Command) Close() error {
}
}
// newStatsClient creates a stats client from the config
func newStatsClient(name string, host string, namespace string) (stats.StatsClient, error) {
switch name {
case "expvar":
return stats.NewExpvarStatsClient(), nil
case "statsd":
return statsd.NewStatsClient(host, namespace)
case "prometheus":
return prometheus.NewPrometheusClient(
prometheus.OptClientNamespace(namespace),
)
case "nop", "none":
return stats.NopStatsClient, nil
default:
return nil, errors.Errorf("'%v' not a valid stats client, choose from [expvar, statsd, prometheus, none].", name)
}
}
// getListener gets a net.Listener based on the config.
func getListener(uri pnet.URI, tlsconf *tls.Config) (ln net.Listener, err error) {
// If bind URI has the https scheme, enable TLS

Some files were not shown because too many files have changed in this diff Show more