Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve Multi perf on SQL state stores with 1 op only #3300

Merged
merged 4 commits into from
Jan 4, 2024
Merged

Improve Multi perf on SQL state stores with 1 op only #3300

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0798594
Improve Multi perf on SQL state stores with 1 op only
ItalyPaleAle Jan 2, 2024
c22622b
Update common/component/postgresql/v1/postgresql.go
ItalyPaleAle Jan 3, 2024
7cf1953
Update state/postgresql/v2/postgresql.go
ItalyPaleAle Jan 3, 2024
f33d5ee
Merge branch 'main' into multi-single-op
ItalyPaleAle Jan 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
47 changes: 29 additions & 18 deletions common/component/postgresql/v1/postgresql.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -438,29 +438,40 @@ func (p *PostgreSQL) doDelete(parentCtx context.Context, db pginterfaces.DBQueri
}

func (p *PostgreSQL) Multi(parentCtx context.Context, request *state.TransactionalStateRequest) error {
_, err := pgtransactions.ExecuteInTransaction[struct{}](parentCtx, p.logger, p.db, p.metadata.Timeout, func(ctx context.Context, tx pgx.Tx) (res struct{}, err error) {
for _, o := range request.Operations {
switch x := o.(type) {
case state.SetRequest:
err = p.doSet(parentCtx, tx, &x)
if request == nil {
return nil
}

// If there's only 1 operation, skip starting a transaction
switch len(request.Operations) {
case 0:
return nil
case 1:
return p.execMultiOperation(parentCtx, request.Operations[0], p.db)
default:
_, err := pgtransactions.ExecuteInTransaction[struct{}](parentCtx, p.logger, p.db, p.metadata.Timeout, func(ctx context.Context, tx pgx.Tx) (res struct{}, err error) {
for _, op := range request.Operations {
err = p.execMultiOperation(ctx, op, tx)
if err != nil {
return res, err
}

case state.DeleteRequest:
err = p.doDelete(parentCtx, tx, &x)
if err != nil {
return res, err
}

default:
return res, fmt.Errorf("unsupported operation: %s", o.Operation())
}
}

return res, nil
})
return err
return res, nil
})
return err
}
}

func (p *PostgreSQL) execMultiOperation(ctx context.Context, op state.TransactionalStateOperation, db pginterfaces.DBQuerier) error {
switch x := op.(type) {
case state.SetRequest:
return p.doSet(ctx, db, &x)
case state.DeleteRequest:
return p.doDelete(ctx, db, &x)
default:
return fmt.Errorf("unsupported operation: %s", op.Operation())
}
}

func (p *PostgreSQL) CleanupExpired() error {
Expand Down
107 changes: 72 additions & 35 deletions common/component/postgresql/v1/postgresql_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,16 +44,9 @@ func TestMultiWithNoRequests(t *testing.T) {
m, _ := mockDatabase(t)
defer m.db.Close()

m.db.ExpectBegin()
m.db.ExpectCommit()
// There's also a rollback called after a commit, which is expected and will not have effect
m.db.ExpectRollback()

var operations []state.TransactionalStateOperation

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
Operations: nil,
})

// Assert
Expand All @@ -66,24 +59,46 @@ func TestValidSetRequest(t *testing.T) {
defer m.db.Close()

setReq := createSetRequest()
operations := []state.TransactionalStateOperation{setReq}
val, _ := json.Marshal(setReq.Value)

m.db.ExpectBegin()
m.db.ExpectExec("INSERT INTO").
WithArgs(setReq.Key, string(val), false).
WillReturnResult(pgxmock.NewResult("INSERT", 1))
m.db.ExpectCommit()
// There's also a rollback called after a commit, which is expected and will not have effect
m.db.ExpectRollback()
t.Run("single op", func(t *testing.T) {
operations := []state.TransactionalStateOperation{setReq}

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
m.db.ExpectExec("INSERT INTO").
WithArgs(setReq.Key, string(val), false).
WillReturnResult(pgxmock.NewResult("INSERT", 1))

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
})

// Assert
require.NoError(t, err)
})

// Assert
require.NoError(t, err)
t.Run("multiple ops", func(t *testing.T) {
operations := []state.TransactionalStateOperation{setReq, setReq}

m.db.ExpectBegin()
m.db.ExpectExec("INSERT INTO").
WithArgs(setReq.Key, string(val), false).
WillReturnResult(pgxmock.NewResult("INSERT", 1))
m.db.ExpectExec("INSERT INTO").
WithArgs(setReq.Key, string(val), false).
WillReturnResult(pgxmock.NewResult("INSERT", 1))
m.db.ExpectCommit()
// There's also a rollback called after a commit, which is expected and will not have effect
m.db.ExpectRollback()

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
})

// Assert
require.NoError(t, err)
})
}

