drsybren/Flamenco
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Worker: separate contexts of upstream buffer

Browse Source 
The upstream buffer takes care of two things: communication with Flamenco
Manager (first context) and buffering things in a SQLite database (second
context). This commit separates those two contexts, so that shutting down
the application isn't going to prevent buffering things in the database.
This commit is contained in:
Sybren A. Stüvel 2022-04-08 16:49:53 +02:00
parent 959a235c59
commit cc20b850ec
2 changed files with 47 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
internal/worker/upstream_buffer.go
View File

@ -21,6 +21,13 @@ import (
// TODO: pull the SQLite stuff out of this file into a more global place, so // TODO: pull the SQLite stuff out of this file into a more global place, so
// that other areas of Flamenco Worker can also use it. // that other areas of Flamenco Worker can also use it.
// Note that there are two contexts used in this file. One (`dbCtx`) is for
// database access and is a local, short-lived, background context. The other
// (`ctx`) is the one that's passed from the caller, which should indicate the
// global worker context. If that context is done, queueing updates in the
// database will still work, but all communication with Flamenco Manager will
// halt.
// UpstreamBufferDB implements the UpstreamBuffer interface using a database as backend. // UpstreamBufferDB implements the UpstreamBuffer interface using a database as backend.
type UpstreamBufferDB struct { type UpstreamBufferDB struct {
db *sql.DB db *sql.DB
@ -35,6 +42,7 @@ type UpstreamBufferDB struct {
} }
const defaultUpstreamFlushInterval = 30 * time.Second const defaultUpstreamFlushInterval = 30 * time.Second
const databaseContextTimeout = 10 * time.Second
var _ UpstreamBuffer = (*UpstreamBufferDB)(nil) var _ UpstreamBuffer = (*UpstreamBufferDB)(nil)
@ -54,7 +62,7 @@ func NewUpstreamBuffer(client FlamencoClient, clock TimeService) (*UpstreamBuffe
} }
// OpenDB opens the database. Must be called once before using. // OpenDB opens the database. Must be called once before using.
func (ub *UpstreamBufferDB) OpenDB(ctx context.Context, databaseFilename string) error { func (ub *UpstreamBufferDB) OpenDB(dbCtx context.Context, databaseFilename string) error {
if ub.db != nil { if ub.db != nil {
return errors.New("upstream buffer database already opened") return errors.New("upstream buffer database already opened")
} }
@ -64,13 +72,13 @@ func (ub *UpstreamBufferDB) OpenDB(ctx context.Context, databaseFilename string)
return fmt.Errorf("opening %s: %w", databaseFilename, err) return fmt.Errorf("opening %s: %w", databaseFilename, err)
} }
if err := db.PingContext(ctx); err != nil { if err := db.PingContext(dbCtx); err != nil {
return fmt.Errorf("accessing %s: %w", databaseFilename, err) return fmt.Errorf("accessing %s: %w", databaseFilename, err)
} }
ub.db = db ub.db = db
if err := ub.prepareDatabase(ctx); err != nil { if err := ub.prepareDatabase(dbCtx); err != nil {
return err return err
} }
@ -84,7 +92,7 @@ func (ub *UpstreamBufferDB) SendTaskUpdate(ctx context.Context, taskID string, u
ub.dbMutex.Lock() ub.dbMutex.Lock()
defer ub.dbMutex.Unlock() defer ub.dbMutex.Unlock()
queueSize, err := ub.queueSize(ctx) queueSize, err := ub.queueSize()
if err != nil { if err != nil {
return fmt.Errorf("unable to determine upstream queue size: %w", err) return fmt.Errorf("unable to determine upstream queue size: %w", err)
} }
@ -93,7 +101,7 @@ func (ub *UpstreamBufferDB) SendTaskUpdate(ctx context.Context, taskID string, u
if queueSize > 0 { if queueSize > 0 {
log.Debug().Int("queueSize", queueSize). log.Debug().Int("queueSize", queueSize).
Msg("task updates already queued, immediately queueing new update") Msg("task updates already queued, immediately queueing new update")
return ub.queueTaskUpdate(ctx, taskID, update) return ub.queueTaskUpdate(taskID, update)
} }
// Try to deliver the update. // Try to deliver the update.
@ -101,7 +109,7 @@ func (ub *UpstreamBufferDB) SendTaskUpdate(ctx context.Context, taskID string, u
if err != nil { if err != nil {
log.Warn().Err(err).Str("task", taskID). log.Warn().Err(err).Str("task", taskID).
Msg("error communicating with Manager, going to queue task update for sending later") Msg("error communicating with Manager, going to queue task update for sending later")
return ub.queueTaskUpdate(ctx, taskID, update) return ub.queueTaskUpdate(taskID, update)
} }
// The Manager responded, so no need to queue this update, even when there was an error. // The Manager responded, so no need to queue this update, even when there was an error.
@ -117,7 +125,7 @@ func (ub *UpstreamBufferDB) SendTaskUpdate(ctx context.Context, taskID string, u
} }
// Close releases the database. It does not try to flush any pending items. // Close releases the database. It does not try to flush any pending items.
func (ub *UpstreamBufferDB) Close(ctx context.Context) error { func (ub *UpstreamBufferDB) Close() error {
if ub.db == nil { if ub.db == nil {
return nil return nil
} }
@ -131,11 +139,11 @@ func (ub *UpstreamBufferDB) Close(ctx context.Context) error {
} }
// prepareDatabase creates the database schema, if necessary. // prepareDatabase creates the database schema, if necessary.
func (ub *UpstreamBufferDB) prepareDatabase(ctx context.Context) error { func (ub *UpstreamBufferDB) prepareDatabase(dbCtx context.Context) error {
ub.dbMutex.Lock() ub.dbMutex.Lock()
defer ub.dbMutex.Unlock() defer ub.dbMutex.Unlock()
tx, err := ub.db.BeginTx(ctx, nil) tx, err := ub.db.BeginTx(dbCtx, nil)
if err != nil { if err != nil {
return fmt.Errorf("beginning database transaction: %w", err) return fmt.Errorf("beginning database transaction: %w", err)
} }
@ -144,7 +152,7 @@ func (ub *UpstreamBufferDB) prepareDatabase(ctx context.Context) error {
stmt := `CREATE TABLE IF NOT EXISTS task_update_queue(task_id VARCHAR(36), payload BLOB)` stmt := `CREATE TABLE IF NOT EXISTS task_update_queue(task_id VARCHAR(36), payload BLOB)`
log.Debug().Str("sql", stmt).Msg("creating database table") log.Debug().Str("sql", stmt).Msg("creating database table")
if _, err := tx.ExecContext(ctx, stmt); err != nil { if _, err := tx.ExecContext(dbCtx, stmt); err != nil {
return fmt.Errorf("creating database table: %w", err) return fmt.Errorf("creating database table: %w", err)
} }
@ -155,15 +163,18 @@ func (ub *UpstreamBufferDB) prepareDatabase(ctx context.Context) error {
return nil return nil
} }
func (ub *UpstreamBufferDB) queueSize(ctx context.Context) (int, error) { func (ub *UpstreamBufferDB) queueSize() (int, error) {
if ub.db == nil { if ub.db == nil {
log.Panic().Msg("no database opened, unable to inspect upstream queue") log.Panic().Msg("no database opened, unable to inspect upstream queue")
} }
dbCtx, dbCtxCancel := context.WithTimeout(context.Background(), databaseContextTimeout)
defer dbCtxCancel()
var queueSize int var queueSize int
err := ub.db. err := ub.db.
QueryRowContext(ctx, "SELECT count(*) FROM task_update_queue"). QueryRowContext(dbCtx, "SELECT count(*) FROM task_update_queue").
Scan(&queueSize) Scan(&queueSize)
switch { switch {
@ -176,12 +187,15 @@ func (ub *UpstreamBufferDB) queueSize(ctx context.Context) (int, error) {
} }
} }
func (ub *UpstreamBufferDB) queueTaskUpdate(ctx context.Context, taskID string, update api.TaskUpdateJSONRequestBody) error { func (ub *UpstreamBufferDB) queueTaskUpdate(taskID string, update api.TaskUpdateJSONRequestBody) error {
if ub.db == nil { if ub.db == nil {
log.Panic().Msg("no database opened, unable to queue task updates") log.Panic().Msg("no database opened, unable to queue task updates")
} }
tx, err := ub.db.BeginTx(ctx, nil) dbCtx, dbCtxCancel := context.WithTimeout(context.Background(), databaseContextTimeout)
defer dbCtxCancel()
tx, err := ub.db.BeginTx(dbCtx, nil)
if err != nil { if err != nil {
return fmt.Errorf("beginning database transaction: %w", err) return fmt.Errorf("beginning database transaction: %w", err)
} }
@ -195,7 +209,7 @@ func (ub *UpstreamBufferDB) queueTaskUpdate(ctx context.Context, taskID string,
stmt := `INSERT INTO task_update_queue (task_id, payload) VALUES (?, ?)` stmt := `INSERT INTO task_update_queue (task_id, payload) VALUES (?, ?)`
log.Debug().Str("sql", stmt).Str("task", taskID).Msg("inserting task update") log.Debug().Str("sql", stmt).Str("task", taskID).Msg("inserting task update")
if _, err := tx.ExecContext(ctx, stmt, taskID, blob); err != nil { if _, err := tx.ExecContext(dbCtx, stmt, taskID, blob); err != nil {
return fmt.Errorf("queueing task update: %w", err) return fmt.Errorf("queueing task update: %w", err)
} }
@ -215,7 +229,7 @@ func (ub *UpstreamBufferDB) Flush(ctx context.Context) error {
} }
// See if we need to flush at all. // See if we need to flush at all.
queueSize, err := ub.queueSize(ctx) queueSize, err := ub.queueSize()
switch { switch {
case err != nil: case err != nil:
return fmt.Errorf("unable to determine queue size: %w", err) return fmt.Errorf("unable to determine queue size: %w", err)
@ -237,7 +251,10 @@ func (ub *UpstreamBufferDB) Flush(ctx context.Context) error {
} }
func (ub *UpstreamBufferDB) flushFirstItem(ctx context.Context) (done bool, err error) { func (ub *UpstreamBufferDB) flushFirstItem(ctx context.Context) (done bool, err error) {
tx, err := ub.db.BeginTx(ctx, nil) dbCtx, dbCtxCancel := context.WithTimeout(context.Background(), databaseContextTimeout)
defer dbCtxCancel()
tx, err := ub.db.BeginTx(dbCtx, nil)
if err != nil { if err != nil {
return false, fmt.Errorf("beginning database transaction: %w", err) return false, fmt.Errorf("beginning database transaction: %w", err)
} }
@ -250,7 +267,7 @@ func (ub *UpstreamBufferDB) flushFirstItem(ctx context.Context) (done bool, err
var taskID string var taskID string
var blob []byte var blob []byte
err = tx.QueryRowContext(ctx, stmt).Scan(&rowID, &taskID, &blob) err = tx.QueryRowContext(dbCtx, stmt).Scan(&rowID, &taskID, &blob)
switch { switch {
case err == sql.ErrNoRows: case err == sql.ErrNoRows:
// Flush operation is done. // Flush operation is done.
@ -268,7 +285,7 @@ func (ub *UpstreamBufferDB) flushFirstItem(ctx context.Context) (done bool, err
// than to discard it and ignore it ever happened. // than to discard it and ignore it ever happened.
logger.Warn().Err(err). logger.Warn().Err(err).
Msg("unable to unmarshal queued task update, discarding") Msg("unable to unmarshal queued task update, discarding")
if err := ub.discardRow(ctx, tx, rowID); err != nil { if err := ub.discardRow(tx, rowID); err != nil {
return false, err return false, err
} }
return false, tx.Commit() return false, tx.Commit()
@ -295,17 +312,20 @@ func (ub *UpstreamBufferDB) flushFirstItem(ctx context.Context) (done bool, err
Msg("queued task update discarded by Manager, unknown reason") Msg("queued task update discarded by Manager, unknown reason")
} }
if err := ub.discardRow(ctx, tx, rowID); err != nil { if err := ub.discardRow(tx, rowID); err != nil {
return false, err return false, err
} }
return false, tx.Commit() return false, tx.Commit()
} }
func (ub *UpstreamBufferDB) discardRow(ctx context.Context, tx *sql.Tx, rowID int64) error { func (ub *UpstreamBufferDB) discardRow(tx *sql.Tx, rowID int64) error {
dbCtx, dbCtxCancel := context.WithTimeout(context.Background(), databaseContextTimeout)
defer dbCtxCancel()
stmt := `DELETE FROM task_update_queue WHERE rowid = ?` stmt := `DELETE FROM task_update_queue WHERE rowid = ?`
log.Trace().Str("sql", stmt).Int64("rowID", rowID).Msg("un-queueing task update") log.Trace().Str("sql", stmt).Int64("rowID", rowID).Msg("un-queueing task update")
_, err := tx.ExecContext(ctx, stmt, rowID) _, err := tx.ExecContext(dbCtx, stmt, rowID)
if err != nil { if err != nil {
return fmt.Errorf("un-queueing task update: %w", err) return fmt.Errorf("un-queueing task update: %w", err)
} }
@ -334,7 +354,7 @@ func (ub *UpstreamBufferDB) periodicFlushLoop() {
} }
func rollbackTransaction(tx *sql.Tx) { func rollbackTransaction(tx *sql.Tx) {
if err := tx.Rollback(); err != nil { if err := tx.Rollback(); err != nil && err != sql.ErrTxDone {
log.Error().Err(err).Msg("rolling back transaction") log.Error().Err(err).Msg("rolling back transaction")
} }
} }

8
internal/worker/upstream_buffer_test.go
View File

@ -46,7 +46,7 @@ func TestUpstreamBufferCloseUnopened(t *testing.T) {
defer mockCtrl.Finish() defer mockCtrl.Finish()
ub, _ := mockUpstreamBufferDB(t, mockCtrl) ub, _ := mockUpstreamBufferDB(t, mockCtrl)
err := ub.Close(context.Background()) err := ub.Close()
assert.NoError(t, err, "Closing without opening should be OK") assert.NoError(t, err, "Closing without opening should be OK")
} }
@ -76,7 +76,7 @@ func TestUpstreamBufferManagerUnavailable(t *testing.T) {
assert.NoError(t, err) assert.NoError(t, err)
// Check the queue size, it should have an item queued. // Check the queue size, it should have an item queued.
queueSize, err := ub.queueSize(ctx) queueSize, err := ub.queueSize()
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, 1, queueSize) assert.Equal(t, 1, queueSize)
@ -94,10 +94,10 @@ func TestUpstreamBufferManagerUnavailable(t *testing.T) {
// Queue should be empty now. // Queue should be empty now.
ub.dbMutex.Lock() ub.dbMutex.Lock()
queueSize, err = ub.queueSize(ctx) queueSize, err = ub.queueSize()
ub.dbMutex.Unlock() ub.dbMutex.Unlock()
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, 0, queueSize) assert.Equal(t, 0, queueSize)
assert.NoError(t, ub.Close(ctx)) assert.NoError(t, ub.Close())
} }