Skip to content

Commit

Permalink
fix: improve the performance of list artifacts
Browse files Browse the repository at this point in the history 
1. Change the query for listing tasks of scan which can use the db
   index.
2. Add the gin index for task.extra_attrs.report_uuids

Fixes: #18013

Signed-off-by: chlins <[email protected]>
  • Loading branch information
@chlins
chlins committed Apr 27, 2023
1 parent bf6389e commit f12c209
Show file tree
Hide file tree
Showing 10 changed files with 175 additions and 31 deletions.
1 change: 1 addition & 0 deletions make/migrations/postgresql/0120_2.9.0_schema.up.sql
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
CREATE INDEX IF NOT EXISTS idx_task_extra_attrs_report_uuids ON task USING gin ((extra_attrs::jsonb->'report_uuids'));
18 changes: 7 additions & 11 deletions src/controller/scan/base_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -964,15 +964,6 @@ func (bc *basicController) launchScanJob(ctx context.Context, param *launchScanJ
reportUUIDsKey: reportUUIDs,
}

// NOTE: due to the limitation of the beego's orm, the List method of the task manager not support ?! operator for the jsonb field,
// we cann't list the tasks for scan reports of uuid1, uuid2 by SQL `SELECT * FROM task WHERE (extra_attrs->'report_uuids')::jsonb ?| array['uuid1', 'uuid2']`
// or by `SELECT * FROM task WHERE id IN (SELECT id FROM task WHERE (extra_attrs->'report_uuids')::jsonb ?| array['uuid1', 'uuid2'])`
// so save {"report:uuid1": "1", "report:uuid2": "2"} in the extra_attrs of the task, and then list it with
// SQL `SELECT * FROM task WHERE extra_attrs->>'report:uuid1' = '1'` in loop
for _, reportUUID := range reportUUIDs {
extraAttrs["report:"+reportUUID] = "1"
}

_, err = bc.taskMgr.Create(ctx, param.ExecutionID, j, extraAttrs)
return err
}
Expand Down Expand Up @@ -1022,11 +1013,16 @@ func (bc *basicController) listScanTasks(ctx context.Context, reportUUIDs []stri
}

func (bc *basicController) getScanTask(ctx context.Context, reportUUID string) (*task.Task, error) {
query := q.New(q.KeyWords{"extra_attrs." + "report:" + reportUUID: "1"})
tasks, err := bc.taskMgr.List(bc.cloneCtx(ctx), query)
// NOTE: the SQL uses the postgres' unique grammar and should consider here if support other database in the future.
// Due to the limitation of the beego's orm, the SQL cannot be converted by orm framework,
// so we can only execute the query by raw SQL, the SQL filters the task contains the report uuid in the column extra_attrs,
// consider from performance side which can using indexes to speed up queries.
sql := fmt.Sprintf(`SELECT * FROM task WHERE extra_attrs::jsonb->'report_uuids' @> '["%s"]'`, reportUUID)
tasks, err := bc.taskMgr.ListByRaw(ctx, sql)
if err != nil {
return nil, err
}

if len(tasks) == 0 {
return nil, errors.NotFoundError(nil).WithMessage("task for report %s not found", reportUUID)
}
Expand Down
42 changes: 22 additions & 20 deletions src/controller/scan/base_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -321,7 +321,7 @@ func (suite *ControllerTestSuite) TestScanControllerScan() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListByRaw").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"), Status: "Success"},
}, nil).Once()

Expand All @@ -343,7 +343,7 @@ func (suite *ControllerTestSuite) TestScanControllerScan() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListByRaw").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"), Status: "Success"},
}, nil).Once()

Expand All @@ -360,7 +360,7 @@ func (suite *ControllerTestSuite) TestScanControllerScan() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListByRaw").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"), Status: "Running"},
}, nil).Once()

Expand Down Expand Up @@ -409,37 +409,40 @@ func (suite *ControllerTestSuite) TestScanControllerStop() {

// TestScanControllerGetReport ...
func (suite *ControllerTestSuite) TestScanControllerGetReport() {
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
mock.OnAnything(suite.ar, "Walk").Return(nil).Run(func(args mock.Arguments) {
walkFn := args.Get(2).(func(*artifact.Artifact) error)
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListByRaw").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001")},
}, nil).Once()
mock.OnAnything(suite.accessoryMgr, "List").Return(nil, nil)
rep, err := suite.c.GetReport(context.TODO(), suite.artifact, []string{v1.MimeTypeNativeReport})
rep, err := suite.c.GetReport(ctx, suite.artifact, []string{v1.MimeTypeNativeReport})
require.NoError(suite.T(), err)
assert.Equal(suite.T(), 1, len(rep))
}