func TestInvalidMultiSetRequestNoKey(t *testing.T) {
Expand Down Expand Up @@ -113,23 +128,45 @@ func TestValidMultiDeleteRequest(t *testing.T) {
defer m.db.Close()

deleteReq := createDeleteRequest()
operations := []state.TransactionalStateOperation{deleteReq}

m.db.ExpectBegin()
m.db.ExpectExec("DELETE FROM").
WithArgs(deleteReq.Key).
WillReturnResult(pgxmock.NewResult("DELETE", 1))
m.db.ExpectCommit()
// There's also a rollback called after a commit, which is expected and will not have effect
m.db.ExpectRollback()
t.Run("single op", func(t *testing.T) {
operations := []state.TransactionalStateOperation{deleteReq}

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
m.db.ExpectExec("DELETE FROM").
WithArgs(deleteReq.Key).
WillReturnResult(pgxmock.NewResult("DELETE", 1))

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
})

// Assert
require.NoError(t, err)
})

// Assert
require.NoError(t, err)
t.Run("multiple ops", func(t *testing.T) {
operations := []state.TransactionalStateOperation{deleteReq, deleteReq}

m.db.ExpectBegin()
m.db.ExpectExec("DELETE FROM").
WithArgs(deleteReq.Key).
WillReturnResult(pgxmock.NewResult("DELETE", 1))
m.db.ExpectExec("DELETE FROM").
WithArgs(deleteReq.Key).
WillReturnResult(pgxmock.NewResult("DELETE", 1))
m.db.ExpectCommit()
// There's also a rollback called after a commit, which is expected and will not have effect
m.db.ExpectRollback()

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Operations: operations,
})

// Assert
require.NoError(t, err)
})
}

func TestInvalidMultiDeleteRequestNoKey(t *testing.T) {
Expand All @@ -140,7 +177,7 @@ func TestInvalidMultiDeleteRequestNoKey(t *testing.T) {
m.db.ExpectBegin()
m.db.ExpectRollback()

operations := []state.TransactionalStateOperation{state.DeleteRequest{}} // Delete request without key is not valid for Delete operation
operations := []state.TransactionalStateOperation{state.DeleteRequest{}, state.DeleteRequest{}} // Delete request without key is not valid for Delete operation

// Act
err := m.pg.Multi(context.Background(), &state.TransactionalStateRequest{
Expand Down
56 changes: 33 additions & 23 deletions state/mysql/mysql.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -768,40 +768,50 @@ func readRow(row interface{ Scan(dest ...any) error }) (key string, value []byte
return key, value, &etag, expireTime, nil
}

// Multi handles multiple transactions.
// TransactionalStore Interface.
// Multi handles multiple operations in batch.
// Implements the TransactionalStore Interface.
func (m *MySQL) Multi(ctx context.Context, request *state.TransactionalStateRequest) error {
tx, err := m.db.Begin()
if err != nil {
return err
if request == nil {
return nil
}
defer func() {
rollbackErr := tx.Rollback()
if rollbackErr != nil && !errors.Is(rollbackErr, sql.ErrTxDone) {
m.logger.Errorf("Error rolling back transaction: %v", rollbackErr)
}
}()

for _, o := range request.Operations {
switch req := o.(type) {
case state.SetRequest:
err = m.setValue(ctx, tx, &req)
if err != nil {
return err
// If there's only 1 operation, skip starting a transaction
switch len(request.Operations) {
case 0:
return nil
case 1:
return m.execMultiOperation(ctx, request.Operations[0], m.db)
default:
tx, err := m.db.Begin()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should probably use BeginTx here like we do in sqlite.

Suggested change
tx, err := m.db.Begin()
tx, err := m.db.BeginTx(ctx, nil)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly a good idea, but I'd like to not make this change in this PR and maybe do a follow-up PR

Begin and BeginTx have some difference in behavior, specifically (unlike in the pgx driver for Postgres) the context controls the lifecycle of the entire transaction. See: https://pkg.go.dev/database/sql#DB.BeginTx

if err != nil {
return err
}
defer func() {
rollbackErr := tx.Rollback()
if rollbackErr != nil && !errors.Is(rollbackErr, sql.ErrTxDone) {
m.logger.Errorf("Error rolling back transaction: %v", rollbackErr)
}
}()

case state.DeleteRequest:
err = m.deleteValue(ctx, tx, &req)
for _, op := range request.Operations {
err = m.execMultiOperation(ctx, op, tx)
if err != nil {
return err
}

default:
return fmt.Errorf("unsupported operation: %s", req.Operation())
}
return tx.Commit()
}
}

return tx.Commit()
func (m *MySQL) execMultiOperation(ctx context.Context, op state.TransactionalStateOperation, db querier) error {
switch req := op.(type) {
case state.SetRequest:
return m.setValue(ctx, db, &req)
case state.DeleteRequest:
return m.deleteValue(ctx, db, &req)
default:
return fmt.Errorf("unsupported operation: %s", op.Operation())
}
}

// Close implements io.Closer.
Expand Down