104 lines
3.5 KiB
Go
104 lines
3.5 KiB
Go
|
package container
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"database/sql"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"gorm.io/driver/sqlite"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
const pendingOpsDDL = `
|
||
|
|
CREATE TABLE IF NOT EXISTS task_pending_ops (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
tenant_id INTEGER NOT NULL DEFAULT 0,
|
||
|
|
task_type VARCHAR(64) NOT NULL DEFAULT '',
|
||
|
|
scope VARCHAR(32) NOT NULL DEFAULT '',
|
||
|
|
scope_id VARCHAR(64) NOT NULL DEFAULT '',
|
||
|
|
op VARCHAR(32) NOT NULL DEFAULT '',
|
||
|
|
dedup_key VARCHAR(128) NOT NULL DEFAULT '',
|
||
|
|
payload TEXT,
|
||
|
|
fail_count INTEGER NOT NULL DEFAULT 0,
|
||
|
|
enqueued_at DATETIME,
|
||
|
|
claimed_at DATETIME
|
||
|
|
);
|
||
|
|
`
|
||
|
|
|
||
|
|
func openPendingOpsDB(t *testing.T) *gorm.DB {
|
||
|
|
t.Helper()
|
||
|
|
db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared&_uniq="+t.Name()), &gorm.Config{})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, db.Exec(pendingOpsDDL).Error)
|
||
|
|
t.Cleanup(func() {
|
||
|
|
if sqlDB, err := db.DB(); err == nil {
|
||
|
|
_ = sqlDB.Close()
|
||
|
|
}
|
||
|
|
})
|
||
|
|
return db
|
||
|
|
}
|
||
|
|
|
||
|
|
// claimOf reads the claim column as nullable, so a released claim is distinguishable
|
||
|
|
// from a timestamp.
|
||
|
|
func claimOf(t *testing.T, db *gorm.DB, id int64) (time.Time, bool) {
|
||
|
|
t.Helper()
|
||
|
|
var got sql.NullTime
|
||
|
|
require.NoError(t, db.Table("task_pending_ops").Select("claimed_at").Where("id = ?", id).Scan(&got).Error)
|
||
|
|
return got.Time, got.Valid
|
||
|
|
}
|
||
|
|
|
||
|
|
func seedPendingOp(t *testing.T, db *gorm.DB, taskType, scopeID string, claimedAt *time.Time) int64 {
|
||
|
|
t.Helper()
|
||
|
|
row := types.TaskPendingOp{
|
||
|
|
TenantID: 10059, TaskType: taskType, Scope: types.TaskScopeKnowledgeBase,
|
||
|
|
ScopeID: scopeID, Op: "ingest", EnqueuedAt: time.Now().UTC(), ClaimedAt: claimedAt,
|
||
|
|
}
|
||
|
|
require.NoError(t, db.Create(&row).Error)
|
||
|
|
return row.ID
|
||
|
|
}
|
||
|
|
|
||
|
|
// A claim stamped by a worker that has since died must be released at startup, so a
|
||
|
|
// restarted deployment resumes instead of waiting out the 90 minute stale threshold.
|
||
|
|
func TestSweepOrphanedWikiClaimsReleasesStaleClaimsOnly(t *testing.T) {
|
||
|
|
db := openPendingOpsDB(t)
|
||
|
|
now := time.Now().UTC()
|
||
|
|
stale := now.Add(-20 * time.Minute)
|
||
|
|
fresh := now.Add(-1 * time.Minute)
|
||
|
|
nonWikiStale := now.Add(-45 * time.Minute)
|
||
|
|
|
||
|
|
staleID := seedPendingOp(t, db, types.TypeWikiIngest, "kb-stale", &stale)
|
||
|
|
freshID := seedPendingOp(t, db, types.TypeWikiIngest, "kb-fresh", &fresh)
|
||
|
|
unclaimedID := seedPendingOp(t, db, types.TypeWikiIngest, "kb-unclaimed", nil)
|
||
|
|
nonWikiID := seedPendingOp(t, db, "some:other:task", "kb-other", &nonWikiStale)
|
||
|
|
|
||
|
|
released, err := sweepOrphanedWikiClaims(context.Background(), db, now)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.EqualValues(t, 1, released)
|
||
|
|
|
||
|
|
_, staleStillClaimed := claimOf(t, db, staleID)
|
||
|
|
require.False(t, staleStillClaimed, "the stale wiki claim is released")
|
||
|
|
|
||
|
|
_, freshStillClaimed := claimOf(t, db, freshID)
|
||
|
|
require.True(t, freshStillClaimed, "a live worker's fresh claim is untouched")
|
||
|
|
|
||
|
|
_, unclaimedStillClaimed := claimOf(t, db, unclaimedID)
|
||
|
|
require.False(t, unclaimedStillClaimed)
|
||
|
|
|
||
|
|
_, nonWikiStillClaimed := claimOf(t, db, nonWikiID)
|
||
|
|
require.True(t, nonWikiStillClaimed, "non-wiki queues keep their own claim policy")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Without the durable table the sweep is a no-op rather than an error, because the
|
||
|
|
// Lite deployment may run without the task queue at all.
|
||
|
|
func TestSweepOrphanedWikiClaimsWithoutTableIsNoop(t *testing.T) {
|
||
|
|
db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared&_uniq=notable"), &gorm.Config{})
|
||
|
|
require.NoError(t, err)
|
||
|
|
|
||
|
|
released, err := sweepOrphanedWikiClaims(context.Background(), db, time.Now())
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.EqualValues(t, 0, released)
|
||
|
|
}
|