// TestScanControllerGetSummary ...
func (suite *ControllerTestSuite) TestScanControllerGetSummary() {
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
mock.OnAnything(suite.accessoryMgr, "List").Return([]accessoryModel.Accessory{}, nil).Once()
mock.OnAnything(suite.ar, "Walk").Return(nil).Run(func(args mock.Arguments) {
walkFn := args.Get(2).(func(*artifact.Artifact) error)
walkFn(suite.artifact)
}).Once()
mock.OnAnything(suite.taskMgr, "List").Return(nil, nil).Once()
mock.OnAnything(suite.taskMgr, "ListByRaw").Return(nil, nil).Once()

sum, err := suite.c.GetSummary(context.TODO(), suite.artifact, []string{v1.MimeTypeNativeReport})
sum, err := suite.c.GetSummary(ctx, suite.artifact, []string{v1.MimeTypeNativeReport})
require.NoError(suite.T(), err)
assert.Equal(suite.T(), 1, len(sum))
}

// TestScanControllerGetScanLog ...
func (suite *ControllerTestSuite) TestScanControllerGetScanLog() {
mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
mock.OnAnything(suite.taskMgr, "ListByRaw").Return([]*task.Task{
{
ID: 1,
ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"),
Expand All @@ -448,7 +451,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetScanLog() {

mock.OnAnything(suite.taskMgr, "GetLog").Return([]byte("log"), nil).Once()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, "rp-uuid-001")
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, "rp-uuid-001")
require.NoError(suite.T(), err)
assert.Condition(suite.T(), func() (success bool) {
success = len(bytes) > 0
Expand All @@ -457,8 +460,8 @@ func (suite *ControllerTestSuite) TestScanControllerGetScanLog() {
}

func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
kw1 := q.KeyWords{"extra_attrs.report:rp-uuid-001": "1"}
suite.taskMgr.On("List", context.TODO(), q.New(kw1)).Return([]*task.Task{
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
suite.taskMgr.On("ListByRaw", ctx, `SELECT * FROM task WHERE extra_attrs::jsonb->'report_uuids' @> '["rp-uuid-001"]'`).Return([]*task.Task{
{
ID: 1,
ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"),
Expand All @@ -469,8 +472,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
walkFn(suite.artifact)
})
mock.OnAnything(suite.accessoryMgr, "List").Return(nil, nil)
kw2 := q.KeyWords{"extra_attrs.report:rp-uuid-002": "1"}
suite.taskMgr.On("List", context.TODO(), q.New(kw2)).Return([]*task.Task{
suite.taskMgr.On("ListByRaw", ctx, `SELECT * FROM task WHERE extra_attrs::jsonb->'report_uuids' @> '["rp-uuid-002"]'`).Return([]*task.Task{
{
ID: 2,
ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-002"),
Expand All @@ -480,7 +482,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
// Both success
mock.OnAnything(suite.taskMgr, "GetLog").Return([]byte("log"), nil).Twice()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Nil(err)
suite.NotEmpty(bytes)
suite.Contains(string(bytes), "Logs of report rp-uuid-001")
Expand All @@ -489,10 +491,10 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {

{
// One successfully, one failed
suite.taskMgr.On("GetLog", context.TODO(), int64(1)).Return([]byte("log"), nil).Once()
suite.taskMgr.On("GetLog", context.TODO(), int64(2)).Return(nil, fmt.Errorf("failed")).Once()
suite.taskMgr.On("GetLog", ctx, int64(1)).Return([]byte("log"), nil).Once()
suite.taskMgr.On("GetLog", ctx, int64(2)).Return(nil, fmt.Errorf("failed")).Once()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Nil(err)
suite.NotEmpty(bytes)
suite.NotContains(string(bytes), "Logs of report rp-uuid-001")
Expand All @@ -502,7 +504,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
// Both failed
mock.OnAnything(suite.taskMgr, "GetLog").Return(nil, fmt.Errorf("failed")).Twice()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Error(err)
suite.Empty(bytes)
}
Expand All @@ -511,7 +513,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
// Both empty
mock.OnAnything(suite.taskMgr, "GetLog").Return(nil, nil).Twice()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Nil(err)
suite.Empty(bytes)
}
Expand Down Expand Up @@ -560,7 +562,7 @@ func (suite *ControllerTestSuite) TestScanAll() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return(nil, nil).Once()
mock.OnAnything(suite.taskMgr, "ListByRaw").Return(nil, nil).Once()

mock.OnAnything(suite.reportMgr, "Delete").Return(nil).Once()
mock.OnAnything(suite.reportMgr, "Create").Return("uuid", nil).Once()
Expand Down
17 changes: 17 additions & 0 deletions src/pkg/task/dao/task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ type TaskDAO interface {
// List the tasks according to the query
// Query the "ExtraAttrs" by setting 'query.Keywords["ExtraAttrs.key"]="value"'
List(ctx context.Context, query *q.Query) (tasks []*Task, err error)
// ListByRaw lists the tasks by raw sql
ListByRaw(ctx context.Context, sql string, args ...interface{}) (tasks []*Task, err error)
// Get the specified task
Get(ctx context.Context, id int64) (task *Task, err error)
// Create a task
Expand Down Expand Up @@ -88,6 +90,21 @@ func (t *taskDAO) List(ctx context.Context, query *q.Query) ([]*Task, error) {
return tasks, nil
}

func (t *taskDAO) ListByRaw(ctx context.Context, sql string, args ...interface{}) ([]*Task, error) {
ormer, err := orm.FromContext(ctx)
if err != nil {
return nil, err
}

tasks := []*Task{}
_, err = ormer.Raw(sql, args...).QueryRows(&tasks)
if err != nil {
return nil, err
}

return tasks, nil
}

func (t *taskDAO) Get(ctx context.Context, id int64) (*Task, error) {
task := &Task{
ID: id,
Expand Down
12 changes: 12 additions & 0 deletions src/pkg/task/dao/task_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,18 @@ func (t *taskDAOTestSuite) TestList() {
t.Require().Len(tasks, 0)
}

func (t *taskDAOTestSuite) TestListByRaw() {
// exist case
tasks, err := t.taskDAO.ListByRaw(t.ctx, "select * from task where id = ?", t.taskID)
t.Require().Nil(err)
t.Require().Len(tasks, 1)
t.Equal(t.taskID, tasks[0].ID)
// not exist case
tasks, err = t.taskDAO.ListByRaw(t.ctx, "select * from task where id = ?", -1)
t.Require().Nil(err)
t.Require().Len(tasks, 0)
}

func (t *taskDAOTestSuite) TestGet() {
// not exist
_, err := t.taskDAO.Get(t.ctx, 10000)
Expand Down
29 changes: 29 additions & 0 deletions src/pkg/task/mock_task_dao_test.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

29 changes: 29 additions & 0 deletions src/pkg/task/mock_task_manager_test.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

16 changes: 16 additions & 0 deletions src/pkg/task/task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ type Manager interface {
// List the tasks according to the query
// Query the "ExtraAttrs" by setting 'query.Keywords["ExtraAttrs.key"]="value"'
List(ctx context.Context, query *q.Query) (tasks []*Task, err error)
// ListByRaw lists the tasks by raw sql
ListByRaw(ctx context.Context, sql string, args ...interface{}) (tasks []*Task, err error)
// Update the extra attributes of the specified task
UpdateExtraAttrs(ctx context.Context, id int64, extraAttrs map[string]interface{}) (err error)
// Get the log of the specified task
Expand Down Expand Up @@ -234,6 +236,20 @@ func (m *manager) List(ctx context.Context, query *q.Query) ([]*Task, error) {
return ts, nil
}

func (m *manager) ListByRaw(ctx context.Context, sql string, args ...interface{}) ([]*Task, error) {
tasks, err := m.dao.ListByRaw(ctx, sql, args...)
if err != nil {
return nil, err
}
var ts []*Task
for _, task := range tasks {
t := &Task{}
t.From(task)
ts = append(ts, t)
}
return ts, nil
}

func (m *manager) UpdateExtraAttrs(ctx context.Context, id int64, extraAttrs map[string]interface{}) error {
data, err := json.Marshal(extraAttrs)
if err != nil {
Expand Down
13 changes: 13 additions & 0 deletions src/pkg/task/task_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,19 @@ func (t *taskManagerTestSuite) TestList() {
t.dao.AssertExpectations(t.T())
}

func (t *taskManagerTestSuite) TestListByRaw() {
t.dao.On("ListByRaw", mock.Anything, mock.Anything).Return([]*dao.Task{
{
ID: 1,
},
}, nil)
tasks, err := t.mgr.ListByRaw(nil, "select 1")
t.Require().Nil(err)
t.Require().Len(tasks, 1)
t.Equal(int64(1), tasks[0].ID)
t.dao.AssertExpectations(t.T())
}

func TestTaskManagerTestSuite(t *testing.T) {
suite.Run(t, &taskManagerTestSuite{})
}
29 changes: 29 additions & 0 deletions src/testing/pkg/task/manager.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit f12c209

Please sign in to comment.