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 committed Apr 29, 2023
1 parent c09e539 commit 7884e3c
Show file tree
Hide file tree
Showing 10 changed files with 177 additions and 31 deletions.
1 change: 1 addition & 0 deletions make/migrations/postgresql/0120_2.9.0_schema.up.sql
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'));
14 changes: 3 additions & 11 deletions src/controller/scan/base_controller.go
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,12 @@ 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 method uses the postgres' unique operations and should consider here if support other database in the future.
tasks, err := bc.taskMgr.ListScanTasksByReportUUID(ctx, reportUUID)
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
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, "ListScanTasksByReportUUID").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, "ListScanTasksByReportUUID").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, "ListScanTasksByReportUUID").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, "ListScanTasksByReportUUID").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, "ListScanTasksByReportUUID").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, "ListScanTasksByReportUUID").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("ListScanTasksByReportUUID", ctx, "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("ListScanTasksByReportUUID", ctx, "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, "ListScanTasksByReportUUID").Return(nil, nil).Once()

mock.OnAnything(suite.reportMgr, "Delete").Return(nil).Once()
mock.OnAnything(suite.reportMgr, "Create").Return("uuid", nil).Once()
Expand Down
22 changes: 22 additions & 0 deletions src/pkg/task/dao/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,9 @@ type TaskDAO interface {
UpdateStatusInBatch(ctx context.Context, jobIDs []string, status string, batchSize int) (err error)
// ExecutionIDsByVendorAndStatus retrieve the execution id by vendor status
ExecutionIDsByVendorAndStatus(ctx context.Context, vendorType, status string) ([]int64, error)
// ListScanTasksByReportUUID lists scan tasks by report uuid, although it's a specific case but it will be
// more suitable to support multi database in the future.
ListScanTasksByReportUUID(ctx context.Context, uuid string) (tasks []*Task, err error)
}

// NewTaskDAO returns an instance of TaskDAO
Expand Down Expand Up @@ -88,6 +91,25 @@ func (t *taskDAO) List(ctx context.Context, query *q.Query) ([]*Task, error) {
return tasks, nil
}

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

tasks := []*Task{}
// 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"]'`, uuid)
_, err = ormer.Raw(sql).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
21 changes: 21 additions & 0 deletions src/pkg/task/dao/task_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,27 @@ func (t *taskDAOTestSuite) TestList() {
t.Require().Len(tasks, 0)
}

func (t *taskDAOTestSuite) TestListScanTasksByReportUUID() {
// should not exist if non set
tasks, err := t.taskDAO.ListScanTasksByReportUUID(t.ctx, "fake-report-uuid")
t.Require().Nil(err)
t.Require().Len(tasks, 0)
// create one with report uuid
taskID, err := t.taskDAO.Create(t.ctx, &Task{
ExecutionID: t.executionID,
Status: "success",
StatusCode: 1,
ExtraAttrs: `{"report_uuids": ["fake-report-uuid"]}`,
})
t.Require().Nil(err)
defer t.taskDAO.Delete(t.ctx, taskID)
// should exist as created
tasks, err = t.taskDAO.ListScanTasksByReportUUID(t.ctx, "fake-report-uuid")
t.Require().Nil(err)
t.Require().Len(tasks, 1)
t.Equal(taskID, tasks[0].ID)
}

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

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

26 changes: 26 additions & 0 deletions src/pkg/task/mock_task_manager_test.go

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

17 changes: 17 additions & 0 deletions src/pkg/task/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ type Manager interface {
UpdateStatusInBatch(ctx context.Context, jobIDs []string, status string, batchSize int) error
// ExecutionIDsByVendorAndStatus retrieve execution id by vendor type and status
ExecutionIDsByVendorAndStatus(ctx context.Context, vendorType, status string) ([]int64, error)
// ListScanTasksByReportUUID lists scan tasks by report uuid, although it's a specific case but it will be
// more suitable to support multi database in the future.
ListScanTasksByReportUUID(ctx context.Context, uuid string) (tasks []*Task, err error)
}

// NewManager creates an instance of the default task manager
Expand Down Expand Up @@ -234,6 +237,20 @@ func (m *manager) List(ctx context.Context, query *q.Query) ([]*Task, error) {
return ts, nil
}

func (m *manager) ListScanTasksByReportUUID(ctx context.Context, uuid string) ([]*Task, error) {
tasks, err := m.dao.ListScanTasksByReportUUID(ctx, uuid)
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
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) TestListScanTasksByReportUUID() {
t.dao.On("ListScanTasksByReportUUID", mock.Anything, mock.Anything).Return([]*dao.Task{
{
ID: 1,
},
}, nil)
tasks, err := t.mgr.ListScanTasksByReportUUID(nil, "uuid")
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{})
}
26 changes: 26 additions & 0 deletions src/testing/pkg/task/manager.go

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

0 comments on commit 7884e3c

Please sign in to comment.