refactor calc recom score cron job

OpenCSGs · Jan 22, 2025 · 02cce9e · 02cce9e
1 parent a8d401b
commit 02cce9e
Show file tree

Hide file tree

Showing 7 changed files with 86 additions and 38 deletions.
diff --git a/_mocks/opencsg.com/csghub-server/component/mock_RecomComponent.go b/_mocks/opencsg.com/csghub-server/component/mock_RecomComponent.go
diff --git a/api/workflow/activity/calc_recom_score.go b/api/workflow/activity/calc_recom_score.go
@@ -5,6 +5,5 @@ import (
 )
 
 func (a *Activities) CalcRecomScore(ctx context.Context) error {
-	a.recom.CalculateRecomScore(context.Background())
-	return nil
+	return a.recom.CalculateRecomScore(ctx, 0)
 }
diff --git a/api/workflow/cron_calc_recom_score.go b/api/workflow/cron_calc_recom_score.go
@@ -16,7 +16,7 @@ func CalcRecomScoreWorkflow(ctx workflow.Context) error {
 	}
 
 	options := workflow.ActivityOptions{
-		StartToCloseTimeout: time.Hour * 1,
+		StartToCloseTimeout: time.Hour * 2,
 		RetryPolicy:         retryPolicy,
 	}
 
@@ -26,5 +26,6 @@ func CalcRecomScoreWorkflow(ctx workflow.Context) error {
 		logger.Error("failed to calc recom score", "error", err)
 		return err
 	}
+	logger.Info("calc recom score workflow completed")
 	return nil
 }
diff --git a/api/workflow/schedule_ce_test.go b/api/workflow/schedule_ce_test.go
@@ -3,6 +3,7 @@
 package workflow_test
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/mock"
@@ -12,13 +13,25 @@ import (
 )
 
 func TestSchedule_CalcRecomScoreWorkflow(t *testing.T) {
-	tester, err := newWorkflowTester(t)
-	require.NoError(t, err)
-
-	tester.mocks.recom.EXPECT().CalculateRecomScore(mock.Anything).Return()
-	tester.scheduler.Execute("calc-recom-score-schedule", tester.cronEnv)
-	require.True(t, tester.cronEnv.IsWorkflowCompleted())
-	require.NoError(t, tester.cronEnv.GetWorkflowError())
+	t.Run("normal", func(t *testing.T) {
+		tester, err := newWorkflowTester(t)
+		require.NoError(t, err)
+
+		tester.mocks.recom.EXPECT().CalculateRecomScore(mock.Anything, 0).Return(nil)
+		tester.scheduler.Execute("calc-recom-score-schedule", tester.cronEnv)
+		require.True(t, tester.cronEnv.IsWorkflowCompleted())
+		require.NoError(t, tester.cronEnv.GetWorkflowError())
+	})
+
+	t.Run("error", func(t *testing.T) {
+		tester, err := newWorkflowTester(t)
+		require.NoError(t, err)
+
+		tester.mocks.recom.EXPECT().CalculateRecomScore(mock.Anything, 0).Return(errors.New("error"))
+		tester.scheduler.Execute("calc-recom-score-schedule", tester.cronEnv)
+		require.True(t, tester.cronEnv.IsWorkflowCompleted())
+		require.Error(t, tester.cronEnv.GetWorkflowError())
+	})
 }
 
 func TestSchedule_SyncAsClient(t *testing.T) {

diff --git a/cmd/csghub-server/cmd/cron/calc_recom_score.go b/cmd/csghub-server/cmd/cron/calc_recom_score.go
@@ -45,6 +45,9 @@ var cmdCalcRecomScore = &cobra.Command{
 			slog.Error("failed to create recom component", "err", err)
 			return
 		}
-		c.CalculateRecomScore(cmd.Context())
+		err = c.CalculateRecomScore(cmd.Context(), 0)
+		if err != nil {
+			slog.Error("failed to calculate recom score", "err", err)
+		}
 	},
 }
