Repository: oom-ai/oomstore (go)

Issue:
bug:  offline export has the possibility of memory leaks
https://github.com/oom-ai/oomstore/blob/6827eaab2484c5891e567b8462280ab5c74bff0e/internal/database/offline/sqlutil/export.go#L144-L163

If there is no place to consume the stream, the goroutine will be leaked

Pull request description:
refactor(offline/import): refactor offline export with context for avoid goroutine leck
fix #1157 

Patch:
diff --git a/internal/database/metadata/sqlutil/feature.go b/internal/database/metadata/sqlutil/feature.go
index a0f7f7e1..aa7c5381 100644
--- a/internal/database/metadata/sqlutil/feature.go
+++ b/internal/database/metadata/sqlutil/feature.go
@@ -114,14 +114,9 @@ func ListFeature(ctx context.Context, sqlxCtx metadata.SqlxContext, opt metadata
 	}
 
 	// filter by entity
-	if opt.EntityIDs != nil {
+	if opt.EntityID != nil {
 		features = features.Filter(func(f *types.Feature) bool {
-			for _, id := range *opt.EntityIDs {
-				if f.Group.EntityID == id {
-					return true
-				}
-			}
-			return false
+			return f.Group.EntityID == *opt.EntityID
 		})
 	}
 	return features, nil
@@ -131,11 +126,8 @@ func buildListFeatureCond(opt metadata.ListFeatureOpt) ([]string, []interface{},
 	in := make(map[string]interface{})
 	and := make(map[string]interface{})
 
-	if opt.GroupIDs != nil {
-		if len(*opt.GroupIDs) == 0 {
-			return []string{"false"}, nil, nil
-		}
-		in["group_id"] = *opt.GroupIDs
+	if opt.GroupID != nil {
+		and["group_id"] = *opt.GroupID
 	}
 
 	if opt.FeatureIDs != nil {
diff --git a/internal/database/metadata/types.go b/internal/database/metadata/types.go
index 3f31a173..5a51cd0c 100644
--- a/internal/database/metadata/types.go
+++ b/internal/database/metadata/types.go
@@ -61,8 +61,8 @@ type UpdateRevisionOpt struct {
 }
 
 type ListFeatureOpt struct {
-	EntityIDs  *[]int
-	GroupIDs   *[]int
+	EntityID   *int
+	GroupID    *int
 	FeatureIDs *[]int
 }
 
diff --git a/internal/database/offline/bigquery/export.go b/internal/database/offline/bigquery/export.go
index 6e8e9dd0..2b9d75b5 100644
--- a/internal/database/offline/bigquery/export.go
+++ b/internal/database/offline/bigquery/export.go
@@ -36,27 +36,41 @@ func bigqueryQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt off
 
 	go func() {
 		defer close(stream)
+
 		rows, err := dbOpt.BigQueryDB.Query(query).Read(ctx)
 		if err != nil {
-			stream <- types.ExportRecord{Error: err}
-			return
+			select {
+			case stream <- types.ExportRecord{Error: err}:
+				return
+			case <-ctx.Done():
+				return
+			}
 		}
 		for {
 			recordMap := make(map[string]bigquery.Value)
 			err = rows.Next(&recordMap)
 			if err == iterator.Done {
-				break
+				return
 			}
 			if err != nil {
-				stream <- types.ExportRecord{Error: errdefs.Errorf("failed at rows.Next, err=%v", err)}
-				return
+				select {
+				case stream <- types.ExportRecord{Error: errdefs.Errorf("failed at rows.Next, err=%v", err)}:
+					return
+				case <-ctx.Done():
+					return
+				}
 			}
 			record := make([]interface{}, 0, len(recordMap))
 			record = append(record, recordMap[opt.EntityName])
 			for _, feature := range features {
 				record = append(record, recordMap[feature.DBFullName(Backend)])
 			}
-			stream <- types.ExportRecord{Record: record}
+			select {
+			case stream <- types.ExportRecord{Record: record}:
+				// nothing to do
+			case <-ctx.Done():
+				return
+			}
 		}
 	}()
 	header := append([]string{opt.EntityName}, features.FullNames()...)
diff --git a/internal/database/offline/sqlutil/export.go b/internal/database/offline/sqlutil/export.go
index 9d28bac0..a5314488 100644
--- a/internal/database/offline/sqlutil/export.go
+++ b/internal/database/offline/sqlutil/export.go
@@ -34,13 +34,6 @@ type DoExportOpt struct {
 }
 
 func DoExport(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOpt) (*types.ExportResult, error) {
-	var (
-		emptyStream = make(chan types.ExportRecord)
-		errs        = make(chan error, 1) // at most 1 error
-	)
-	defer close(emptyStream)
-	defer close(errs)
-
 	// Step 0: prepare variables
 	var (
 		snapshotTables = make([]string, 0, len(opt.SnapshotTables))
@@ -129,23 +122,35 @@ func sqlxQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline
 		defer close(stream)
 		stmt, err := dbOpt.SqlxDB.Preparex(dbOpt.SqlxDB.Rebind(query))
 		if err != nil {
-			stream <- types.ExportRecord{Error: errdefs.WithStack(err)}
-			return
+			select {
+			case stream <- types.ExportRecord{Error: errdefs.WithStack(err)}:
+				return
+			case <-ctx.Done():
+				return
+			}
 		}
 		defer stmt.Close()
 
 		rows, err := stmt.Queryx(args...)
 		if err != nil {
-			stream <- types.ExportRecord{Error: errdefs.WithStack(err)}
-			return
+			select {
+			case stream <- types.ExportRecord{Error: errdefs.WithStack(err)}:
+				return
+			case <-ctx.Done():
+				return
+			}
 		}
 		defer rows.Close()
 
 		for rows.Next() {
 			record, err := rows.SliceScan()
 			if err != nil {
-				stream <- types.ExportRecord{Error: errdefs.Errorf("failed at rows.SliceScan, err=%v", err)}
-				return
+				select {
+				case stream <- types.ExportRecord{Error: errdefs.Errorf("failed at rows.SliceScan, err=%v", err)}:
+					return
+				case <-ctx.Done():
+					return
+				}
 			}
 			record[0] = cast.ToString(record[0])
 			for i, f := range features {
@@ -154,12 +159,21 @@ func sqlxQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline
 				}
 				deserializedValue, err := dbutil.DeserializeByValueType(record[i+1], f.ValueType, dbOpt.Backend)
 				if err != nil {
-					stream <- types.ExportRecord{Error: err}
-					return
+					select {
+					case stream <- types.ExportRecord{Error: err}:
+						return
+					case <-ctx.Done():
+						return
+					}
 				}
 				record[i+1] = deserializedValue
 			}
-			stream <- types.ExportRecord{Record: record, Error: nil}
+			select {
+			case stream <- types.ExportRecord{Record: record, Error: nil}:
+				// nothing to do
+			case <-ctx.Done():
+				return
+			}
 		}
 	}()
 	header := append([]string{opt.EntityName}, features.FullNames()...)
diff --git a/oomagent/server.go b/oomagent/server.go
index 10d28266..9b62006e 100644
--- a/oomagent/server.go
+++ b/oomagent/server.go
@@ -242,7 +242,7 @@ func (s *server) ChannelJoin(stream codegen.OomAgent_ChannelJoinServer) error {
 				UnixMilli: req.EntityRow.UnixMilli,
 				Values:    req.EntityRow.Values,
 			}:
-				// nothing to do
+				// nothing todo
 			case <-ctx.Done():
 				return
 			}
@@ -299,7 +299,9 @@ func (s *server) ChannelExport(req *codegen.ChannelExportRequest, stream codegen
 		return nil
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	exportResult, err := s.oomstore.ChannelExport(ctx, types.ChannelExportOpt{
 		FeatureNames: req.Features,
 		UnixMilli:    req.UnixMilli,
diff --git a/oomcli/cmd/export_helper.go b/oomcli/cmd/export_helper.go
index dce1215c..2d164de2 100644
--- a/oomcli/cmd/export_helper.go
+++ b/oomcli/cmd/export_helper.go
@@ -15,6 +15,9 @@ import (
 )
 
 func export(ctx context.Context, store *oomstore.OomStore, opt types.ChannelExportOpt, output string) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	exportResult, err := store.ChannelExport(ctx, opt)
 	if err != nil {
 		return err
diff --git a/pkg/oomstore/apply.go b/pkg/oomstore/apply.go
index d58dba3a..50bb1d41 100644
--- a/pkg/oomstore/apply.go
+++ b/pkg/oomstore/apply.go
@@ -192,7 +192,7 @@ func (s *OomStore) applyFeature(ctx context.Context, tx metadata.DBStore, newFea
 
 		if group.Category == types.CategoryStream {
 			features, err := tx.ListFeature(ctx, metadata.ListFeatureOpt{
-				GroupIDs: &[]int{group.ID},
+				GroupID: &group.ID,
 			})
 			if err != nil {
 				return nil, err
diff --git a/pkg/oomstore/export.go b/pkg/oomstore/export.go
index 5ffe76bd..bc320570 100644
--- a/pkg/oomstore/export.go
+++ b/pkg/oomstore/export.go
@@ -99,6 +99,9 @@ func (s *OomStore) ChannelExport(ctx context.Context, opt types.ChannelExportOpt
 // Export exports the latest feature value up to the given timestamp, it outputs
 // feature values to the given file path.
 func (s *OomStore) Export(ctx context.Context, opt types.ExportOpt) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	exportResult, err := s.ChannelExport(ctx, types.ChannelExportOpt{
 		FeatureNames: opt.FeatureNames,
 		UnixMilli:    opt.UnixMilli,
diff --git a/pkg/oomstore/feature.go b/pkg/oomstore/feature.go
index 6d9f4307..3ca0d094 100644
--- a/pkg/oomstore/feature.go
+++ b/pkg/oomstore/feature.go
@@ -39,14 +39,14 @@ func (s *OomStore) ListFeature(ctx context.Context, opt types.ListFeatureOpt) (t
 		if err != nil {
 			return nil, err
 		}
-		metadataOpt.EntityIDs = &[]int{entity.ID}
+		metadataOpt.EntityID = &entity.ID
 	}
 	if opt.GroupName != nil {
 		group, err := s.metadata.GetGroupByName(ctx, *opt.GroupName)
 		if err != nil {
 			return nil, err
 		}
-		metadataOpt.GroupIDs = &[]int{group.ID}
+		metadataOpt.GroupID = &group.ID
 	}
 	features, err := s.metadata.ListFeature(ctx, metadataOpt)
 	if err != nil {
@@ -100,7 +100,7 @@ func (s *OomStore) CreateFeature(ctx context.Context, opt types.CreateFeatureOpt
 
 	if group.Category == types.CategoryStream {
 		features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
-			GroupIDs: &[]int{group.ID},
+			GroupID: &group.ID,
 		})
 		if err != nil {
 			return 0, err
diff --git a/pkg/oomstore/import.go b/pkg/oomstore/import.go
index bb830006..4c01bf0b 100644
--- a/pkg/oomstore/import.go
+++ b/pkg/oomstore/import.go
@@ -67,7 +67,7 @@ func (s *OomStore) parseImportOpt(ctx context.Context, opt types.ImportOpt) (*im
 	}
 
 	features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
-		GroupIDs: &[]int{group.ID},
+		GroupID: &group.ID,
 	})
 	if err != nil {
 		return nil, err
diff --git a/pkg/oomstore/revision.go b/pkg/oomstore/revision.go
index 49a9e995..3e651719 100644
--- a/pkg/oomstore/revision.go
+++ b/pkg/oomstore/revision.go
@@ -99,7 +99,7 @@ func (s *OomStore) createSnapshotAndCdcTable(ctx context.Context, revisionID int
 
 	// Create snapshot table in offline store
 	features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
-		GroupIDs: &[]int{revision.GroupID},
+		GroupID: &revision.GroupID,
 	})
 	if err != nil {
 		return err
diff --git a/pkg/oomstore/snapshot.go b/pkg/oomstore/snapshot.go
index 07215a70..cd65cd14 100644
--- a/pkg/oomstore/snapshot.go
+++ b/pkg/oomstore/snapshot.go
@@ -19,7 +19,7 @@ func (s *OomStore) Snapshot(ctx context.Context, groupName string) error {
 	}
 
 	features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
-		GroupIDs: &[]int{group.ID},
+		GroupID: &group.ID,
 	})
 	if err != nil {
 		return err
diff --git a/pkg/oomstore/sync.go b/pkg/oomstore/sync.go
index 6853ae6b..e9fa904c 100644
--- a/pkg/oomstore/sync.go
+++ b/pkg/oomstore/sync.go
@@ -30,6 +30,9 @@ func (s *OomStore) Sync(ctx context.Context, opt types.SyncOpt) error {
 
 // syncBatch syncs batch feature group from offline store to online store.
 func (s *OomStore) syncBatch(ctx context.Context, opt types.SyncOpt, group *types.Group, revision *types.Revision) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	prevOnlineRevisionID := group.OnlineRevisionID
 	if prevOnlineRevisionID != nil && *prevOnlineRevisionID == revision.ID {
 		return errdefs.Errorf("the specific revision was synced to the online store, won't do it again this time")
@@ -98,6 +101,9 @@ func (s *OomStore) syncBatch(ctx context.Context, opt types.SyncOpt, group *type
 
 // syncStream syncs stream feature group from offline store to online store.
 func (s *OomStore) syncStream(ctx context.Context, opt types.SyncOpt) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	features, err := s.ListFeature(ctx, types.ListFeatureOpt{
 		GroupName: &opt.GroupName,
 	})


Test patch:
diff --git a/internal/database/metadata/test_impl/feature.go b/internal/database/metadata/test_impl/feature.go
index 1248f066..7c6a4b8b 100644
--- a/internal/database/metadata/test_impl/feature.go
+++ b/internal/database/metadata/test_impl/feature.go
@@ -233,7 +233,7 @@ func TestListFeature(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 
 	// case 4: list features by EntityID and FeatureIDs
 	features, err = store.ListFeature(ctx, metadata.ListFeatureOpt{
-		EntityIDs:  &[]int{entityID + 1},
+		EntityID:   intPtr(entityID + 1),
 		FeatureIDs: &[]int{featureID},
 	})
 	assert.NoError(t, err)
@@ -241,7 +241,7 @@ func TestListFeature(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 
 	// case 5: list features by GroupID and FeatureIDs
 	features, err = store.ListFeature(ctx, metadata.ListFeatureOpt{
-		GroupIDs:   &[]int{groupID + 1},
+		GroupID:    intPtr(groupID + 1),
 		FeatureIDs: &[]int{featureID},
 	})
 	assert.NoError(t, err)
@@ -249,32 +249,18 @@ func TestListFeature(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 
 	// case 6: list features by EntityID and empty FeatureIDs, return no feature
 	features, err = store.ListFeature(ctx, metadata.ListFeatureOpt{
-		EntityIDs:  &[]int{entityID},
+		EntityID:   &entityID,
 		FeatureIDs: &[]int{},
 	})
 	assert.NoError(t, err)
 	assert.Equal(t, 0, len(features))
 
-	// case 7: list features by EntityIDs
+	// case 7: list features by EntityID
 	features, err = store.ListFeature(ctx, metadata.ListFeatureOpt{
-		EntityIDs: &[]int{entityID, entityID + 1},
+		EntityID: &entityID,
 	})
 	assert.NoError(t, err)
 	assert.Equal(t, 1, len(features))
-
-	// case 8: list features by empty EntityIDs
-	features, err = store.ListFeature(ctx, metadata.ListFeatureOpt{
-		EntityIDs: &[]int{},
-	})
-	assert.NoError(t, err)
-	assert.Equal(t, 0, len(features))
-
-	// case 9: list features by empty GroupIDs
-	features, err = store.ListFeature(ctx, metadata.ListFeatureOpt{
-		GroupIDs: &[]int{},
-	})
-	assert.NoError(t, err)
-	assert.Equal(t, 0, len(features))
 }
 
 func TestUpdateFeature(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyStoreFn) {
diff --git a/internal/database/metadata/test_impl/helper.go b/internal/database/metadata/test_impl/helper.go
index 4467438a..feeece08 100644
--- a/internal/database/metadata/test_impl/helper.go
+++ b/internal/database/metadata/test_impl/helper.go
@@ -4,6 +4,10 @@ func boolPtr(i bool) *bool {
 	return &i
 }
 
+func intPtr(i int) *int {
+	return &i
+}
+
 func stringPtr(s string) *string {
 	return &s
 }