featurebase/dax/controller/balancer/balancer_test.go
Matthew Jaffee be4f365eaf
Cloud 1358 bolt postgres (#2286)
Switch Serverless from using BoltDB to Postgres as metadata store.

Previously, the controller stored all metadata to BoltDB. This implements SQLDB (currently Postgres flavored) as the backing store for metadata. This will allow us to have multiple instances of the controller running for HA, and to easily inspect and repair the contents of the metadata store.

Unfortunately, it was not straightforward to keep the BoltDB implementation working alongside the SQL one, so it will be removed in a later patch. Once that's done, the SQL implementation should allow for a number of simplifications of the schemar and balancer interfaces.

Database migration is built directly into the application by embedding the migration files and logic from the `soda` command line tool. When connecting to the RDBMS, the app will always attempt to create the necessary database and apply any outstanding migrations.

Integration tests truncate all tables upon start, but *not* at the end, so the state of the database can be inspected after integration tests.

Had to refactor some of the controller's background tasks to make sure they get properly shut down on controller exit.
2023-03-20 09:02:22 -05:00

1097 lines
24 KiB
Go

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/controller"
"github.com/featurebasedb/featurebase/v3/dax/controller/balancer/boltdb"
schemardb "github.com/featurebasedb/featurebase/v3/dax/controller/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"
"github.com/stretchr/testify/require"
)
func newBoltBalancer(t *testing.T) (*daxbolt.DB, func()) {
db := testbolt.MustOpenDB(t)
require.NoError(t, db.InitializeBuckets(boltdb.BalancerBuckets...))
require.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) {
assert.NoError(t, schemar.SetDatabaseOption(tx, qdb.QualifiedID(), dax.DatabaseOptionWorkersMin, "3"))
})
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) {
assert.NoError(t, schemar.SetDatabaseOption(tx, qdb.QualifiedID(), dax.DatabaseOptionWorkersMin, "2"))
})
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) {
assert.NoError(t, schemar.SetDatabaseOption(tx, qdb.QualifiedID(), dax.DatabaseOptionWorkersMin, "3"))
})
// 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) {
assert.NoError(t, schemar.SetDatabaseOption(tx, qdb.QualifiedID(), dax.DatabaseOptionWorkersMin, "3"))
})
// 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())
})
}