diff --git a/component/recom.go b/component/recom.go
@@ -2,6 +2,7 @@ package component
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"time"
@@ -22,7 +23,7 @@ type recomComponentImpl struct {
 type RecomComponent interface {
 	SetOpWeight(ctx context.Context, repoID, weight int64) error
 	// loop through repositories and calculate the recom score of the repository
-	CalculateRecomScore(ctx context.Context)
+	CalculateRecomScore(ctx context.Context, batchSize int) error
 	CalcTotalScore(ctx context.Context, repo *database.Repository, weights map[string]string) float64
 }
 
@@ -48,26 +49,38 @@ func (rc *recomComponentImpl) SetOpWeight(ctx context.Context, repoID, weight in
 }
 
 // loop through repositories and calculate the recom score of the repository
-func (rc *recomComponentImpl) CalculateRecomScore(ctx context.Context) {
+func (rc *recomComponentImpl) CalculateRecomScore(ctx context.Context, batchSize int) error {
 	weights, err := rc.loadWeights()
 	if err != nil {
-		slog.Error("Error loading weights", "error", err)
-		return
+		return errors.New("error loading weights")
 	}
-	repos, err := rc.repoStore.All(ctx)
-	if err != nil {
-		slog.Error("Error fetching repositories", "error", err)
-		return
+	if batchSize <= 0 {
+		batchSize = 1000
 	}
-	for _, repo := range repos {
-		repoID := repo.ID
-		score := rc.CalcTotalScore(ctx, repo, weights)
-		err := rc.recomStore.UpsertScore(ctx, repoID, score)
+	batch := 0
+	for {
+		repos, err := rc.repoStore.FindWithBatch(ctx, batchSize, batch)
 		if err != nil {
-			slog.Error("Error updating recom score", slog.Int64("repo_id", repoID), slog.Float64("score", score),
-				slog.String("error", err.Error()))
+			return errors.New("error fetching repositories")
+		}
+		for _, repo := range repos {
+			repoID := repo.ID
+			score := rc.CalcTotalScore(ctx, &repo, weights)
+			err := rc.recomStore.UpsertScore(ctx, repoID, score)
+			if err != nil {
+				slog.Error("Error updating recom score", slog.Int64("repo_id", repoID), slog.Float64("score", score),
+					slog.String("error", err.Error()))
+			}
+		}
+
+		if len(repos) < batchSize {
+			break
 		}
+
+		batch++
 	}
+
+	return nil
 }
 
 func (rc *recomComponentImpl) CalcTotalScore(ctx context.Context, repo *database.Repository, weights map[string]string) float64 {
@@ -83,7 +96,7 @@ func (rc *recomComponentImpl) CalcTotalScore(ctx context.Context, repo *database
 
 	qualityScore, err := rc.calcQualityScore(ctx, repo)
 	if err != nil {
-		slog.Error("failed to calculate quality score", slog.Any("error", err))
+		slog.Error("failed to calculate quality score", slog.Any("error", err), slog.Int64("repo_id", repo.ID), slog.String("repo_path", repo.Path))
 	} else {
 		score += qualityScore
 	}

diff --git a/component/recom_test.go b/component/recom_test.go
@@ -31,16 +31,21 @@ func TestRecomComponent_CalculateRecomScore(t *testing.T) {
 	rc.mocks.stores.RecomMock().EXPECT().LoadWeights(mock.Anything).Return(
 		[]*database.RecomWeight{{Name: "freshness", WeightExp: "score = 12.34"}}, nil,
 	)
-	rc.mocks.stores.RepoMock().EXPECT().All(ctx).Return([]*database.Repository{
+	// loop 1
+	rc.mocks.stores.RepoMock().EXPECT().FindWithBatch(ctx, 1, 0).Return([]database.Repository{
 		{ID: 1, Path: "foo/bar"},
 	}, nil)
+	// loop 2
+	rc.mocks.stores.RepoMock().EXPECT().FindWithBatch(ctx, 1, 1).Return([]database.Repository{}, nil)
+
 	rc.mocks.stores.RecomMock().EXPECT().UpsertScore(ctx, int64(1), 12.34).Return(nil)
 	rc.mocks.gitServer.EXPECT().GetRepoFileTree(mock.Anything, gitserver.GetRepoInfoByPathReq{
 		Namespace: "foo",
 		Name:      "bar",
 	}).Return(nil, nil)
 
-	rc.CalculateRecomScore(ctx)
+	err := rc.CalculateRecomScore(ctx, 1)
+	require.NoError(t, err)
 }
 
 func TestRecomComponent_CalculateTotalScore(t *testing.T) {