diff --git a/api.go b/api.go index ae61eb99a..5e7bfa418 100644 --- a/api.go +++ b/api.go @@ -3354,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 { @@ -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 +} diff --git a/api_directive_test.go b/api_directive_test.go index fffa1ce77..a90a3b1ba 100644 --- a/api_directive_test.go +++ b/api_directive_test.go @@ -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) { diff --git a/dax/address.go b/dax/address.go index 99e65ff7d..ab5e45ec6 100644 --- a/dax/address.go +++ b/dax/address.go @@ -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 { diff --git a/dax/boltdb/boltdb.go b/dax/boltdb/boltdb.go index 1ddaf795c..aa50bd27d 100644 --- a/dax/boltdb/boltdb.go +++ b/dax/boltdb/boltdb.go @@ -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 } diff --git a/dax/boltdb/directiveversion.go b/dax/boltdb/directiveversion.go index b00452c6a..c4c04fe44 100644 --- a/dax/boltdb/directiveversion.go +++ b/dax/boltdb/directiveversion.go @@ -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 } diff --git a/dax/boltdb/node.go b/dax/boltdb/node.go index 583ebb01d..b06b3912d 100644 --- a/dax/boltdb/node.go +++ b/dax/boltdb/node.go @@ -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. diff --git a/dax/boltdb/node_test.go b/dax/boltdb/node_test.go index f9791ae5e..591bb6968 100644 --- a/dax/boltdb/node_test.go +++ b/dax/boltdb/node_test.go @@ -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()) }) } diff --git a/dax/directive.go b/dax/directive.go index a0eccad26..c7ffde19a 100644 --- a/dax/directive.go +++ b/dax/directive.go @@ -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] } diff --git a/dax/errors.go b/dax/errors.go index d5bee2627..728b297ea 100644 --- a/dax/errors.go +++ b/dax/errors.go @@ -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", + ) +} diff --git a/dax/mds/client/client.go b/dax/mds/client/client.go index d105d1f1a..3f2c1c163 100644 --- a/dax/mds/client/client.go +++ b/dax/mds/client/client.go @@ -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)) diff --git a/dax/mds/controller/balancer.go b/dax/mds/controller/balancer.go index a6402c6fb..567c4b8c7 100644 --- a/dax/mds/controller/balancer.go +++ b/dax/mds/controller/balancer.go @@ -8,31 +8,38 @@ import ( ) 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 +52,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 -} diff --git a/dax/mds/controller/balancer/balancer.go b/dax/mds/controller/balancer/balancer.go new file mode 100644 index 000000000..f29813819 --- /dev/null +++ b/dax/mds/controller/balancer/balancer.go @@ -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) +} diff --git a/dax/mds/controller/balancer/balancer_test.go b/dax/mds/controller/balancer/balancer_test.go new file mode 100644 index 000000000..e538ef388 --- /dev/null +++ b/dax/mds/controller/balancer/balancer_test.go @@ -0,0 +1,1104 @@ +package balancer_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" + "github.com/featurebasedb/featurebase/v3/dax/mds/controller/balancer/boltdb" + schemardb "github.com/featurebasedb/featurebase/v3/dax/mds/schemar/boltdb" + daxtest "github.com/featurebasedb/featurebase/v3/dax/test" + 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.BalancerBuckets...)) + assert.NoError(t, db.InitializeBuckets(schemardb.SchemarBuckets...)) + + return db, func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + } +} + +type runner interface { + run(tx dax.Transaction, bal controller.Balancer) ([]dax.WorkerDiff, error) +} + +type addWorker struct { + node *dax.Node +} + +func (r *addWorker) run(tx dax.Transaction, bal controller.Balancer) ([]dax.WorkerDiff, error) { + return bal.AddWorker(tx, r.node) +} + +type removeWorker struct { + addr dax.Address +} + +func (r *removeWorker) run(tx dax.Transaction, bal controller.Balancer) ([]dax.WorkerDiff, error) { + return bal.RemoveWorker(tx, r.addr) +} + +type addJob struct { + roleType dax.RoleType + qtid dax.QualifiedTableID + job dax.Job +} + +func (r *addJob) run(tx dax.Transaction, bal controller.Balancer) ([]dax.WorkerDiff, error) { + return bal.AddJobs(tx, r.roleType, r.qtid, r.job) +} + +type removeJob struct { + roleType dax.RoleType + qtid dax.QualifiedTableID + job dax.Job +} + +func (r *removeJob) run(tx dax.Transaction, bal controller.Balancer) ([]dax.WorkerDiff, error) { + return bal.RemoveJobs(tx, r.roleType, r.qtid, r.job) +} + +type balanceDatabase struct { + qdbid dax.QualifiedDatabaseID +} + +func (r *balanceDatabase) run(tx dax.Transaction, bal controller.Balancer) ([]dax.WorkerDiff, error) { + return bal.BalanceDatabase(tx, r.qdbid) +} + +func TestBalancer(t *testing.T) { + ctx := context.Background() + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + dbName := dax.DatabaseName("db1name") + tableID := dax.TableID("tbl1") + qtid := dax.NewQualifiedTableID(dax.NewQualifiedDatabaseID(orgID, dbID), tableID) + + t.Run("SingleWorker", func(t *testing.T) { + qdb := &dax.QualifiedDatabase{ + OrganizationID: orgID, + Database: dax.Database{ + ID: dbID, + Name: dbName, + Options: dax.DatabaseOptions{ + WorkersMin: 1, + WorkersMax: 1, + }, + }, + } + + schemar, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer(db, schemar, logger.NewStandardLogger(os.Stderr)) + + tx, err := db.BeginTx(ctx, true) + assert.NoError(t, err) + defer tx.Rollback() + + assert.NoError(t, schemar.CreateDatabase(tx, qdb)) + + role := dax.RoleTypeCompute + + tests := []struct { + runner runner + expDiff []dax.WorkerDiff + expState []dax.WorkerInfo + }{ + { + // Add job. + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s2", + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }, + { + // Add worker. + runner: &addWorker{ + node: &dax.Node{ + Address: "w1", + RoleTypes: dax.RoleTypes{role}, + }, + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s2"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s2"}, + }, + }, + }, + { + // Add another job out of order. + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s1", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s1", "s2"}, + }, + }, + }, + { + // Add another job. + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s3", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s3"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s1", "s2", "s3"}, + }, + }, + }, + { + // Add a duplicate job. + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s2", + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s1", "s2", "s3"}, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + diff, err := test.runner.run(tx, bal) + assert.NoError(t, err) + assert.Equal(t, test.expDiff, diff) + + state, err := bal.CurrentState(tx, role, qdb.QualifiedID()) + assert.NoError(t, err) + assert.Equal(t, test.expState, state) + }) + } + + assert.NoError(t, tx.Commit()) + }) + + t.Run("MultipleWorkers", func(t *testing.T) { + dbOptions := dax.DatabaseOptions{ + WorkersMin: 2, + WorkersMax: 2, + } + + qdb := &dax.QualifiedDatabase{ + OrganizationID: orgID, + Database: dax.Database{ + ID: dbID, + Name: dbName, + Options: dbOptions, + }, + } + + schemar, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer(db, schemar, logger.NewStandardLogger(os.Stderr)) + + tx, err := db.BeginTx(ctx, true) + assert.NoError(t, err) + defer tx.Rollback() + + assert.NoError(t, schemar.CreateDatabase(tx, qdb)) + + role := dax.RoleTypeCompute + + type testPart struct { + name string + runner runner + expDiff []dax.WorkerDiff + expState []dax.WorkerInfo + } + + runTestPart := func(tp testPart) { + t.Run(fmt.Sprintf("test-%s", tp.name), func(t *testing.T) { + diff, err := tp.runner.run(tx, bal) + assert.NoError(t, err) + assert.Equal(t, tp.expDiff, diff) + + state, err := bal.CurrentState(tx, role, qdb.QualifiedID()) + assert.NoError(t, err) + assert.Equal(t, tp.expState, state) + }) + } + + runTestPart(testPart{ + name: "balance when empty", + runner: &balanceDatabase{ + qdbid: qdb.QualifiedID(), + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }) + + runTestPart(testPart{ + name: "add worker", + runner: &addWorker{ + node: &dax.Node{ + Address: "w2", + RoleTypes: dax.RoleTypes{role}, + }, + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }) + + runTestPart(testPart{ + name: "add worker again", + runner: &addWorker{ + node: &dax.Node{ + Address: "w2", + RoleTypes: dax.RoleTypes{role}, + }, + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }) + + runTestPart(testPart{ + name: "add a second worker", + runner: &addWorker{ + node: &dax.Node{ + Address: "w1", + RoleTypes: dax.RoleTypes{role}, + }, + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }) + + runTestPart(testPart{ + name: "add job 2", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s2", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s2"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s2"}, + }, + { + Address: "w2", + Jobs: []dax.Job{}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 3", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s3", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w2", + AddedJobs: []dax.Job{"s3"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s2"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 1", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s1", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w1", + Jobs: []dax.Job{"s1", "s2"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s3"}, + }, + }, + }) + + // Update database options on schemar so min worker for db is 3. Here we + // are updating the database option first, and then adding a worker to + // satisfy those options. + t.Run(fmt.Sprintf("test-%s", "set database options db min 3"), func(t *testing.T) { + dbOptions.WorkersMin = 3 + dbOptions.WorkersMax = 3 + assert.NoError(t, schemar.SetDatabaseOptions(tx, qdb.QualifiedID(), dbOptions)) + }) + + runTestPart(testPart{ + name: "add a third worker", + runner: &addWorker{ + node: &dax.Node{ + Address: "w0", + RoleTypes: dax.RoleTypes{role}, + }, + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w0", + AddedJobs: []dax.Job{"s2"}, + RemovedJobs: []dax.Job{}, + }, + { + Address: "w1", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"s2"}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 4", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s4", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w0", + AddedJobs: []dax.Job{"s4"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 5", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s5", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s5"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1", "s5"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 0", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s0", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w2", + AddedJobs: []dax.Job{"s0"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1", "s5"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s0", "s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 6", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s6", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w0", + AddedJobs: []dax.Job{"s6"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s6"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1", "s5"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s0", "s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "add job 7", + runner: &addJob{ + roleType: role, + qtid: qtid, + job: "s7", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w1", + AddedJobs: []dax.Job{"s7"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s6"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1", "s5", "s7"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s0", "s3"}, + }, + }, + }) + + //////////////////// Remove ///////////////////////// + + runTestPart(testPart{ + name: "remove nonexistent worker", + runner: &removeWorker{ + addr: "nonexistent", + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s6"}, + }, + { + Address: "w1", + Jobs: []dax.Job{"s1", "s5", "s7"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s0", "s3"}, + }, + }, + }) + + runTestPart(testPart{ + name: "remove worker", + runner: &removeWorker{ + addr: "w1", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w0", + AddedJobs: []dax.Job{"s5"}, + RemovedJobs: []dax.Job{}, + }, + { + Address: "w1", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"s1", "s5", "s7"}, + }, + { + Address: "w2", + AddedJobs: []dax.Job{"s1", "s7"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s5", "s6"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s0", "s1", "s3", "s7"}, + }, + }, + }) + + runTestPart(testPart{ + name: "remove active job", + runner: &removeJob{ + roleType: role, + qtid: qtid, + job: "s0", + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w2", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"s0"}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s5", "s6"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s1", "s3", "s7"}, + }, + }, + }) + + // Set WorkersMin back to 2 so we can test the change to 3 again. + t.Run(fmt.Sprintf("test-%s", "set database options db min back to 2"), func(t *testing.T) { + dbOptions.WorkersMin = 2 + dbOptions.WorkersMax = 2 + assert.NoError(t, schemar.SetDatabaseOptions(tx, qdb.QualifiedID(), dbOptions)) + }) + + runTestPart(testPart{ + name: "add a fourth worker", + runner: &addWorker{ + node: &dax.Node{ + Address: "w3", + RoleTypes: dax.RoleTypes{role}, + }, + }, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s5", "s6"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s1", "s3", "s7"}, + }, + }, + }) + + // Update database options on schemar so min worker for db is 4. Here we + // have added a worker which will satisfy this option, and then updated + // the option. + t.Run(fmt.Sprintf("test-%s", "set database options db min back to 3"), func(t *testing.T) { + dbOptions.WorkersMin = 3 + dbOptions.WorkersMax = 3 + assert.NoError(t, schemar.SetDatabaseOptions(tx, qdb.QualifiedID(), dbOptions)) + }) + + // This implies that there is a condition where the database does not + // get repaired to its workerMin: if workers per database has dropped + // below its min, and there are no available workers to replace the + // missing workers (and bring it back to min), then the database will + // operate below min. If, then, a worker is added to the pool, we do not + // currently have a process to automatically assign that worker to a + // database under min. That happens with an explicit call to + // assignMinWorkers or BalanceDatabase. + runTestPart(testPart{ + name: "balance to include new third worker", + runner: &balanceDatabase{ + qdbid: qdb.QualifiedID(), + }, + expDiff: []dax.WorkerDiff{ + { + Address: "w0", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"s6"}, + }, + { + Address: "w2", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"s7"}, + }, + { + Address: "w3", + AddedJobs: []dax.Job{"s6", "s7"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + Address: "w0", + Jobs: []dax.Job{"s2", "s4", "s5"}, + }, + { + Address: "w2", + Jobs: []dax.Job{"s1", "s3"}, + }, + { + Address: "w3", + Jobs: []dax.Job{"s6", "s7"}, + }, + }, + }) + + assert.NoError(t, tx.Commit()) + }) + + t.Run("WorkersForJobs", func(t *testing.T) { + qdb := &dax.QualifiedDatabase{ + OrganizationID: orgID, + Database: dax.Database{ + ID: dbID, + Name: dbName, + Options: dax.DatabaseOptions{ + WorkersMin: 2, + WorkersMax: 2, + }, + }, + } + qdbid := qdb.QualifiedID() + + schemar, scleanup := daxtest.NewSchemar(t) + defer scleanup() + + role := dax.RoleTypeCompute + + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer(db, schemar, logger.NewStandardLogger(os.Stderr)) + + tx, err := db.BeginTx(ctx, true) + assert.NoError(t, err) + defer tx.Rollback() + + assert.NoError(t, schemar.CreateDatabase(tx, qdb)) + + node1 := &dax.Node{ + Address: "n1", + RoleTypes: []dax.RoleType{role}, + } + + node2 := &dax.Node{ + Address: "n2", + RoleTypes: []dax.RoleType{role}, + } + + _, err = bal.AddWorker(tx, node1) + assert.NoError(t, err) + _, err = bal.AddWorker(tx, node2) + assert.NoError(t, err) + for i := 0; i < 12; i++ { + _, err = bal.AddJobs(tx, role, qtid, dax.Job(fmt.Sprintf("s%d", i))) + assert.NoError(t, err) + } + + exp := dax.WorkerInfo{ + Address: "n1", + Jobs: []dax.Job{"s0", "s10", "s2", "s4", "s6", "s8"}, + } + ws, err := bal.WorkerState(tx, role, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + Address: "n2", + Jobs: []dax.Job{"s1", "s11", "s3", "s5", "s7", "s9"}, + } + ws, err = bal.WorkerState(tx, role, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + tests := []struct { + jobs []dax.Job + exp []dax.WorkerInfo + }{ + { + jobs: []dax.Job{"s0"}, + exp: []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{"s0"}}, + }, + }, + { + jobs: []dax.Job{"s0", "s4"}, + exp: []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{"s0", "s4"}}, + }, + }, + { + jobs: []dax.Job{"s0", "s4", "s999"}, + exp: []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{"s0", "s4"}}, + }, + }, + { + jobs: []dax.Job{"s0", "s1"}, + exp: []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{"s0"}}, + {Address: "n2", Jobs: []dax.Job{"s1"}}, + }, + }, + { + jobs: []dax.Job{"s5", "s0", "s1", "s8"}, + exp: []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{"s0", "s8"}}, + {Address: "n2", Jobs: []dax.Job{"s1", "s5"}}, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + workers, err := bal.WorkersForJobs(tx, role, qdbid, test.jobs...) + assert.NoError(t, err) + assert.Equal(t, test.exp, workers) + }) + } + + assert.NoError(t, tx.Commit()) + }) + + t.Run("WorkersForTable", func(t *testing.T) { + qdb := &dax.QualifiedDatabase{ + OrganizationID: orgID, + Database: dax.Database{ + ID: dbID, + Name: dbName, + Options: dax.DatabaseOptions{ + WorkersMin: 2, + WorkersMax: 2, + }, + }, + } + + qdbid := dax.NewQualifiedDatabaseID(orgID, dbID) + + // Table 1. + tbl1 := &dax.Table{ + ID: "id1", + Name: "table1", + } + qtbl1 := dax.NewQualifiedTable(qdbid, tbl1) + qtid1 := qtbl1.QualifiedID() + + // Table 2. + tbl2 := &dax.Table{ + ID: "id2", + Name: "table2", + } + qtbl2 := dax.NewQualifiedTable(qdbid, tbl2) + qtid2 := qtbl2.QualifiedID() + + schemar, scleanup := daxtest.NewSchemar(t) + defer scleanup() + + role := dax.RoleTypeCompute + + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer(db, schemar, logger.NewStandardLogger(os.Stderr)) + + tx, err := db.BeginTx(ctx, true) + assert.NoError(t, err) + defer tx.Rollback() + + assert.NoError(t, schemar.CreateDatabase(tx, qdb)) + + node1 := &dax.Node{ + Address: "n1", + RoleTypes: []dax.RoleType{role}, + } + + node2 := &dax.Node{ + Address: "n2", + RoleTypes: []dax.RoleType{role}, + } + + _, err = bal.AddWorker(tx, node1) + assert.NoError(t, err) + _, err = bal.AddWorker(tx, node2) + assert.NoError(t, err) + for i := 0; i < 12; i++ { + job := fmt.Sprintf("%s|s%d", qtid1.Key(), i) + _, err = bal.AddJobs(tx, role, qtid1, dax.Job(job)) + assert.NoError(t, err) + } + + for i := 9; i < 16; i++ { + job := fmt.Sprintf("%s|s%d", qtid2.Key(), i) + _, err = bal.AddJobs(tx, role, qtid2, dax.Job(job)) + assert.NoError(t, err) + } + + // fn1 and fn2 are just helper functions used to make the tests easier + // to read. + fn1 := func(s string) dax.Job { + return dax.Job(fmt.Sprintf("tbl__acme__db1__id1|%s", s)) + } + fn2 := func(s string) dax.Job { + return dax.Job(fmt.Sprintf("tbl__acme__db1__id2|%s", s)) + } + + // Table 1 + workers, err := bal.WorkersForTable(tx, role, qtid1) + assert.NoError(t, err) + assert.ElementsMatch(t, []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{fn1("s0"), fn1("s10"), fn1("s2"), fn1("s4"), fn1("s6"), fn1("s8")}}, + {Address: "n2", Jobs: []dax.Job{fn1("s1"), fn1("s11"), fn1("s3"), fn1("s5"), fn1("s7"), fn1("s9")}}, + }, workers) + + // Table 2 + workers, err = bal.WorkersForTable(tx, role, qtid2) + assert.NoError(t, err) + assert.ElementsMatch(t, []dax.WorkerInfo{ + {Address: "n1", Jobs: []dax.Job{fn2("s11"), fn2("s13"), fn2("s15"), fn2("s9")}}, + {Address: "n2", Jobs: []dax.Job{fn2("s10"), fn2("s12"), fn2("s14")}}, + }, workers) + + // No match + qtid0 := dax.NewQualifiedTableID(qdbid, "bad") + workers, err = bal.WorkersForTable(tx, role, qtid0) + assert.NoError(t, err) + assert.ElementsMatch(t, []dax.WorkerInfo{}, workers) + + assert.NoError(t, tx.Commit()) + }) + + t.Run("Balance", func(t *testing.T) { + dbOptions := dax.DatabaseOptions{ + WorkersMin: 2, + WorkersMax: 2, + } + qdb := &dax.QualifiedDatabase{ + OrganizationID: orgID, + Database: dax.Database{ + ID: dbID, + Name: dbName, + Options: dbOptions, + }, + } + + qdbid := qdb.QualifiedID() + + schemar, scleanup := daxtest.NewSchemar(t) + defer scleanup() + + role := dax.RoleTypeCompute + + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer(db, schemar, logger.NewStandardLogger(os.Stderr)) + + tx, err := db.BeginTx(ctx, true) + assert.NoError(t, err) + defer tx.Rollback() + + assert.NoError(t, schemar.CreateDatabase(tx, qdb)) + + node1 := &dax.Node{ + Address: "n1", + RoleTypes: []dax.RoleType{role}, + } + + node2 := &dax.Node{ + Address: "n2", + RoleTypes: []dax.RoleType{role}, + } + + node3 := &dax.Node{ + Address: "n3", + RoleTypes: []dax.RoleType{role}, + } + + // Add two workers with some jobs evenly spread across them. + _, err = bal.AddWorker(tx, node1) + assert.NoError(t, err) + _, err = bal.AddWorker(tx, node2) + assert.NoError(t, err) + for i := 0; i < 13; i++ { + job := fmt.Sprintf("s%d", i) + _, err = bal.AddJobs(tx, role, qtid, dax.Job(job)) + assert.NoError(t, err) + } + + exp := dax.WorkerInfo{ + Address: "n1", + Jobs: []dax.Job{"s0", "s10", "s12", "s2", "s4", "s6", "s8"}, + } + ws, err := bal.WorkerState(tx, role, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + Address: "n2", + Jobs: []dax.Job{"s1", "s11", "s3", "s5", "s7", "s9"}, + } + ws, err = bal.WorkerState(tx, role, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Update database options on schemar so min worker for db is 3. + t.Run(fmt.Sprintf("test-%s", "set database options db min 3"), func(t *testing.T) { + dbOptions.WorkersMin = 3 + dbOptions.WorkersMax = 3 + assert.NoError(t, schemar.SetDatabaseOptions(tx, qdb.QualifiedID(), dbOptions)) + }) + + // Now, add a worker and confirm that it has received some jobs. + _, err = bal.AddWorker(tx, node3) + assert.NoError(t, err) + exp = dax.WorkerInfo{ + Address: "n3", + Jobs: []dax.Job{"s6", "s7", "s8", "s9"}, + } + ws, err = bal.WorkerState(tx, role, "n3") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Finally, call Balance() and confirm that the appropriate jobs got are + // as expected (actually, this is no longer needed here because we + // automatically balance when we add workers, but calling it should + // effectively be a no-op). + _, err = bal.BalanceDatabase(tx, qdbid) + assert.NoError(t, err) + + exp = dax.WorkerInfo{ + Address: "n1", + Jobs: []dax.Job{"s0", "s10", "s12", "s2", "s4"}, + } + ws, err = bal.WorkerState(tx, role, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + Address: "n2", + Jobs: []dax.Job{"s1", "s11", "s3", "s5"}, + } + ws, err = bal.WorkerState(tx, role, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + Address: "n3", + Jobs: []dax.Job{"s6", "s7", "s8", "s9"}, + } + ws, err = bal.WorkerState(tx, role, "n3") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + assert.NoError(t, tx.Commit()) + }) +} diff --git a/dax/mds/controller/balancer/boltdb/balancer.go b/dax/mds/controller/balancer/boltdb/balancer.go new file mode 100644 index 000000000..bb2725718 --- /dev/null +++ b/dax/mds/controller/balancer/boltdb/balancer.go @@ -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) +} diff --git a/dax/mds/controller/naive/types.go b/dax/mds/controller/balancer/types.go similarity index 74% rename from dax/mds/controller/naive/types.go rename to dax/mds/controller/balancer/types.go index 6659b7f9d..90fcafa6e 100644 --- a/dax/mds/controller/naive/types.go +++ b/dax/mds/controller/balancer/types.go @@ -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++ diff --git a/dax/mds/controller/config.go b/dax/mds/controller/config.go index 1d77fc73c..e67d3ddb0 100644 --- a/dax/mds/controller/config.go +++ b/dax/mds/controller/config.go @@ -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 diff --git a/dax/mds/controller/controller.go b/dax/mds/controller/controller.go index df71b5bda..abb4332f3 100644 --- a/dax/mds/controller/controller.go +++ b/dax/mds/controller/controller.go @@ -5,7 +5,6 @@ import ( "context" "fmt" "sort" - "sync" "time" "github.com/featurebasedb/featurebase/v3/dax" @@ -17,22 +16,17 @@ import ( ) type Controller struct { - // mu is primarily to protect against conflicting reads/writes to the nodes - // map. The balancers map is currently never written to after - // initialization. - mu sync.RWMutex - // Schemar used by the controller to get table information. The controller // should NOT call Schemar methods which modify data. Schema mutations are // made outside of the controller (at this point that happens in MDS). Schemar schemar.Schemar - // nodes is the map of nodes, by address, which have registered with the - // controller. + // nodeService is the interface to working with nodes, by address, which + // have registered with the controller. nodeService dax.NodeService - ComputeBalancer Balancer - TranslateBalancer Balancer + boltDB *boltdb.DB + Balancer Balancer // Director is used to send directives to computer workers. Director Director @@ -61,8 +55,8 @@ func New(cfg Config) *Controller { c := &Controller{ Schemar: schemar.NewNopSchemar(), - ComputeBalancer: cfg.ComputeBalancer, - TranslateBalancer: cfg.TranslateBalancer, + boltDB: cfg.BoltDB, + Balancer: cfg.Balancer, Director: NewNopDirector(), @@ -119,24 +113,16 @@ func (c *Controller) Stop() { close(c.stopping) } -func (c *Controller) balancerForRole(rt dax.RoleType) (Balancer, error) { - var bal Balancer - if rt == dax.RoleTypeCompute { - bal = c.ComputeBalancer - } else if rt == dax.RoleTypeTranslate { - bal = c.TranslateBalancer - } else { - return nil, errors.Errorf("unknown role type: '%s'", rt) - } - return bal, nil -} - // RegisterNodes adds nodes to the controller's list of registered // nodes. func (c *Controller) RegisterNodes(ctx context.Context, nodes ...*dax.Node) error { - c.logger.Printf("c.RegisterNodes(): %+v", nodes) - c.mu.Lock() - defer c.mu.Unlock() + c.logger.Printf("c.RegisterNodes(): %s", dax.Nodes(nodes)) + + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // Validate input. for _, n := range nodes { @@ -159,8 +145,8 @@ func (c *Controller) RegisterNodes(ctx context.Context, nodes ...*dax.Node) erro // Create node if we don't already have it for _, n := range nodes { - if node, _ := c.nodeService.ReadNode(ctx, n.Address); node == nil { - if err := c.nodeService.CreateNode(ctx, n.Address, n); err != nil { + if node, _ := c.nodeService.ReadNode(tx, n.Address); node == nil { + if err := c.nodeService.CreateNode(tx, n.Address, n); err != nil { return errors.Wrapf(err, "creating node: %s", n.Address) } @@ -175,33 +161,24 @@ func (c *Controller) RegisterNodes(ctx context.Context, nodes ...*dax.Node) erro // diffByAddr keeps track of the diffs that have been applied to each // specific address. + // TODO(tlt): I don't understand why we're keeping track of the + // dax.WorkerDiff here (as opposed to just the unique Address) because it + // doesn't ever seem to be used. diffByAddr := make(map[dax.Address]dax.WorkerDiff) for _, n := range nodes { - for _, rt := range n.RoleTypes { - balancer, err := c.balancerForRole(rt) - if err != nil { - return errors.Wrap(err, "getting balancer") - } - adiffs, err := balancer.AddWorker(ctx, n.Address) - if err != nil { - return errors.Wrap(err, "adding worker") - } + adiffs, err := c.Balancer.AddWorker(tx, n) + if err != nil { + return errors.Wrap(err, "adding worker") + } - // Rebalance so existing jobs can be spread evenly across all nodes, - // including the node being registered. - bdiffs, err := balancer.Balance(ctx) - if err != nil { - return errors.Wrap(err, "balancing") - } - for _, diff := range append(adiffs, bdiffs...) { - existingDiff, ok := diffByAddr[dax.Address(diff.WorkerID)] - if !ok { - existingDiff.WorkerID = diff.WorkerID - } - existingDiff.Add(diff) - diffByAddr[dax.Address(diff.WorkerID)] = existingDiff + for _, diff := range adiffs { + existingDiff, ok := diffByAddr[dax.Address(diff.Address)] + if !ok { + existingDiff.Address = diff.Address } + existingDiff.Add(diff) + diffByAddr[dax.Address(diff.Address)] = existingDiff } } @@ -217,7 +194,7 @@ func (c *Controller) RegisterNodes(ctx context.Context, nodes ...*dax.Node) erro } // Tell the poller about the new nodes. - if err := c.poller.AddAddresses(ctx, addrs...); err != nil { + if err := c.poller.AddAddresses(tx.Context(), addrs...); err != nil { return NewErrInternal(err.Error()) } @@ -241,20 +218,17 @@ func (c *Controller) RegisterNodes(ctx context.Context, nodes ...*dax.Node) erro // Get the current job assignments for this worker and send that to the node // as a Directive. - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } - return nil + return tx.Commit() } // RegisterNode adds a node to the controller's list of registered // nodes. It makes no guarantees about when the node will actually be // used for anything or assigned any jobs. func (c *Controller) RegisterNode(ctx context.Context, n *dax.Node) error { - c.mu.Lock() - defer c.mu.Unlock() - // Validate input. if n.Address == "" { return NewErrNodeKeyInvalid(n.Address) @@ -268,7 +242,13 @@ func (c *Controller) RegisterNode(ctx context.Context, n *dax.Node) error { } } - if node, _ := c.nodeService.ReadNode(ctx, n.Address); node != nil { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if node, _ := c.nodeService.ReadNode(tx, n.Address); node != nil { return nil } @@ -284,15 +264,18 @@ func (c *Controller) RegisterNode(ctx context.Context, n *dax.Node) error { // from its list (perhaps due to a network fault) and therefore the node needs // to be re-registered. func (c *Controller) CheckInNode(ctx context.Context, n *dax.Node) error { - c.mu.RLock() - defer c.mu.RUnlock() + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // If we already know about this node, just no-op. In the future, we may // want this check-in payload to include things like the compute node's // Directive; then we could check that the compute node is actually doing // what we expect it to be doing. But for now, we're just checking that we // know about the compute node at all. - if node, _ := c.nodeService.ReadNode(ctx, n.Address); node != nil { + if node, _ := c.nodeService.ReadNode(tx, n.Address); node != nil { return nil } @@ -304,8 +287,11 @@ func (c *Controller) CheckInNode(ctx context.Context, n *dax.Node) error { // DeregisterNodes removes nodes from the controller's list of registered nodes. // It sends directives to the removed nodes, but ignores errors. func (c *Controller) DeregisterNodes(ctx context.Context, addresses ...dax.Address) error { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // workerSet maintains the set of workers which have a job assignment change // and therefore need to be sent an updated Directive. @@ -325,40 +311,21 @@ func (c *Controller) DeregisterNodes(ctx context.Context, addresses ...dax.Addre // address this confusion. // workerSet.Add(address) - // Ensure the host:port is currently registered. - n, err := c.nodeService.ReadNode(ctx, address) + rdiffs, err := c.Balancer.RemoveWorker(tx, address) if err != nil { - return errors.Wrapf(err, "reading the node for address: %s", address) + return errors.Wrapf(err, "removing worker: %s", address) } - for _, rt := range n.RoleTypes { - balancer, err := c.balancerForRole(rt) - if err != nil { - c.logger.Printf("Unsupported role type in DeregisterNode: '%s'", rt) - // Skip any role types which aren't currently supported by a balancer. - continue - } - rdiffs, err := balancer.RemoveWorker(ctx, address) - if err != nil { - return errors.Wrap(err, "removing worker") - } - // Rebalance so any jobs that were assigned to the node being deregistered - // get assigned to another node. - bdiffs, err := balancer.Balance(ctx) - if err != nil { - return errors.Wrap(err, "balancing") - } - // we assume that the job names are different between the - // different role types so we don't have to track each - // role separately which would be annoying. - for _, diff := range append(rdiffs, bdiffs...) { - existingDiff, ok := diffByAddr[dax.Address(diff.WorkerID)] - if !ok { - existingDiff.WorkerID = diff.WorkerID - } - existingDiff.Add(diff) - diffByAddr[dax.Address(diff.WorkerID)] = existingDiff + // we assume that the job names are different between the + // different role types so we don't have to track each + // role separately which would be annoying. + for _, diff := range rdiffs { + existingDiff, ok := diffByAddr[dax.Address(diff.Address)] + if !ok { + existingDiff.Address = diff.Address } + existingDiff.Add(diff) + diffByAddr[dax.Address(diff.Address)] = existingDiff } } @@ -374,18 +341,18 @@ func (c *Controller) DeregisterNodes(ctx context.Context, addresses ...dax.Addre } for _, address := range addresses { - if err := c.nodeService.DeleteNode(ctx, address); err != nil { + if err := c.nodeService.DeleteNode(tx, address); err != nil { return errors.Wrapf(err, "deleting node at address: %s", address) } } - if err := c.poller.RemoveAddresses(ctx, addresses...); err != nil { + if err := c.poller.RemoveAddresses(tx.Context(), addresses...); err != nil { return NewErrInternal(err.Error()) } // No need to send Directives if nothing has ultimately changed. if len(workerSet) == 0 { - return nil + return tx.Commit() } // Convert the slice of addresses into a slice of addressMethod containing @@ -403,234 +370,92 @@ func (c *Controller) DeregisterNodes(ctx context.Context, addresses ...dax.Addre // Get the current job assignments for these workers and send them to the // nodes as Directives. - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } - return nil + return tx.Commit() } -// Nodes returns the list of assigned nodes responsible for the jobs included in -// the given role. If createMissing is true, the Controller will create new jobs -// for any of which it isn't currently aware. -func (c *Controller) Nodes(ctx context.Context, role dax.Role, createMissing bool) ([]dax.AssignedNode, error) { - nodes := []dax.AssignedNode{} - var err error +func (c *Controller) nodesTranslateReadOrWrite(tx dax.Transaction, role *dax.TranslateRole, qdbid dax.QualifiedDatabaseID, createMissing bool, asWrite bool) ([]dax.AssignedNode, bool, error) { + qtid := role.TableKey.QualifiedTableID() + roleType := dax.RoleTypeTranslate - switch v := role.(type) { - case *dax.ComputeRole: - nodes, err = c.nodesCompute(ctx, v, createMissing) - if err != nil { - return nil, errors.Wrap(err, "getting compute nodes") - } - - case *dax.TranslateRole: - nodes, err = c.nodesTranslate(ctx, v, createMissing) - if err != nil { - return nil, errors.Wrap(err, "getting translate nodes") - } - } - - return nodes, nil -} - -// nodesTranslate is like nodesCompute. See the comments there. -func (c *Controller) nodesTranslate(ctx context.Context, role *dax.TranslateRole, createMissing bool) ([]dax.AssignedNode, error) { - // Try calling c.nodesTranslate as a read first. If we don't have to actually - // create any missing partitions, then we won't have to obtain a write lock. - translateNodes, retryAsWrite, err := c.nodesTranslateReadOrWrite(ctx, role, createMissing, false) - if err != nil { - return nil, errors.Wrap(err, "getting translate nodes read or write") - } - - if retryAsWrite { - translateNodes, _, err = c.nodesTranslateReadOrWrite(ctx, role, createMissing, retryAsWrite) - if err != nil { - return nil, errors.Wrap(err, "getting translate nodes read or write retry") - } - } - - return translateNodes, nil -} - -func (c *Controller) nodesTranslateReadOrWrite(ctx context.Context, role *dax.TranslateRole, createMissing bool, asWrite bool) ([]dax.AssignedNode, bool, error) { - if asWrite { - c.mu.Lock() - defer c.mu.Unlock() - } else { - c.mu.RLock() - defer c.mu.RUnlock() - } - - nodes := []dax.AssignedNode{} - - bal := c.TranslateBalancer - - //inJobs := NewStringSet() inJobs := dax.NewSet[dax.Job]() for _, p := range role.Partitions { partitionString := partition(role.TableKey, p).String() inJobs.Add(dax.Job(partitionString)) } - workers, err := bal.WorkersForJobs(ctx, inJobs.Sorted()) + workers, err := c.Balancer.WorkersForJobs(tx, roleType, qdbid, inJobs.Sorted()...) if err != nil { return nil, false, errors.Wrap(err, "getting workers for jobs") } - if createMissing { - // If any provided jobs were not returned in the WorkersForJobs - // request, then create those. - outJobs := dax.NewSet[dax.Job]() - for _, worker := range workers { - for _, job := range worker.Jobs { - outJobs.Add(job) - } - } - - missed := inJobs.Minus(outJobs).Sorted() - - if len(missed) > 0 { - // If we are currently under a read lock, and we get to this point, - // it means that we have partitions which need to be assigned (and - // directives sent) to workers. In that case, we need to abort this - // method run and notify the caller to rety as a write. - if !asWrite { - return nil, true, nil - } - - sort.Slice(missed, func(i, j int) bool { return missed[i] < missed[j] }) - - workerSet := NewAddressSet() - for _, job := range missed { - j, err := decodePartition(job) - if err != nil { - return nil, false, NewErrInternal(err.Error()) - } - diffs, err := bal.AddJobs(ctx, j) - if err != nil { - return nil, false, errors.Wrap(err, "adding job") - } - for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) - } - } - - // Convert the slice of addresses into a slice of addressMethod containing - // the appropriate method. - addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - - if err := c.sendDirectives(ctx, addressMethods...); err != nil { - return nil, false, NewErrDirectiveSendFailure(err.Error()) - } - - // Re-run WorkersForJobs. - workers, err = bal.WorkersForJobs(ctx, inJobs.Sorted()) - if err != nil { - return nil, false, errors.Wrap(err, "getting workers for jobs") - } + outJobs := dax.NewSet[dax.Job]() + for _, worker := range workers { + for _, job := range worker.Jobs { + outJobs.Add(job) } } - for _, worker := range workers { - // covert worker.Jobs []string to map[string][]Partition - translateMap := make(map[dax.TableKey]dax.PartitionNums) - for _, job := range worker.Jobs { + missed := inJobs.Minus(outJobs).Sorted() + if !createMissing && len(missed) > 0 { + return nil, false, NewErrUnassignedJobs(missed) + } + + // If any provided jobs were not returned in the WorkersForJobs request, + // then create those. + if createMissing && len(missed) > 0 { + // If we are currently under a read lock, and we get to this point, it + // means that we have partitions which need to be assigned (and + // directives sent) to workers. In that case, we need to abort this + // method run and notify the caller to rety as a write. + if !asWrite { + return nil, true, nil + } + + sort.Slice(missed, func(i, j int) bool { return missed[i] < missed[j] }) + + workerSet := NewAddressSet() + for _, job := range missed { j, err := decodePartition(job) if err != nil { return nil, false, NewErrInternal(err.Error()) } - - translateMap[j.table()] = append(translateMap[j.table()], j.partitionNum()) + diffs, err := c.Balancer.AddJobs(tx, roleType, qtid, j.Job()) + if err != nil { + return nil, false, errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.Address)) + } } - for table, partitions := range translateMap { - // Sort the partitions int slice before returning it. - sort.Sort(partitions) + // Convert the slice of addresses into a slice of addressMethod + // containing the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - nodes = append(nodes, dax.AssignedNode{ - Address: dax.Address(worker.ID), - Role: &dax.TranslateRole{ - TableKey: table, - Partitions: partitions, - }, - }) + if err := c.sendDirectives(tx, addressMethods...); err != nil { + return nil, false, NewErrDirectiveSendFailure(err.Error()) } - } - return nodes, false, nil -} - -// nodesCompute tries to get the list of compute nodes under a read lock. If the -// call into c.nodesComputeReadOrWrite() comes back with `retryAsWrite = true`, -// then it gets called again but with a write lock so that sendDirective can -// happen, and the receiving compute node can get apply the latest schema, -// without encountering race conditions. -// -// Really, we shouldn't have to rely on the directive being applied within a -// mu.Lock(). Instead, if a client (or the IDK) tries to perform some action on -// a compute node (for example, ingesting data to an index/field), if that -// action fails because the schema on the compute node is not in sync, or if the -// compute node is completely unavailable, the client should ask mds for updated -// node information and keep retrying. Basically, what I'm saying is that a lot -// of the mu.Lock()s in this file can be changed back to mu.RLock()s, and the -// SendDirective() can happen asyncronously without worring about race -// conditions. -// -// The race condition happened when concurrent requests to Nodes() occurred and -// the order of events was: -// - req1 wants node for [idx, 0] -// - req2 wants node for [idx, 0] -// - (req1) [idx, 0] registered in controller to node A -// - directive sent to node A to create index [idx] -// -// - (req2) receives: node A -// - (req2) tries to ingest data to node A [idx, 0] -// **** RACE: [idx] does not exist because (req1) directive step is not compete -func (c *Controller) nodesCompute(ctx context.Context, role *dax.ComputeRole, createMissing bool) ([]dax.AssignedNode, error) { - if len(role.Shards) == 0 { - return c.nodesForTableKey(ctx, role.TableKey) - } - // Try calling c.nodesCompute as a read first. If we don't have to actually - // create any missing shards, then we won't have to obtain a write lock. - computeNodes, retryAsWrite, err := c.nodesComputeReadOrWrite(ctx, role, createMissing, false) - if err != nil { - return nil, errors.Wrap(err, "getting compute nodes read or write") - } - - if retryAsWrite { - computeNodes, _, err = c.nodesComputeReadOrWrite(ctx, role, createMissing, retryAsWrite) + // Re-run WorkersForJobs. + workers, err = c.Balancer.WorkersForJobs(tx, roleType, qdbid, inJobs.Sorted()...) if err != nil { - return nil, errors.Wrap(err, "getting compute nodes read or write retry") + return nil, false, errors.Wrap(err, "getting workers for jobs") } } - return computeNodes, nil -} - -func (c *Controller) nodesForTableKey(ctx context.Context, tk dax.TableKey) ([]dax.AssignedNode, error) { - bal := c.ComputeBalancer - workers, err := bal.WorkersForJobPrefix(ctx, string(tk)) - if err != nil { - return nil, errors.Wrapf(err, "getting workers for table: '%s'", tk) - } - - return c.workersToAssignedNodes(ctx, workers) - + nodes, err := c.translateWorkersToAssignedNodes(tx, workers) + return nodes, false, errors.Wrap(err, "converting to assigned nodes") } // nodesComputeReadOrWrite contains the logic for the c.nodesCompute() method, // but it supports being called with either a read or write lock. -func (c *Controller) nodesComputeReadOrWrite(ctx context.Context, role *dax.ComputeRole, createMissing bool, asWrite bool) ([]dax.AssignedNode, bool, error) { - if asWrite { - c.mu.Lock() - defer c.mu.Unlock() - } else { - c.mu.RLock() - defer c.mu.RUnlock() - } - - bal := c.ComputeBalancer +func (c *Controller) nodesComputeReadOrWrite(tx dax.Transaction, role *dax.ComputeRole, qdbid dax.QualifiedDatabaseID, createMissing bool, asWrite bool) ([]dax.AssignedNode, bool, error) { + qtid := role.TableKey.QualifiedTableID() + roleType := dax.RoleTypeCompute inJobs := dax.NewSet[dax.Job]() for _, s := range role.Shards { @@ -638,7 +463,7 @@ func (c *Controller) nodesComputeReadOrWrite(ctx context.Context, role *dax.Comp inJobs.Add(dax.Job(shardString)) } - workers, err := bal.WorkersForJobs(ctx, inJobs.Sorted()) + workers, err := c.Balancer.WorkersForJobs(tx, roleType, qdbid, inJobs.Sorted()...) if err != nil { return nil, false, errors.Wrap(err, "getting workers for jobs") } @@ -656,56 +481,54 @@ func (c *Controller) nodesComputeReadOrWrite(ctx context.Context, role *dax.Comp return nil, false, NewErrUnassignedJobs(missed) } - if createMissing { - // If any provided jobs were not returned in the WorkersForJobs - // request, then create those. - if len(missed) > 0 { - // If we are currently under a read lock, and we get to this point, - // it means that we have shards which need to be assigned (and - // directives sent) to workers. In that case, we need to abort this - // method run and notify the caller to rety as a write. - if !asWrite { - return nil, true, nil - } + // If any provided jobs were not returned in the WorkersForJobs request, + // then create those. + if createMissing && len(missed) > 0 { + // If we are currently under a read lock, and we get to this point, it + // means that we have shards which need to be assigned (and directives + // sent) to workers. In that case, we need to abort this method run and + // notify the caller to rety as a write. + if !asWrite { + return nil, true, nil + } - sort.Slice(missed, func(i, j int) bool { return missed[i] < missed[j] }) + sort.Slice(missed, func(i, j int) bool { return missed[i] < missed[j] }) - workerSet := NewAddressSet() - for _, job := range missed { - j, err := decodeShard(job) - if err != nil { - return nil, false, NewErrInternal(err.Error()) - } - diffs, err := bal.AddJobs(ctx, j) - if err != nil { - return nil, false, errors.Wrap(err, "adding job") - } - for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) - } - } - - // Convert the slice of addresses into a slice of addressMethod containing - // the appropriate method. - addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - - if err := c.sendDirectives(ctx, addressMethods...); err != nil { - return nil, false, NewErrDirectiveSendFailure(err.Error()) - } - - // Re-run WorkersForJobs. - workers, err = bal.WorkersForJobs(ctx, inJobs.Sorted()) + workerSet := NewAddressSet() + for _, job := range missed { + j, err := decodeShard(job) if err != nil { - return nil, false, errors.Wrap(err, "getting workers for jobs") + return nil, false, NewErrInternal(err.Error()) } + diffs, err := c.Balancer.AddJobs(tx, roleType, qtid, j.Job()) + if err != nil { + return nil, false, errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.Address)) + } + } + + // Convert the slice of addresses into a slice of addressMethod + // containing the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(tx, addressMethods...); err != nil { + return nil, false, NewErrDirectiveSendFailure(err.Error()) + } + + // Re-run WorkersForJobs. + workers, err = c.Balancer.WorkersForJobs(tx, roleType, qdbid, inJobs.Sorted()...) + if err != nil { + return nil, false, errors.Wrap(err, "getting workers for jobs") } } - nodes, err := c.workersToAssignedNodes(ctx, workers) + nodes, err := c.computeWorkersToAssignedNodes(tx, workers) return nodes, false, errors.Wrap(err, "converting to assigned nodes") } -func (c *Controller) workersToAssignedNodes(ctx context.Context, workers []dax.WorkerInfo) ([]dax.AssignedNode, error) { +func (c *Controller) computeWorkersToAssignedNodes(tx dax.Transaction, workers []dax.WorkerInfo) ([]dax.AssignedNode, error) { nodes := []dax.AssignedNode{} for _, worker := range workers { // convert worker.Jobs []string to map[TableName][]Shard @@ -724,7 +547,7 @@ func (c *Controller) workersToAssignedNodes(ctx context.Context, workers []dax.W sort.Sort(shards) nodes = append(nodes, dax.AssignedNode{ - Address: dax.Address(worker.ID), + Address: dax.Address(worker.Address), Role: &dax.ComputeRole{ TableKey: table, Shards: shards, @@ -735,11 +558,100 @@ func (c *Controller) workersToAssignedNodes(ctx context.Context, workers []dax.W return nodes, nil } -// CreateTable adds a table to the versionStore and schemar, and then sends directives -// to all affected nodes based on the change. +func (c *Controller) translateWorkersToAssignedNodes(tx dax.Transaction, workers []dax.WorkerInfo) ([]dax.AssignedNode, error) { + nodes := []dax.AssignedNode{} + for _, worker := range workers { + // covert worker.Jobs []string to map[string][]Partition + translateMap := make(map[dax.TableKey]dax.PartitionNums) + for _, job := range worker.Jobs { + j, err := decodePartition(job) + if err != nil { + return nil, NewErrInternal(err.Error()) + } + + translateMap[j.table()] = append(translateMap[j.table()], j.partitionNum()) + } + + for table, partitions := range translateMap { + // Sort the partitions int slice before returning it. + sort.Sort(partitions) + + nodes = append(nodes, dax.AssignedNode{ + Address: dax.Address(worker.Address), + Role: &dax.TranslateRole{ + TableKey: table, + Partitions: partitions, + }, + }) + } + } + return nodes, nil +} + +// CreateDatabase adds a database to the schemar. +func (c *Controller) CreateDatabase(ctx context.Context, qdb *dax.QualifiedDatabase) error { + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if err := c.Schemar.CreateDatabase(tx, qdb); err != nil { + return errors.Wrap(err, "creating database in schemar") + } + + return tx.Commit() +} + +// DatabaseByID returns the database for the given id. +func (c *Controller) DatabaseByID(ctx context.Context, qdbid dax.QualifiedDatabaseID) (*dax.QualifiedDatabase, error) { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + qdb, err := c.Schemar.DatabaseByID(tx, qdbid) + if err != nil { + return nil, errors.Wrap(err, "getting database from schemar") + } + + return qdb, nil +} + +// SetDatabaseOptions sets the options on the given database. +func (c *Controller) SetDatabaseOptions(ctx context.Context, qdbid dax.QualifiedDatabaseID, opts dax.DatabaseOptions) error { + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if err := c.Schemar.SetDatabaseOptions(tx, qdbid, opts); err != nil { + return errors.Wrap(err, "setting database options") + } + + return tx.Commit() +} + +// CreateTable adds a table to the schemar, and then sends directives to all +// affected nodes based on the change. func (c *Controller) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + // Create Table ID. + if _, err := qtbl.CreateID(); err != nil { + return errors.Wrap(err, "creating table ID") + } + + // Create the table in schemar. + if err := c.Schemar.CreateTable(tx, qtbl); err != nil { + return errors.Wrapf(err, "creating table: %s", qtbl) + } // If the table is keyed, add partitions to the balancer. if qtbl.StringKeys() { @@ -753,24 +665,24 @@ func (c *Controller) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) partitionsToAdd[partitionNum] = dax.PartitionNum(partitionNum) } - stringers := make([]fmt.Stringer, 0, len(partitionsToAdd)) + jobs := make([]dax.Job, 0, len(partitionsToAdd)) for _, p := range partitionsToAdd { - stringers = append(stringers, partition(qtbl.Key(), p)) + jobs = append(jobs, partition(qtbl.Key(), p).Job()) } - diffs, err := c.TranslateBalancer.AddJobs(ctx, stringers...) + diffs, err := c.Balancer.AddJobs(tx, dax.RoleTypeTranslate, qtbl.QualifiedID(), jobs...) if err != nil { return errors.Wrap(err, "adding job") } for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) + workerSet.Add(dax.Address(diff.Address)) } // Convert the slice of addresses into a slice of addressMethod containing // the appropriate method. addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } } @@ -790,34 +702,37 @@ func (c *Controller) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) // We don't currently use the returned diff, other than to determine // which worker was affected, because we send the full Directive // every time. - diffs, err := c.TranslateBalancer.AddJobs(ctx, partition(qtbl.Key(), p)) + diffs, err := c.Balancer.AddJobs(tx, dax.RoleTypeTranslate, qtbl.QualifiedID(), partition(qtbl.Key(), p).Job()) if err != nil { return errors.Wrap(err, "adding job") } for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) + workerSet.Add(dax.Address(diff.Address)) } // Convert the slice of addresses into a slice of addressMethod containing // the appropriate method. addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } } - return nil + return tx.Commit() } // DropTable removes a table from the schema and sends directives to all affected // nodes based on the change. func (c *Controller) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // Get the table from the schemar. - if _, err := c.Schemar.Table(ctx, qtid); err != nil { + if _, err := c.Schemar.Table(tx, qtid); err != nil { return errors.Wrapf(err, "table not in schemar: %s", qtid) } @@ -825,90 +740,70 @@ func (c *Controller) DropTable(ctx context.Context, qtid dax.QualifiedTableID) e // and therefore need to be sent an updated Directive. workerSet := NewAddressSet() - diffs, err := c.ComputeBalancer.RemoveJobs(ctx, string(qtid.Key())) + diffs, err := c.Balancer.RemoveJobs(tx, dax.RoleTypeCompute, qtid) if err != nil { - return errors.Wrap(err, "removing jobs") + return errors.Wrapf(err, "removing compute jobs for table: %s", qtid) } for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) + workerSet.Add(dax.Address(diff.Address)) } - diffs, err = c.TranslateBalancer.RemoveJobs(ctx, string(qtid.Key())) + diffs, err = c.Balancer.RemoveJobs(tx, dax.RoleTypeTranslate, qtid) if err != nil { - return errors.Wrap(err, "removing job") + return errors.Wrapf(err, "removing translate jobs for table: %s", qtid) } for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) + workerSet.Add(dax.Address(diff.Address)) } // Convert the slice of addresses into a slice of addressMethod containing // the appropriate method. addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } - return nil + // Remove table from schemar. + if err := c.Schemar.DropTable(tx, qtid); err != nil { + return errors.Wrapf(err, "dropping table from schemar: %s", qtid) + } + + return tx.Commit() } // Table returns a table by quaified table id. func (c *Controller) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // Get the table from the schemar. - return c.Schemar.Table(ctx, qtid) + return c.Schemar.Table(tx, qtid) } // Tables returns a list of tables by name. -func (c *Controller) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { - c.mu.Lock() - defer c.mu.Unlock() +func (c *Controller) Tables(ctx context.Context, qdbid dax.QualifiedDatabaseID, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // Get the tables from the schemar. - return c.Schemar.Tables(ctx, qual, ids...) -} - -// AddShards registers the table/shard combinations with the controller and -// sends the necessary directive. -func (c *Controller) AddShards(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.ShardNum) error { - c.mu.Lock() - defer c.mu.Unlock() - - // workerSet maintains the set of workers which have a job assignment change - // and therefore need to be sent an updated Directive. - workerSet := NewAddressSet() - - for _, s := range shards { - // We don't currently use the returned diff, other than to determine - // which worker was affected, because we send the full Directive every - // time. - diffs, err := c.ComputeBalancer.AddJobs(ctx, shard(qtid.Key(), s)) - if err != nil { - return errors.Wrap(err, "adding job") - } - for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) - } - } - - // Convert the slice of addresses into a slice of addressMethod containing - // the appropriate method. - addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - - if err := c.sendDirectives(ctx, addressMethods...); err != nil { - return NewErrDirectiveSendFailure(err.Error()) - } - - return nil + return c.Schemar.Tables(tx, qdbid, ids...) } // RemoveShards deregisters the table/shard combinations with the controller and // sends the necessary directives. func (c *Controller) RemoveShards(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.ShardNum) error { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() // workerSet maintains the set of workers which have a job assignment change // and therefore need to be sent an updated Directive. @@ -918,12 +813,12 @@ func (c *Controller) RemoveShards(ctx context.Context, qtid dax.QualifiedTableID // We don't currently use the returned diff, other than to determine // which worker was affected, because we send the full Directive every // time. - diffs, err := c.ComputeBalancer.RemoveJob(ctx, shard(qtid.Key(), s)) + diffs, err := c.Balancer.RemoveJobs(tx, dax.RoleTypeCompute, qtid, shard(qtid.Key(), s).Job()) if err != nil { return errors.Wrap(err, "removing job") } for _, diff := range diffs { - workerSet.Add(dax.Address(diff.WorkerID)) + workerSet.Add(dax.Address(diff.Address)) } } @@ -931,22 +826,22 @@ func (c *Controller) RemoveShards(ctx context.Context, qtid dax.QualifiedTableID // the appropriate method. addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } - return nil + return tx.Commit() } // sendDirectives sends a directive (based on the current balancer state) to // each of the nodes provided. -func (c *Controller) sendDirectives(ctx context.Context, addrs ...addressMethod) error { +func (c *Controller) sendDirectives(tx dax.Transaction, addrs ...addressMethod) error { // If nodes is empty, return early. if len(addrs) == 0 { return nil } - directives, err := c.buildDirectives(ctx, addrs) + directives, err := c.buildDirectives(tx, addrs) if err != nil { return errors.Wrap(err, "building directives") } @@ -957,7 +852,7 @@ func (c *Controller) sendDirectives(ctx context.Context, addrs ...addressMethod) i := i dir := dir eg.Go(func() error { - errs[i] = c.Director.SendDirective(ctx, dir) + errs[i] = c.Director.SendDirective(tx.Context(), dir) if dir.IsEmpty() { errs[i] = nil } @@ -1016,11 +911,11 @@ func applyAddressMethod(addrs []dax.Address, method dax.DirectiveMethod) []addre // buildDirectives builds a list of directives for the given addrs (i.e. nodes) // using information (i.e. current state) from the balancers. -func (c *Controller) buildDirectives(ctx context.Context, addrs []addressMethod) ([]*dax.Directive, error) { +func (c *Controller) buildDirectives(tx dax.Transaction, addrs []addressMethod) ([]*dax.Directive, error) { directives := make([]*dax.Directive, len(addrs)) for i, addressMethod := range addrs { - dVersion, err := c.directiveVersion.Increment(ctx, 1) + dVersion, err := c.directiveVersion.Increment(tx, 1) if err != nil { return nil, errors.Wrap(err, "incrementing directive version") } @@ -1053,11 +948,7 @@ func (c *Controller) buildDirectives(ctx context.Context, addrs []addressMethod) ownsPartition0 := make(map[dax.TableKey]struct{}, 0) for _, rt := range supportedRoleTypes { - bal, err := c.balancerForRole(rt) - if err != nil { - return nil, errors.Wrap(err, "getting balancer") - } - w, err := bal.WorkerState(ctx, dax.Worker(addressMethod.address.String())) + w, err := c.Balancer.WorkerState(tx, rt, addressMethod.address) if err != nil { return nil, errors.Wrapf(err, "getting worker state: %s", addressMethod.address) } @@ -1130,7 +1021,7 @@ func (c *Controller) buildDirectives(ctx context.Context, addrs []addressMethod) // everything that way for now. for tkey := range ownsPartition0 { qtid := tkey.QualifiedTableID() - table, err := c.Schemar.Table(ctx, qtid) + table, err := c.Schemar.Table(tx, qtid) if err != nil { return nil, errors.Wrapf(err, "getting table: %s", tkey) } @@ -1165,7 +1056,7 @@ func (c *Controller) buildDirectives(ctx context.Context, addrs []addressMethod) if len(tableSet) > 0 { dTables := make([]*dax.QualifiedTable, 0) for qual, tblIDs := range tableSet.QualifiedSortedSlice() { - qtbls, err := c.Schemar.Tables(ctx, qual, tblIDs...) + qtbls, err := c.Schemar.Tables(tx, qual, tblIDs...) if err != nil { return nil, errors.Wrapf(err, "getting directive tables for qual: %s", qual) } @@ -1194,7 +1085,13 @@ func (c *Controller) SetPoller(poller dax.AddressManager) { // This is useful in the case where MDS has restarted (or has been replaced) and // its poller is emtpy (i.e. it doesn't know about any nodes). func (c *Controller) InitializePoller(ctx context.Context) error { - nodes, err := c.nodeService.Nodes(context.Background()) + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + nodes, err := c.nodeService.Nodes(tx) if err != nil { return errors.Wrap(err, "initializing poller") } @@ -1219,10 +1116,21 @@ func (c *Controller) SnapshotTable(ctx context.Context, qtid dax.QualifiedTableI // snapshot that shard, then increment its shard version for logs written to the // WriteLogger. func (c *Controller) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) error { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return c.snapshotShardData(tx, qtid, shardNum) +} + +func (c *Controller) snapshotShardData(tx dax.Transaction, qtid dax.QualifiedTableID, shardNum dax.ShardNum) error { + qdbid := qtid.QualifiedDatabaseID + // Get the node responsible for the shard. - bal := c.ComputeBalancer - job := shard(qtid.Key(), shardNum) - workers, err := bal.WorkersForJobs(ctx, []dax.Job{dax.Job(job.String())}) + job := shard(qtid.Key(), shardNum).Job() + workers, err := c.Balancer.WorkersForJobs(tx, dax.RoleTypeCompute, qdbid, job) if err != nil { return errors.Wrapf(err, "getting workers for jobs: %s", job) } @@ -1231,7 +1139,7 @@ func (c *Controller) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTa return nil } - addr := dax.Address(workers[0].ID) + addr := dax.Address(workers[0].Address) // Send the node a snapshot request. req := &dax.SnapshotShardDataRequest{ @@ -1240,7 +1148,7 @@ func (c *Controller) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTa ShardNum: shardNum, } - if err := c.Director.SendSnapshotShardDataRequest(ctx, req); err != nil { + if err := c.Director.SendSnapshotShardDataRequest(tx.Context(), req); err != nil { return NewErrInternal(err.Error()) } @@ -1251,10 +1159,21 @@ func (c *Controller) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTa // partition to snapshot the table keys for that partition, then increment its // version for logs written to the WriteLogger. func (c *Controller) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) error { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return c.snapshotTableKeys(tx, qtid, partitionNum) +} + +func (c *Controller) snapshotTableKeys(tx dax.Transaction, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) error { + qdbid := qtid.QualifiedDatabaseID + // Get the node responsible for the partition. - bal := c.TranslateBalancer - job := partition(qtid.Key(), partitionNum) - workers, err := bal.WorkersForJobs(ctx, []dax.Job{dax.Job(job.String())}) + job := partition(qtid.Key(), partitionNum).Job() + workers, err := c.Balancer.WorkersForJobs(tx, dax.RoleTypeTranslate, qdbid, job) if err != nil { return errors.Wrapf(err, "getting workers for jobs: %s", job) } @@ -1263,7 +1182,7 @@ func (c *Controller) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTa return nil } - addr := dax.Address(workers[0].ID) + addr := dax.Address(workers[0].Address) // Send the node a snapshot request. req := &dax.SnapshotTableKeysRequest{ @@ -1272,7 +1191,7 @@ func (c *Controller) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTa PartitionNum: partitionNum, } - if err := c.Director.SendSnapshotTableKeysRequest(ctx, req); err != nil { + if err := c.Director.SendSnapshotTableKeysRequest(tx.Context(), req); err != nil { return NewErrInternal(err.Error()) } @@ -1283,13 +1202,24 @@ func (c *Controller) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTa // to snapshot the keys for that field, then increment its version for logs // written to the WriteLogger. func (c *Controller) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName) error { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return c.snapshotFieldKeys(tx, qtid, field) +} + +func (c *Controller) snapshotFieldKeys(tx dax.Transaction, qtid dax.QualifiedTableID, field dax.FieldName) error { + qdbid := qtid.QualifiedDatabaseID + // Get the node responsible for the field. - bal := c.TranslateBalancer // Field translation is currently handled by partition 0. partitionNum := dax.PartitionNum(0) - job := partition(qtid.Key(), partitionNum) + job := partition(qtid.Key(), partitionNum).Job() - workers, err := bal.WorkersForJobs(ctx, []dax.Job{dax.Job(job.String())}) + workers, err := c.Balancer.WorkersForJobs(tx, dax.RoleTypeTranslate, qdbid, job) if err != nil { return errors.Wrapf(err, "getting workers for jobs: %s", job) } @@ -1298,7 +1228,7 @@ func (c *Controller) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTa return nil } - addr := dax.Address(workers[0].ID) + addr := dax.Address(workers[0].Address) // Send the node a snapshot request. req := &dax.SnapshotFieldKeysRequest{ @@ -1307,7 +1237,7 @@ func (c *Controller) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTa Field: field, } - if err := c.Director.SendSnapshotFieldKeysRequest(ctx, req); err != nil { + if err := c.Director.SendSnapshotFieldKeysRequest(tx.Context(), req); err != nil { return NewErrInternal(err.Error()) } @@ -1316,17 +1246,48 @@ func (c *Controller) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTa ///////////// -func (c *Controller) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shards dax.ShardNums, isWrite bool) ([]dax.ComputeNode, error) { - inRole := &dax.ComputeRole{ +// ComputeNodes returns the compute nodes for the given table/shards. It always +// uses a read transaction. The writable equivalent to this method is +// `IngestShard`. +func (c *Controller) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shards dax.ShardNums) ([]dax.ComputeNode, error) { + role := &dax.ComputeRole{ TableKey: qtid.Key(), Shards: shards, } + qdbid := qtid.QualifiedDatabaseID - nodes, err := c.Nodes(ctx, inRole, isWrite) + tx, err := c.boltDB.BeginTx(ctx, false) if err != nil { - return nil, errors.Wrap(err, "getting compute nodes") + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + // If no shards are provided, get the nodes responsible for the entire + // table. + if len(role.Shards) == 0 { + assignedNodes, err := c.nodesForTable(tx, dax.RoleTypeCompute, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting nodes for table") + } + computeNodes, err := c.assignedToComputeNodes(assignedNodes) + if err != nil { + return nil, errors.Wrap(err, "converting assigned to compute nodes") + } + return computeNodes, nil } + assignedNodes, _, err := c.nodesComputeReadOrWrite(tx, role, qdbid, false, false) + if err != nil { + return nil, errors.Wrap(err, "getting compute nodes read or write") + } + + return c.assignedToComputeNodes(assignedNodes) +} + +// assignedToComputeNodes converts the provided []dax.AssignedNode to +// []dax.ComputeNode. If any of the assigned nodes are not for RoleType +// "compute", an error will be returned. +func (c *Controller) assignedToComputeNodes(nodes []dax.AssignedNode) ([]dax.ComputeNode, error) { computeNodes := make([]dax.ComputeNode, 0) for _, node := range nodes { @@ -1347,17 +1308,63 @@ func (c *Controller) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID return computeNodes, nil } -func (c *Controller) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitions dax.PartitionNums, isWrite bool) ([]dax.TranslateNode, error) { - inRole := &dax.TranslateRole{ +// TranslateNodes returns the translate nodes for the given table/partitions. It +// always uses a read transaction. The writable equivalent to this method is +// `IngestPartition`. +func (c *Controller) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitions dax.PartitionNums) ([]dax.TranslateNode, error) { + role := &dax.TranslateRole{ TableKey: qtid.Key(), Partitions: partitions, } + qdbid := qtid.QualifiedDatabaseID - nodes, err := c.Nodes(ctx, inRole, isWrite) + tx, err := c.boltDB.BeginTx(ctx, false) if err != nil { - return nil, errors.Wrap(err, "getting translate nodes") + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + // If no partitions are provided, get the nodes responsible for the entire + // table. + if len(role.Partitions) == 0 { + assignedNodes, err := c.nodesForTable(tx, dax.RoleTypeTranslate, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting nodes for table") + } + translateNodes, err := c.assignedToTranslateNodes(assignedNodes) + if err != nil { + return nil, errors.Wrap(err, "converting assigned to translate nodes") + } + return translateNodes, nil } + assignedNodes, _, err := c.nodesTranslateReadOrWrite(tx, role, qdbid, false, false) + if err != nil { + return nil, errors.Wrap(err, "getting translate nodes read or write") + } + + return c.assignedToTranslateNodes(assignedNodes) +} + +func (c *Controller) nodesForTable(tx dax.Transaction, roleType dax.RoleType, qtid dax.QualifiedTableID) ([]dax.AssignedNode, error) { + workers, err := c.Balancer.WorkersForTable(tx, roleType, qtid) + if err != nil { + return nil, errors.Wrapf(err, "getting workers for table: '%s'", qtid) + } + switch roleType { + case dax.RoleTypeCompute: + return c.computeWorkersToAssignedNodes(tx, workers) + case dax.RoleTypeTranslate: + return c.translateWorkersToAssignedNodes(tx, workers) + default: + return nil, errors.Errorf("unsupported role type: %s", roleType) + } +} + +// assignedToTranslateNodes converts the provided []dax.AssignedNode to +// []dax.TranslateNode. If any of the assigned nodes are not for RoleType +// "translate", an error will be returned. +func (c *Controller) assignedToTranslateNodes(nodes []dax.AssignedNode) ([]dax.TranslateNode, error) { translateNodes := make([]dax.TranslateNode, 0) for _, node := range nodes { @@ -1379,25 +1386,63 @@ func (c *Controller) TranslateNodes(ctx context.Context, qtid dax.QualifiedTable } func (c *Controller) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum) (dax.Address, error) { - c.mu.RLock() - defer c.mu.RUnlock() + role := &dax.TranslateRole{ + TableKey: qtid.Key(), + Partitions: dax.PartitionNums{partition}, + } + qdbid := qtid.QualifiedDatabaseID - partitions := dax.PartitionNums{partition} - - nodes, err := c.TranslateNodes(ctx, qtid, partitions, true) + // Try with a read transaction first. + tx, err := c.boltDB.BeginTx(ctx, false) if err != nil { + return "", errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if err := c.sanitizeQTID(tx, &qtid); err != nil { + return "", errors.Wrap(err, "sanitizing") + } + + // Verify that the table exists. + if _, err := c.Schemar.Table(tx, qtid); err != nil { return "", err } - if l := len(nodes); l == 0 { - return "", NewErrNoAvailableNode() - } else if l > 1 { - return "", - NewErrInternal( - fmt.Sprintf("unexpected number of nodes: %d", l)) + nodes, retryAsWrite, err := c.nodesTranslateReadOrWrite(tx, role, qdbid, true, false) + if err != nil { + return "", errors.Wrap(err, "getting translate nodes read or write") } - node := nodes[0] + // If it's writable, and we couldn't find all the partitions with just a + // read, try again with a write transaction. + if retryAsWrite { + tx.Rollback() + + tx, err = c.boltDB.BeginTx(ctx, true) + if err != nil { + return "", errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + nodes, _, err = c.nodesTranslateReadOrWrite(tx, role, qdbid, true, true) + if err != nil { + return "", errors.Wrap(err, "getting translate nodes read or write retry") + } + } + + translateNodes, err := c.assignedToTranslateNodes(nodes) + if err != nil { + return "", errors.Wrap(err, "converting assigned to translate nodes") + } + + if l := len(translateNodes); l == 0 { + return "", NewErrNoAvailableNode() + } else if l > 1 { + return "", NewErrInternal( + fmt.Sprintf("unexpected number of nodes: %d", l)) + } + + node := translateNodes[0] // Verify that the node returned is actually responsible for the partition // requested. @@ -1415,32 +1460,132 @@ func (c *Controller) IngestPartition(ctx context.Context, qtid dax.QualifiedTabl fmt.Sprintf("partition returned (%d) does not match requested (%d)", p, partition)) } + // Only commit if the transaction is writable. + if retryAsWrite { + return node.Address, tx.Commit() + } + + return node.Address, nil +} + +// IngestShard handles an ingest shard request. +func (c *Controller) IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shrdNum dax.ShardNum) (dax.Address, error) { + role := &dax.ComputeRole{ + TableKey: qtid.Key(), + Shards: dax.ShardNums{shrdNum}, + } + qdbid := qtid.QualifiedDatabaseID + + // Try with a read transaction first. + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return "", errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if err := c.sanitizeQTID(tx, &qtid); err != nil { + return "", errors.Wrap(err, "sanitizing") + } + + // Verify that the table exists. + if _, err := c.Schemar.Table(tx, qtid); err != nil { + return "", err + } + + nodes, retryAsWrite, err := c.nodesComputeReadOrWrite(tx, role, qdbid, true, false) + if err != nil { + return "", errors.Wrap(err, "getting compute nodes read or write") + } + + // If it's writable, and we couldn't find all the partitions with just a + // read, try again with a write transaction. + if retryAsWrite { + tx.Rollback() + + tx, err = c.boltDB.BeginTx(ctx, true) + if err != nil { + return "", errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + nodes, _, err = c.nodesComputeReadOrWrite(tx, role, qdbid, true, true) + if err != nil { + return "", errors.Wrap(err, "getting compute nodes read or write retry") + } + } + + computeNodes, err := c.assignedToComputeNodes(nodes) + if err != nil { + return "", errors.Wrap(err, "converting assigned to compute nodes") + } + + if l := len(computeNodes); l == 0 { + return "", NewErrNoAvailableNode() + } else if l > 1 { + return "", NewErrInternal( + fmt.Sprintf("unexpected number of nodes: %d", l)) + } + + node := computeNodes[0] + + // Verify that the node returned is actually responsible for the shard + // requested. + if node.Table != qtid.Key() { + return "", NewErrInternal( + fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid)) + } else if l := len(node.Shards); l != 1 { + return "", NewErrInternal( + fmt.Sprintf("unexpected number of shards returned: %d", l)) + } else if s := node.Shards[0]; s != shrdNum { + return "", NewErrInternal( + fmt.Sprintf("shard returned (%d) does not match requested (%d)", s, shrdNum)) + } + + // Only commit if the transaction is writable. + if retryAsWrite { + return node.Address, tx.Commit() + } + return node.Address, nil } //// func (c *Controller) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if err := c.sanitizeQTID(tx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + // Create the field in schemar. + if err := c.Schemar.CreateField(tx, qtid, fld); err != nil { + return errors.Wrapf(err, "creating field: %s, %s", qtid, fld) + } // workerSet maintains the set of workers which have a job assignment change // and therefore need to be sent an updated Directive. workerSet := NewAddressSet() + qdbid := qtid.QualifiedDatabaseID + // Get the worker(s) responsible for partition 0. - job := partition(qtid.Key(), 0).String() - workers, err := c.TranslateBalancer.WorkersForJobs(ctx, []dax.Job{dax.Job(job)}) + job := partition(qtid.Key(), 0).Job() + workers, err := c.Balancer.WorkersForJobs(tx, dax.RoleTypeTranslate, qdbid, job) if err != nil { return errors.Wrapf(err, "getting workers for job: %s", job) } for _, w := range workers { - workerSet.Add(dax.Address(w.ID)) + workerSet.Add(dax.Address(w.Address)) } // Get the list of workers responsible for shard data for this table. - if state, err := c.ComputeBalancer.CurrentState(ctx); err != nil { + if state, err := c.Balancer.CurrentState(tx, dax.RoleTypeCompute, qdbid); err != nil { return errors.Wrap(err, "getting current compute state") } else { for _, worker := range state { @@ -1448,7 +1593,7 @@ func (c *Controller) CreateField(ctx context.Context, qtid dax.QualifiedTableID, if shard, err := decodeShard(job); err != nil { return errors.Wrapf(err, "decoding shard: %s", job) } else if shard.table() == qtid.Key() { - workerSet.Add(dax.Address(worker.ID)) + workerSet.Add(dax.Address(worker.Address)) break } } @@ -1460,34 +1605,48 @@ func (c *Controller) CreateField(ctx context.Context, qtid dax.QualifiedTableID, addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) // Send a directive to any compute node responsible for this field. - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } - return nil + return tx.Commit() } func (c *Controller) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error { - c.mu.Lock() - defer c.mu.Unlock() + tx, err := c.boltDB.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + if err := c.sanitizeQTID(tx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + // Drop the field from schemar. + if err := c.Schemar.DropField(tx, qtid, fldName); err != nil { + return errors.Wrapf(err, "dropping field: %s, %s", qtid, fldName) + } // workerSet maintains the set of workers which have a job assignment change // and therefore need to be sent an updated Directive. workerSet := NewAddressSet() + qdbid := qtid.QualifiedDatabaseID + // Get the worker(s) responsible for partition 0. - job := partition(qtid.Key(), 0).String() - workers, err := c.TranslateBalancer.WorkersForJobs(ctx, []dax.Job{dax.Job(job)}) + job := partition(qtid.Key(), 0).Job() + workers, err := c.Balancer.WorkersForJobs(tx, dax.RoleTypeTranslate, qdbid, job) if err != nil { return errors.Wrapf(err, "getting workers for job: %s", job) } for _, w := range workers { - workerSet.Add(dax.Address(w.ID)) + workerSet.Add(dax.Address(w.Address)) } // Get the list of workers responsible for shard data for this table. - if state, err := c.ComputeBalancer.CurrentState(ctx); err != nil { + if state, err := c.Balancer.CurrentState(tx, dax.RoleTypeCompute, qdbid); err != nil { return errors.Wrap(err, "getting current compute state") } else { for _, worker := range state { @@ -1495,7 +1654,7 @@ func (c *Controller) DropField(ctx context.Context, qtid dax.QualifiedTableID, f if shard, err := decodeShard(job); err != nil { return errors.Wrapf(err, "decoding shard: %s", job) } else if shard.table() == qtid.Key() { - workerSet.Add(dax.Address(worker.ID)) + workerSet.Add(dax.Address(worker.Address)) break } } @@ -1507,11 +1666,11 @@ func (c *Controller) DropField(ctx context.Context, qtid dax.QualifiedTableID, f addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) // Send a directive to any compute node responsible for this field. - if err := c.sendDirectives(ctx, addressMethods...); err != nil { + if err := c.sendDirectives(tx, addressMethods...); err != nil { return NewErrDirectiveSendFailure(err.Error()) } - return nil + return tx.Commit() } ////////////////////////////////// @@ -1526,5 +1685,35 @@ func (c *Controller) RemoveAddresses(ctx context.Context, addrs ...dax.Address) } func (c *Controller) DebugNodes(ctx context.Context) ([]*dax.Node, error) { - return c.nodeService.Nodes(ctx) + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return c.nodeService.Nodes(tx) +} + +// 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 (c *Controller) sanitizeQTID(tx dax.Transaction, qtid *dax.QualifiedTableID) error { + if qtid.ID == "" { + nqtid, err := c.Schemar.TableID(tx, qtid.QualifiedDatabaseID, qtid.Name) + if err != nil { + return errors.Wrap(err, "getting table ID") + } + qtid.ID = nqtid.ID + } + return nil +} + +// TableID handles a table id (i.e. by name) request. +func (c *Controller) TableID(ctx context.Context, qdbid dax.QualifiedDatabaseID, name dax.TableName) (dax.QualifiedTableID, error) { + tx, err := c.boltDB.BeginTx(ctx, false) + if err != nil { + return dax.QualifiedTableID{}, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return c.Schemar.TableID(tx, qdbid, name) } diff --git a/dax/mds/controller/controller_test.go b/dax/mds/controller/controller_test.go index e57148b16..d135738f5 100644 --- a/dax/mds/controller/controller_test.go +++ b/dax/mds/controller/controller_test.go @@ -19,16 +19,25 @@ import ( func TestController(t *testing.T) { ctx := context.Background() - qual := dax.NewTableQualifier("acme", "db1") + qdbid := dax.NewQualifiedDatabaseID("acme", "db1") t.Run("RegisterNode", func(t *testing.T) { director := newTestDirector() - schemar, cleanup := daxtest.NewSchemar(t) defer cleanup() + + db := testbolt.MustOpenDB(t) + db.InitializeBuckets(balancerdb.BalancerBuckets...) + db.InitializeBuckets(schemardb.SchemarBuckets...) + defer func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + }() + cfg := controller.Config{ Director: director, Schemar: schemar, + BoltDB: db, } con := controller.New(cfg) @@ -61,19 +70,19 @@ func TestController(t *testing.T) { defer cleanup() db := testbolt.MustOpenDB(t) - db.InitializeBuckets(boltdb.NaiveBalancerBuckets...) + db.InitializeBuckets(balancerdb.BalancerBuckets...) + db.InitializeBuckets(schemardb.SchemarBuckets...) defer func() { testbolt.MustCloseDB(t, db) testbolt.CleanupDB(t, db.Path()) }() cfg := controller.Config{ - Director: director, - Schemar: schemar, - BoltDB: db, - StorageMethod: "boltdb", - ComputeBalancer: boltdb.NewBalancer("compute", db, logger.StderrLogger), - TranslateBalancer: boltdb.NewBalancer("translate", db, logger.StderrLogger), + Director: director, + Schemar: schemar, + BoltDB: db, + StorageMethod: "boltdb", + Balancer: balancerdb.NewBalancer(db, schemar, logger.StderrLogger), } con := controller.New(cfg) @@ -100,16 +109,28 @@ func TestController(t *testing.T) { } assert.Equal(t, exp, director.flush()) + // Add a qualified database. + dbOptions := dax.DatabaseOptions{ + WorkersMin: 1, + WorkersMax: 1, + } + qdb1 := daxtest.TestQualifiedDatabaseWithID(t, qdbid.OrganizationID, qdbid.DatabaseID, "dbname1", dbOptions) + assert.NoError(t, con.CreateDatabase(ctx, qdb1)) + + // tbls keeps the sorted list of tables used in tests + var tbls dax.QualifiedTables + // Add a non-keyed table. - tbl0 := daxtest.TestQualifiedTableWithID(t, qual, "2", "foo", 0, false) - assert.NoError(t, schemar.CreateTable(ctx, tbl0)) + tbl0 := daxtest.TestQualifiedTable(t, qdbid, "foo", 0, false) assert.NoError(t, con.CreateTable(ctx, tbl0)) + tbls = append(tbls, tbl0) + exp = []*dax.Directive{} assert.Equal(t, exp, director.flush()) // Add a shard. - assert.NoError(t, con.AddShards(ctx, tbl0.QualifiedID(), 0)) + addShards(t, ctx, con, tbl0.QualifiedID(), 0) exp = []*dax.Directive{ { @@ -130,6 +151,13 @@ func TestController(t *testing.T) { } assert.Equal(t, exp, director.flush()) + // Set WorkersMin to 3 so we can used the added nodes that follow. + { + dbOptions.WorkersMin = 3 + dbOptions.WorkersMax = 3 + assert.NoError(t, con.SetDatabaseOptions(ctx, qdb1.QualifiedID(), dbOptions)) + } + // Register two more nodes. node1 := &dax.Node{ Address: "10.0.0.1:81", @@ -172,11 +200,11 @@ func TestController(t *testing.T) { assert.Equal(t, exp, director.flush()) // Add more shards. - assert.NoError(t, con.AddShards(ctx, tbl0.QualifiedID(), dax.NewShardNums(1, 2, 3, 5, 8)...)) + addShards(t, ctx, con, tbl0.QualifiedID(), dax.NewShardNums(1, 2, 3, 5, 8)...) exp = []*dax.Directive{ { - Address: node0.Address, + Address: node1.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ tbl0, @@ -184,14 +212,14 @@ func TestController(t *testing.T) { ComputeRoles: []dax.ComputeRole{ { TableKey: tbl0.Key(), - Shards: dax.NewShardNums(0, 3), + Shards: dax.NewShardNums(1), }, }, TranslateRoles: []dax.TranslateRole{}, Version: 5, }, { - Address: node1.Address, + Address: node2.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ tbl0, @@ -199,99 +227,151 @@ func TestController(t *testing.T) { ComputeRoles: []dax.ComputeRole{ { TableKey: tbl0.Key(), - Shards: dax.NewShardNums(1, 5), + Shards: dax.NewShardNums(2), }, }, TranslateRoles: []dax.TranslateRole{}, Version: 6, }, - { - Address: node2.Address, - Method: dax.DirectiveMethodDiff, - Tables: []*dax.QualifiedTable{ - tbl0, - }, - ComputeRoles: []dax.ComputeRole{ - { - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(2, 8), - }, - }, - TranslateRoles: []dax.TranslateRole{}, - Version: 7, - }, - } - assert.Equal(t, exp, director.flush()) - - // Add another non-keyed table. - tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "bar", 0, false) - assert.NoError(t, schemar.CreateTable(ctx, tbl1)) - assert.NoError(t, con.CreateTable(ctx, tbl1)) - - // Add more shards. - assert.NoError(t, con.AddShards(ctx, tbl1.QualifiedID(), dax.NewShardNums(3, 5, 8, 13)...)) - - exp = []*dax.Directive{ { Address: node0.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, tbl0, }, ComputeRoles: []dax.ComputeRole{ - { - TableKey: tbl1.Key(), - Shards: dax.NewShardNums(3, 13), - }, { TableKey: tbl0.Key(), Shards: dax.NewShardNums(0, 3), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 8, + Version: 7, }, { Address: node1.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, tbl0, }, ComputeRoles: []dax.ComputeRole{ - { - TableKey: tbl1.Key(), - Shards: dax.NewShardNums(5), - }, { TableKey: tbl0.Key(), Shards: dax.NewShardNums(1, 5), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 9, + Version: 8, }, { Address: node2.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, tbl0, }, ComputeRoles: []dax.ComputeRole{ - { - TableKey: tbl1.Key(), - Shards: dax.NewShardNums(8), - }, { TableKey: tbl0.Key(), Shards: dax.NewShardNums(2, 8), }, }, TranslateRoles: []dax.TranslateRole{}, + Version: 9, + }, + } + assert.Equal(t, exp, director.flush()) + + // Add another non-keyed table. + tbl1 := daxtest.TestQualifiedTable(t, qdbid, "bar", 0, false) + assert.NoError(t, con.CreateTable(ctx, tbl1)) + + tbls = append(tbls, tbl1) + sort.Sort(tbls) + + // Add more shards. + addShards(t, ctx, con, tbl1.QualifiedID(), dax.NewShardNums(3, 5, 8, 13)...) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbls[0], + tbls[1], + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbls[0].Key(), + Shards: dax.NewShardNums(3), + }, + { + TableKey: tbls[1].Key(), + Shards: dax.NewShardNums(0, 3), + }, + }, + TranslateRoles: []dax.TranslateRole{}, Version: 10, }, + { + Address: node1.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbls[0], + tbls[1], + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbls[0].Key(), + Shards: dax.NewShardNums(5), + }, + { + TableKey: tbls[1].Key(), + Shards: dax.NewShardNums(1, 5), + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 11, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbls[0], + tbls[1], + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbls[0].Key(), + Shards: dax.NewShardNums(8), + }, + { + TableKey: tbls[1].Key(), + Shards: dax.NewShardNums(2, 8), + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 12, + }, + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbls[0], + tbls[1], + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbls[0].Key(), + Shards: dax.NewShardNums(3, 13), + }, + { + TableKey: tbls[1].Key(), + Shards: dax.NewShardNums(0, 3), + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 13, + }, } assert.Equal(t, exp, director.flush()) @@ -303,41 +383,41 @@ func TestController(t *testing.T) { Address: node0.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(3, 13), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Shards: dax.NewShardNums(0, 1, 3), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 11, + Version: 14, }, { Address: node2.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(5, 8), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Shards: dax.NewShardNums(2, 5, 8), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 12, + Version: 15, }, } assert.Equal(t, exp, director.flush()) @@ -350,21 +430,21 @@ func TestController(t *testing.T) { Address: node2.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(3, 5, 8, 13), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Shards: dax.NewShardNums(0, 1, 2, 3, 5, 8), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 13, + Version: 16, }, } assert.Equal(t, exp, director.flush()) @@ -375,6 +455,14 @@ func TestController(t *testing.T) { exp = []*dax.Directive{} assert.Equal(t, exp, director.flush()) + // Set WorkersMin to 1 so we can add a single node and have it be used + // (currently just adding 1 node won't satisfy the minimum of 3). + { + dbOptions.WorkersMin = 1 + dbOptions.WorkersMax = 1 + assert.NoError(t, con.SetDatabaseOptions(ctx, qdb1.QualifiedID(), dbOptions)) + } + // Add a new node and ensure that the free shards get assigned to it. node3 := &dax.Node{ Address: "10.0.0.1:83", @@ -389,21 +477,21 @@ func TestController(t *testing.T) { Address: node3.Address, Method: dax.DirectiveMethodReset, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(3, 5, 8, 13), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Shards: dax.NewShardNums(0, 1, 2, 3, 5, 8), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 14, + Version: 17, }, } assert.Equal(t, exp, director.flush()) @@ -416,21 +504,21 @@ func TestController(t *testing.T) { Address: node3.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(3, 5, 8, 13), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Shards: dax.NewShardNums(0, 1, 3, 8), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 15, + Version: 18, }, } assert.Equal(t, exp, director.flush()) @@ -445,21 +533,21 @@ func TestController(t *testing.T) { Address: node3.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(3, 5, 8, 13), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Shards: dax.NewShardNums(0, 1, 8), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 16, + Version: 19, }, } assert.Equal(t, exp, director.flush()) @@ -472,31 +560,27 @@ func TestController(t *testing.T) { Address: node3.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, + tbls[0], }, ComputeRoles: []dax.ComputeRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Shards: dax.NewShardNums(3, 5, 8, 13), }, }, TranslateRoles: []dax.TranslateRole{}, - Version: 17, + Version: 20, }, } assert.Equal(t, exp, director.flush()) // Remove a node which doesn't exist. - err := con.DeregisterNodes(ctx, "invalidNode") - if assert.Error(t, err) { - assert.True(t, errors.Is(err, dax.ErrNodeDoesNotExist)) - } - + assert.NoError(t, con.DeregisterNodes(ctx, "invalidNode")) }) t.Run("TranslateNodes", func(t *testing.T) { invalidQtid := dax.NewQualifiedTableID( - dax.NewTableQualifier("", ""), + dax.NewQualifiedDatabaseID("", ""), dax.TableID("invalidID"), ) @@ -505,19 +589,19 @@ func TestController(t *testing.T) { defer cleanup() db := testbolt.MustOpenDB(t) - db.InitializeBuckets(boltdb.NaiveBalancerBuckets...) + db.InitializeBuckets(balancerdb.BalancerBuckets...) + db.InitializeBuckets(schemardb.SchemarBuckets...) defer func() { testbolt.MustCloseDB(t, db) testbolt.CleanupDB(t, db.Path()) }() cfg := controller.Config{ - Director: director, - Schemar: schemar, - BoltDB: db, - StorageMethod: "boltdb", - ComputeBalancer: boltdb.NewBalancer("compute", db, logger.StderrLogger), - TranslateBalancer: boltdb.NewBalancer("translate", db, logger.StderrLogger), + Director: director, + Schemar: schemar, + BoltDB: db, + StorageMethod: "boltdb", + Balancer: balancerdb.NewBalancer(db, schemar, logger.StderrLogger), } con := controller.New(cfg) @@ -550,11 +634,23 @@ func TestController(t *testing.T) { exp = []*dax.Directive{} assert.Equal(t, exp, director.flush()) + // Add a qualified database. + dbOptions := dax.DatabaseOptions{ + WorkersMin: 1, + WorkersMax: 1, + } + qdb1 := daxtest.TestQualifiedDatabaseWithID(t, qdbid.OrganizationID, qdbid.DatabaseID, "dbname1", dbOptions) + assert.NoError(t, con.CreateDatabase(ctx, qdb1)) + + // tbls keeps the sorted list of tables used in tests + var tbls dax.QualifiedTables + // Add a keyed table. - tbl0 := daxtest.TestQualifiedTableWithID(t, qual, "2", "foo", 8, true) - assert.NoError(t, schemar.CreateTable(ctx, tbl0)) + tbl0 := daxtest.TestQualifiedTable(t, qdbid, "foo", 8, true) assert.NoError(t, con.CreateTable(ctx, tbl0)) + tbls = append(tbls, tbl0) + // Check directives. exp = []*dax.Directive{ { @@ -575,6 +671,13 @@ func TestController(t *testing.T) { } assert.Equal(t, exp, director.flush()) + // Set WorkersMin to 3 so we can used the two added nodes that follow. + { + dbOptions.WorkersMin = 3 + dbOptions.WorkersMax = 3 + assert.NoError(t, con.SetDatabaseOptions(ctx, qdb1.QualifiedID(), dbOptions)) + } + // Register two more nodes. node1 := &dax.Node{ Address: "10.0.0.1:81", @@ -586,34 +689,12 @@ func TestController(t *testing.T) { exp = []*dax.Directive{ { - Address: node0.Address, - Method: dax.DirectiveMethodDiff, - Tables: []*dax.QualifiedTable{ - tbl0, - }, - ComputeRoles: []dax.ComputeRole{}, - TranslateRoles: []dax.TranslateRole{ - { - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(0, 1, 2, 3), - }, - }, - Version: 3, - }, - { - Address: node1.Address, - Method: dax.DirectiveMethodReset, - Tables: []*dax.QualifiedTable{ - tbl0, - }, - ComputeRoles: []dax.ComputeRole{}, - TranslateRoles: []dax.TranslateRole{ - { - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(4, 5, 6, 7), - }, - }, - Version: 4, + Address: node1.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{}, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{}, + Version: 3, }, } assert.Equal(t, exp, director.flush()) @@ -640,7 +721,7 @@ func TestController(t *testing.T) { Partitions: dax.NewPartitionNums(0, 1, 2), }, }, - Version: 5, + Version: 4, }, { Address: node1.Address, @@ -652,10 +733,10 @@ func TestController(t *testing.T) { TranslateRoles: []dax.TranslateRole{ { TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(4, 5, 6), + Partitions: dax.NewPartitionNums(3, 5, 7), }, }, - Version: 6, + Version: 5, }, { Address: node2.Address, @@ -667,10 +748,10 @@ func TestController(t *testing.T) { TranslateRoles: []dax.TranslateRole{ { TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(3, 7), + Partitions: dax.NewPartitionNums(4, 6), }, }, - Version: 7, + Version: 6, }, } assert.Equal(t, exp, director.flush()) @@ -679,71 +760,73 @@ func TestController(t *testing.T) { // Make PartitionN double digit to ensure that partition ints aren't // sorted as strings. Also, it should be large enough to spill over // onto node0. - tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "bar", 24, true) - assert.NoError(t, schemar.CreateTable(ctx, tbl1)) + tbl1 := daxtest.TestQualifiedTable(t, qdbid, "bar", 24, true) assert.NoError(t, con.CreateTable(ctx, tbl1)) + tbls = append(tbls, tbl1) + sort.Sort(tbls) + // Check directives. exp = []*dax.Directive{ { Address: node0.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{}, TranslateRoles: []dax.TranslateRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Partitions: dax.NewPartitionNums(1, 4, 7, 10, 13, 16, 19, 22), }, { - TableKey: tbl0.Key(), + TableKey: tbls[1].Key(), Partitions: dax.NewPartitionNums(0, 1, 2), }, }, - Version: 8, + Version: 7, }, { Address: node1.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{}, TranslateRoles: []dax.TranslateRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Partitions: dax.NewPartitionNums(2, 5, 8, 11, 14, 17, 20, 23), }, { - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(4, 5, 6), + TableKey: tbls[1].Key(), + Partitions: dax.NewPartitionNums(3, 5, 7), }, }, - Version: 9, + Version: 8, }, { Address: node2.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, - tbl0, + tbls[0], + tbls[1], }, ComputeRoles: []dax.ComputeRole{}, TranslateRoles: []dax.TranslateRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Partitions: dax.NewPartitionNums(0, 3, 6, 9, 12, 15, 18, 21), }, { - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(3, 7), + TableKey: tbls[1].Key(), + Partitions: dax.NewPartitionNums(4, 6), }, }, - Version: 10, + Version: 9, }, } assert.Equal(t, exp, director.flush()) @@ -757,46 +840,46 @@ func TestController(t *testing.T) { Address: node0.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, + tbls[0], }, ComputeRoles: []dax.ComputeRole{}, TranslateRoles: []dax.TranslateRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Partitions: dax.NewPartitionNums(1, 4, 7, 10, 13, 16, 19, 22), }, }, - Version: 11, + Version: 10, }, { Address: node1.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, + tbls[0], }, ComputeRoles: []dax.ComputeRole{}, TranslateRoles: []dax.TranslateRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Partitions: dax.NewPartitionNums(2, 5, 8, 11, 14, 17, 20, 23), }, }, - Version: 12, + Version: 11, }, { Address: node2.Address, Method: dax.DirectiveMethodDiff, Tables: []*dax.QualifiedTable{ - tbl1, + tbls[0], }, ComputeRoles: []dax.ComputeRole{}, TranslateRoles: []dax.TranslateRole{ { - TableKey: tbl1.Key(), + TableKey: tbls[0].Key(), Partitions: dax.NewPartitionNums(0, 3, 6, 9, 12, 15, 18, 21), }, }, - Version: 13, + Version: 12, }, } assert.Equal(t, exp, director.flush()) @@ -822,18 +905,18 @@ func TestController(t *testing.T) { defer cleanup() db := testbolt.MustOpenDB(t) - db.InitializeBuckets(boltdb.NaiveBalancerBuckets...) + db.InitializeBuckets(balancerdb.BalancerBuckets...) + db.InitializeBuckets(schemardb.SchemarBuckets...) defer func() { testbolt.MustCloseDB(t, db) testbolt.CleanupDB(t, db.Path()) }() cfg := controller.Config{ - Schemar: schemar, - BoltDB: db, - StorageMethod: "boltdb", - ComputeBalancer: boltdb.NewBalancer("compute", db, logger.StderrLogger), - TranslateBalancer: boltdb.NewBalancer("translate", db, logger.StderrLogger), + Schemar: schemar, + BoltDB: db, + StorageMethod: "boltdb", + Balancer: balancerdb.NewBalancer(db, schemar, logger.StderrLogger), } con := controller.New(cfg) @@ -855,39 +938,42 @@ func TestController(t *testing.T) { } assert.NoError(t, con.RegisterNodes(ctx, node1)) + // Add a qualified database. + dbOptions := dax.DatabaseOptions{ + WorkersMin: 2, + WorkersMax: 2, + } + qdb1 := daxtest.TestQualifiedDatabaseWithID(t, qdbid.OrganizationID, qdbid.DatabaseID, "dbname1", dbOptions) + assert.NoError(t, con.CreateDatabase(ctx, qdb1)) + // Add a keyed table. - tbl0 := daxtest.TestQualifiedTable(t, qual, "foo", 12, true) - assert.NoError(t, schemar.CreateTable(ctx, tbl0)) + tbl0 := daxtest.TestQualifiedTable(t, qdbid, "foo", 12, true) assert.NoError(t, con.CreateTable(ctx, tbl0)) // Add shards. - assert.NoError(t, con.AddShards(ctx, tbl0.QualifiedID(), 0, 1, 2, 3, 11, 12)) + addShards(t, ctx, con, tbl0.QualifiedID(), 0, 1, 2, 3, 11, 12) - t.Run("ComputeRole", func(t *testing.T) { + t.Run("ComputeNodes", func(t *testing.T) { tests := []struct { - role dax.Role + role *dax.ComputeRole isWrite bool - exp []dax.AssignedNode + exp []dax.ComputeNode }{ { role: &dax.ComputeRole{ TableKey: tbl0.Key(), Shards: dax.NewShardNums(0, 1, 2, 3), }, - exp: []dax.AssignedNode{ + exp: []dax.ComputeNode{ { Address: node0.Address, - Role: &dax.ComputeRole{ - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(0, 2), - }, + Table: tbl0.Key(), + Shards: dax.NewShardNums(0, 2), }, { Address: node1.Address, - Role: &dax.ComputeRole{ - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(1, 3), - }, + Table: tbl0.Key(), + Shards: dax.NewShardNums(1, 3), }, }, }, @@ -896,13 +982,11 @@ func TestController(t *testing.T) { TableKey: tbl0.Key(), Shards: dax.NewShardNums(1), }, - exp: []dax.AssignedNode{ + exp: []dax.ComputeNode{ { Address: node1.Address, - Role: &dax.ComputeRole{ - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(1), - }, + Table: tbl0.Key(), + Shards: dax.NewShardNums(1), }, }, }, @@ -913,20 +997,16 @@ func TestController(t *testing.T) { Shards: dax.NewShardNums(1, 888, 889), }, isWrite: true, - exp: []dax.AssignedNode{ + exp: []dax.ComputeNode{ { Address: node0.Address, - Role: &dax.ComputeRole{ - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(888), - }, + Table: tbl0.Key(), + Shards: dax.NewShardNums(888), }, { Address: node1.Address, - Role: &dax.ComputeRole{ - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(1, 889), - }, + Table: tbl0.Key(), + Shards: dax.NewShardNums(1, 889), }, }, }, @@ -936,31 +1016,32 @@ func TestController(t *testing.T) { TableKey: tbl0.Key(), Shards: dax.NewShardNums(2, 11), }, - exp: []dax.AssignedNode{ + exp: []dax.ComputeNode{ { Address: node0.Address, - Role: &dax.ComputeRole{ - TableKey: tbl0.Key(), - Shards: dax.NewShardNums(2, 11), - }, + Table: tbl0.Key(), + Shards: dax.NewShardNums(2, 11), }, }, }, } for i, test := range tests { t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { - nodes, err := con.Nodes(ctx, test.role, test.isWrite) + if test.isWrite { + addShards(t, ctx, con, test.role.TableKey.QualifiedTableID(), test.role.Shards...) + } + nodes, err := con.ComputeNodes(ctx, test.role.TableKey.QualifiedTableID(), test.role.Shards) assert.NoError(t, err) assert.Equal(t, test.exp, nodes) }) } }) - t.Run("TranslateRole", func(t *testing.T) { + t.Run("TranslateNodes", func(t *testing.T) { tests := []struct { - role dax.Role + role *dax.TranslateRole isWrite bool - exp []dax.AssignedNode + exp []dax.TranslateNode expErrCode errors.Code }{ { @@ -969,13 +1050,11 @@ func TestController(t *testing.T) { Partitions: dax.NewPartitionNums(0), }, isWrite: true, - exp: []dax.AssignedNode{ + exp: []dax.TranslateNode{ { - Address: node0.Address, - Role: &dax.TranslateRole{ - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(0), - }, + Address: node0.Address, + Table: tbl0.Key(), + Partitions: dax.NewPartitionNums(0), }, }, }, @@ -985,22 +1064,19 @@ func TestController(t *testing.T) { Partitions: dax.NewPartitionNums(0, 1, 2, 3, 999), }, isWrite: false, - exp: []dax.AssignedNode{ + exp: []dax.TranslateNode{ { - Address: node0.Address, - Role: &dax.TranslateRole{ - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(0, 2), - }, + Address: node0.Address, + Table: tbl0.Key(), + Partitions: dax.NewPartitionNums(0, 2), }, { - Address: node1.Address, - Role: &dax.TranslateRole{ - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(1, 3), - }, + Address: node1.Address, + Table: tbl0.Key(), + Partitions: dax.NewPartitionNums(1, 3), }, }, + expErrCode: controller.ErrCodeUnassignedJobs, }, { role: &dax.TranslateRole{ @@ -1008,13 +1084,11 @@ func TestController(t *testing.T) { Partitions: dax.NewPartitionNums(1), }, isWrite: false, - exp: []dax.AssignedNode{ + exp: []dax.TranslateNode{ { - Address: node1.Address, - Role: &dax.TranslateRole{ - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(1), - }, + Address: node1.Address, + Table: tbl0.Key(), + Partitions: dax.NewPartitionNums(1), }, }, }, @@ -1025,20 +1099,21 @@ func TestController(t *testing.T) { Partitions: dax.NewPartitionNums(2, 10), }, isWrite: false, - exp: []dax.AssignedNode{ + exp: []dax.TranslateNode{ { - Address: node0.Address, - Role: &dax.TranslateRole{ - TableKey: tbl0.Key(), - Partitions: dax.NewPartitionNums(2, 10), - }, + Address: node0.Address, + Table: tbl0.Key(), + Partitions: dax.NewPartitionNums(2, 10), }, }, }, } for i, test := range tests { t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { - nodes, err := con.Nodes(ctx, test.role, test.isWrite) + if test.isWrite { + addPartitions(t, ctx, con, test.role.TableKey.QualifiedTableID(), test.role.Partitions...) + } + nodes, err := con.TranslateNodes(ctx, test.role.TableKey.QualifiedTableID(), test.role.Partitions) if test.expErrCode != "" { assert.True(t, errors.Is(err, test.expErrCode)) @@ -1101,3 +1176,21 @@ func (d *testDirector) flush() []*dax.Directive { return out } + +func addShards(t *testing.T, ctx context.Context, con *controller.Controller, qtid dax.QualifiedTableID, shards ...dax.ShardNum) { + t.Helper() + for _, shard := range shards { + if _, err := con.IngestShard(ctx, qtid, shard); err != nil { + assert.NoError(t, err) + } + } +} + +func addPartitions(t *testing.T, ctx context.Context, con *controller.Controller, qtid dax.QualifiedTableID, partitions ...dax.PartitionNum) { + t.Helper() + for _, parition := range partitions { + if _, err := con.IngestPartition(ctx, qtid, parition); err != nil { + assert.NoError(t, err) + } + } +} diff --git a/dax/mds/controller/errors.go b/dax/mds/controller/errors.go index 51b13cd59..f19bad93d 100644 --- a/dax/mds/controller/errors.go +++ b/dax/mds/controller/errors.go @@ -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), ) } diff --git a/dax/mds/controller/naive/balancer.go b/dax/mds/controller/naive/balancer.go deleted file mode 100644 index 8a8d2a996..000000000 --- a/dax/mds/controller/naive/balancer.go +++ /dev/null @@ -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 -} diff --git a/dax/mds/controller/naive/balancer_test.go b/dax/mds/controller/naive/balancer_test.go deleted file mode 100644 index 2830789e7..000000000 --- a/dax/mds/controller/naive/balancer_test.go +++ /dev/null @@ -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 -} diff --git a/dax/mds/controller/naive/boltdb/balancer.go b/dax/mds/controller/naive/boltdb/balancer.go deleted file mode 100644 index b8e202913..000000000 --- a/dax/mds/controller/naive/boltdb/balancer.go +++ /dev/null @@ -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) -} diff --git a/dax/mds/controller/naive/boltdb/balancer_test.go b/dax/mds/controller/naive/boltdb/balancer_test.go deleted file mode 100644 index b305f43c4..000000000 --- a/dax/mds/controller/naive/boltdb/balancer_test.go +++ /dev/null @@ -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 -} diff --git a/dax/mds/controller/sets.go b/dax/mds/controller/sets.go index deb42a801..1a7313ec6 100644 --- a/dax/mds/controller/sets.go +++ b/dax/mds/controller/sets.go @@ -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. diff --git a/dax/mds/controller/snapping_turtle.go b/dax/mds/controller/snapping_turtle.go index 0af9e2dd7..641555927 100644 --- a/dax/mds/controller/snapping_turtle.go +++ b/dax/mds/controller/snapping_turtle.go @@ -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) } diff --git a/dax/mds/controller/stringers.go b/dax/mds/controller/stringers.go index d47150470..6d56c8173 100644 --- a/dax/mds/controller/stringers.go +++ b/dax/mds/controller/stringers.go @@ -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 } diff --git a/dax/mds/http/handler.go b/dax/mds/http/handler.go index 32b46caf9..9926416d2 100644 --- a/dax/mds/http/handler.go +++ b/dax/mds/http/handler.go @@ -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 diff --git a/dax/mds/mds.go b/dax/mds/mds.go index 666fda0ce..ffa64e441 100644 --- a/dax/mds/mds.go +++ b/dax/mds/mds.go @@ -3,7 +3,6 @@ package mds import ( "context" - "fmt" "os" "sync" "time" @@ -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) } diff --git a/dax/mds/schemar/boltdb/schemar.go b/dax/mds/schemar/boltdb/schemar.go index 14f3c433e..47c9951fd 100644 --- a/dax/mds/schemar/boltdb/schemar.go +++ b/dax/mds/schemar/boltdb/schemar.go @@ -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() } diff --git a/dax/mds/schemar/boltdb/schemar_test.go b/dax/mds/schemar/boltdb/schemar_test.go index 6af10d7bd..20a1d4e27 100644 --- a/dax/mds/schemar/boltdb/schemar_test.go +++ b/dax/mds/schemar/boltdb/schemar_test.go @@ -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) - }) } diff --git a/dax/mds/schemar/errors.go b/dax/mds/schemar/errors.go index 7e1025ec6..710ed9b4c 100644 --- a/dax/mds/schemar/errors.go +++ b/dax/mds/schemar/errors.go @@ -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, diff --git a/dax/mds/schemar/http/handler.go b/dax/mds/schemar/http/handler.go deleted file mode 100644 index eb8dd3f1f..000000000 --- a/dax/mds/schemar/http/handler.go +++ /dev/null @@ -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"` -} diff --git a/dax/mds/schemar/schemar.go b/dax/mds/schemar/schemar.go index 7f365268c..1b45dce9f 100644 --- a/dax/mds/schemar/schemar.go +++ b/dax/mds/schemar/schemar.go @@ -8,23 +8,34 @@ import ( ) 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 +50,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 } diff --git a/dax/mds/schemar/schemar_test.go b/dax/mds/schemar/schemar_test.go deleted file mode 100644 index 8f0904af3..000000000 --- a/dax/mds/schemar/schemar_test.go +++ /dev/null @@ -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) - }) -} diff --git a/dax/node.go b/dax/node.go index bd130ee95..d1ddb3a03 100644 --- a/dax/node.go +++ b/dax/node.go @@ -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 diff --git a/dax/queryer/client/client.go b/dax/queryer/client/client.go index 6e438c4d8..eead12ca8 100644 --- a/dax/queryer/client/client.go +++ b/dax/queryer/client/client.go @@ -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, } diff --git a/dax/queryer/http/handler.go b/dax/queryer/http/handler.go index de344086b..419336bd3 100644 --- a/dax/queryer/http/handler.go +++ b/dax/queryer/http/handler.go @@ -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 diff --git a/dax/queryer/orchestrator.go b/dax/queryer/orchestrator.go index 15b28e16e..5d6ce48c9 100644 --- a/dax/queryer/orchestrator.go +++ b/dax/queryer/orchestrator.go @@ -1840,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{} { @@ -3511,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, } } @@ -3528,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: diff --git a/dax/queryer/queryer.go b/dax/queryer/queryer.go index 1356e8e8d..8b052d2db 100644 --- a/dax/queryer/queryer.go +++ b/dax/queryer/queryer.go @@ -32,7 +32,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 +47,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 +59,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 +77,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, @@ -92,8 +92,8 @@ func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator { logger: q.logger, } - qorch := newQualifiedOrchestrator(orch, qual) - q.orchestrators[qual] = qorch + qorch := newQualifiedOrchestrator(orch, qdbid) + q.orchestrators[qdbid] = qorch return qorch } @@ -132,12 +132,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() { @@ -149,6 +146,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 { @@ -165,19 +175,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 { @@ -229,7 +237,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 { @@ -240,7 +248,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 @@ -250,9 +258,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 } @@ -261,11 +269,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 { @@ -276,14 +310,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") } diff --git a/dax/queryer/schema_api.go b/dax/queryer/schema_api.go index 19405f832..52c638ec0 100644 --- a/dax/queryer/schema_api.go +++ b/dax/queryer/schema_api.go @@ -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) diff --git a/dax/role.go b/dax/role.go index 92712ca64..165356611 100644 --- a/dax/role.go +++ b/dax/role.go @@ -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 diff --git a/dax/schema.go b/dax/schema.go index f7b16abc4..05d31c43e 100644 --- a/dax/schema.go +++ b/dax/schema.go @@ -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 { diff --git a/dax/storage/storage_test.go b/dax/storage/storage_test.go index 0b368569a..081133580 100644 --- a/dax/storage/storage_test.go +++ b/dax/storage/storage_test.go @@ -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", } diff --git a/dax/table.go b/dax/table.go index eb39df22a..577174872 100644 --- a/dax/table.go +++ b/dax/table.go @@ -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, } } diff --git a/dax/table_test.go b/dax/table_test.go index 091722cdd..5706ccf9d 100644 --- a/dax/table_test.go +++ b/dax/table_test.go @@ -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) diff --git a/dax/test/dax/dax_test.go b/dax/test/dax/dax_test.go index 87b510a3b..823f8c1ac 100644 --- a/dax/test/dax/dax_test.go +++ b/dax/test/dax/dax_test.go @@ -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{ @@ -160,7 +179,7 @@ func TestDAXIntegration(t *testing.T) { runTableTests(t, svcmgr.Queryer.Address(), - basicTableTestConfig(qual, tableTests...)..., + basicTableTestConfig(qdbid, tableTests...)..., ) }) @@ -168,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)..., ) }) @@ -182,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") @@ -189,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 @@ -238,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)..., ) }) @@ -268,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, @@ -284,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. @@ -291,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) @@ -313,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, @@ -341,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, @@ -357,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) @@ -365,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 @@ -410,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)..., ) }) @@ -439,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, @@ -452,7 +512,7 @@ func TestDAXIntegration(t *testing.T) { /////////////////////////////////////////////////// type tableTestConfig struct { - qual dax.TableQualifier + qdbid dax.QualifiedDatabaseID test defs.TableTest skipCreate bool skipInsert bool @@ -461,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{ @@ -491,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) } } @@ -499,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) } } @@ -525,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 @@ -602,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 @@ -680,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 diff --git a/dax/test/table.go b/dax/test/table.go index 8d2ba93e8..2bde25a40 100644 --- a/dax/test/table.go +++ b/dax/test/table.go @@ -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, ) } diff --git a/dax/transaction.go b/dax/transaction.go new file mode 100644 index 000000000..0e8e0f598 --- /dev/null +++ b/dax/transaction.go @@ -0,0 +1,9 @@ +package dax + +import "context" + +type Transaction interface { + Commit() error + Context() context.Context + Rollback() error +} diff --git a/dax/workerjob.go b/dax/workerjob.go index a249bd82b..2e0c7b12b 100644 --- a/dax/workerjob.go +++ b/dax/workerjob.go @@ -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) { diff --git a/idk/ingest.go b/idk/ingest.go index 47d33280d..f3cf02485 100644 --- a/idk/ingest.go +++ b/idk/ingest.go @@ -1007,10 +1007,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 { @@ -1019,7 +1019,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) diff --git a/idk/ingest_test.go b/idk/ingest_test.go index 7a25000f2..64aab66e2 100644 --- a/idk/ingest_test.go +++ b/idk/ingest_test.go @@ -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) } } @@ -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) } diff --git a/idk/mds/importer.go b/idk/mds/importer.go index f70c96c34..0d7a1a31c 100644 --- a/idk/mds/importer.go +++ b/idk/mds/importer.go @@ -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, } } @@ -261,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 { diff --git a/idk/mds/schemamanager.go b/idk/mds/schemamanager.go index b73fe4537..d1e691c06 100644 --- a/idk/mds/schemamanager.go +++ b/idk/mds/schemamanager.go @@ -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 } diff --git a/sql3/planner/compilecreateview.go b/sql3/planner/compilecreateview.go index d3a554ee7..b056a8e52 100644 --- a/sql3/planner/compilecreateview.go +++ b/sql3/planner/compilecreateview.go @@ -3,8 +3,8 @@ package planner import ( - "github.com/molecula/featurebase/v3/sql3/parser" - "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/featurebasedb/featurebase/v3/sql3/parser" + "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) // compileCreateViewStatement compiles a parser.CreateViewStatement AST into a PlanOperator diff --git a/sql3/planner/compiledropview.go b/sql3/planner/compiledropview.go index 819c84fac..fec75a82b 100644 --- a/sql3/planner/compiledropview.go +++ b/sql3/planner/compiledropview.go @@ -3,9 +3,9 @@ package planner import ( - "github.com/molecula/featurebase/v3/sql3" - "github.com/molecula/featurebase/v3/sql3/parser" - "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/featurebasedb/featurebase/v3/sql3" + "github.com/featurebasedb/featurebase/v3/sql3/parser" + "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) // compileDropViewStatement compiles a DROP VIEW statement into a PlanOperator. diff --git a/sql3/planner/opalterview.go b/sql3/planner/opalterview.go index 70789acb1..5bf3f9832 100644 --- a/sql3/planner/opalterview.go +++ b/sql3/planner/opalterview.go @@ -6,8 +6,8 @@ import ( "context" "fmt" - "github.com/molecula/featurebase/v3/sql3" - "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/featurebasedb/featurebase/v3/sql3" + "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) // PlanOpAlterView implements the ALTER VIEW operator diff --git a/sql3/planner/opcreateview.go b/sql3/planner/opcreateview.go index 310a6bdfa..5d24656e9 100644 --- a/sql3/planner/opcreateview.go +++ b/sql3/planner/opcreateview.go @@ -6,9 +6,9 @@ import ( "context" "fmt" - "github.com/molecula/featurebase/v3/dax" - "github.com/molecula/featurebase/v3/sql3" - "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/featurebasedb/featurebase/v3/dax" + "github.com/featurebasedb/featurebase/v3/sql3" + "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) // PlanOpCreateView implements the CREATE VIEW operator diff --git a/sql3/planner/opdropview.go b/sql3/planner/opdropview.go index c90b50469..341034f8d 100644 --- a/sql3/planner/opdropview.go +++ b/sql3/planner/opdropview.go @@ -6,8 +6,8 @@ import ( "context" "fmt" - "github.com/molecula/featurebase/v3/sql3" - "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/featurebasedb/featurebase/v3/sql3" + "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) // PlanOpDropView plan operator to drop a view. diff --git a/sql3/planner/opfeaturebasetables.go b/sql3/planner/opfeaturebasetables.go index 59b03375d..2eae4c2cc 100644 --- a/sql3/planner/opfeaturebasetables.go +++ b/sql3/planner/opfeaturebasetables.go @@ -135,11 +135,17 @@ func (i *showTablesRowIter) Next(ctx context.Context) (types.Row, error) { } default: u := i.planner.systemAPI.DataDir() - u = fmt.Sprintf("%s/indexes/%s", u, indexName) - spaceUsed, err = pilosa.GetDiskUsage(u) - if err != nil { - return nil, err + // TODO(tlt): GetDiskUsage needs to be behind an interface because + // this doesn't work in serverless. For now I'm just going to skip + // it based on the emtpy DataDir, but let's do this the right way. + if u != "" { + u = fmt.Sprintf("%s/indexes/%s", u, indexName) + + spaceUsed, err = pilosa.GetDiskUsage(u) + if err != nil { + return nil, err + } } } diff --git a/sql3/planner/systemobjects.go b/sql3/planner/systemobjects.go index d806d7359..909a6e1e0 100644 --- a/sql3/planner/systemobjects.go +++ b/sql3/planner/systemobjects.go @@ -6,11 +6,11 @@ import ( "context" "time" - pilosa "github.com/molecula/featurebase/v3" - "github.com/molecula/featurebase/v3/dax" - "github.com/molecula/featurebase/v3/sql3" - "github.com/molecula/featurebase/v3/sql3/parser" - "github.com/molecula/featurebase/v3/sql3/planner/types" + pilosa "github.com/featurebasedb/featurebase/v3" + "github.com/featurebasedb/featurebase/v3/dax" + "github.com/featurebasedb/featurebase/v3/sql3" + "github.com/featurebasedb/featurebase/v3/sql3/parser" + "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) type viewSystemObject struct {