初始提交：边缘AI算力机统一AI通讯层

2026-08-03 07:44:05 +08:00
commit 93a469061d
51 changed files with 11565 additions and 0 deletions
@@ -0,0 +1,174 @@
+package task
+
+import (
+	"database/sql"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	_ "github.com/mattn/go-sqlite3"
+)
+
+// Store manages task state persistence with SQLite.
+type Store struct {
+	mu sync.Mutex
+	db *sql.DB
+}
+
+// NewStore creates a new task store.
+func NewStore(dbPath string) (*Store, error) {
+	db, err := sql.Open("sqlite3", dbPath)
+	if err != nil {
+		return nil, fmt.Errorf("open task db: %w", err)
+	}
+
+	if err := initTaskDB(db); err != nil {
+		return nil, fmt.Errorf("init task db: %w", err)
+	}
+
+	return &Store{db: db}, nil
+}
+
+func initTaskDB(db *sql.DB) error {
+	schema := `
+	CREATE TABLE IF NOT EXISTS tasks (
+		id TEXT PRIMARY KEY,
+		request_id TEXT NOT NULL,
+		session_id TEXT,
+		app_id TEXT NOT NULL,
+		tenant_id TEXT NOT NULL,
+		logical_model TEXT NOT NULL,
+		actual_model TEXT,
+		priority INTEGER NOT NULL DEFAULT 2,
+		state TEXT NOT NULL,
+		stream INTEGER NOT NULL DEFAULT 0,
+		created_at TEXT NOT NULL,
+		started_at TEXT,
+		completed_at TEXT,
+		cancel_reason TEXT,
+		error_message TEXT,
+		input_tokens INTEGER DEFAULT 0,
+		output_tokens INTEGER DEFAULT 0,
+		node_id TEXT,
+		degraded INTEGER DEFAULT 0
+	);
+	CREATE INDEX IF NOT EXISTS idx_tasks_state ON tasks(state);
+	CREATE INDEX IF NOT EXISTS idx_tasks_app ON tasks(app_id);
+	CREATE INDEX IF NOT EXISTS idx_tasks_tenant ON tasks(tenant_id);`
+	_, err := db.Exec(schema)
+	return err
+}
+
+// Save persists a task to the database.
+func (s *Store) Save(t *Task) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var startedAt, completedAt interface{}
+	if t.StartedAt != nil {
+		startedAt = t.StartedAt.Format(time.RFC3339)
+	}
+	if t.CompletedAt != nil {
+		completedAt = t.CompletedAt.Format(time.RFC3339)
+	}
+
+	streamInt := 0
+	if t.Stream {
+		streamInt = 1
+	}
+	degradedInt := 0
+	if t.Degraded {
+		degradedInt = 1
+	}
+
+	_, err := s.db.Exec(
+		`INSERT OR REPLACE INTO tasks
+		(id, request_id, session_id, app_id, tenant_id, logical_model, actual_model, priority, state, stream, created_at, started_at, completed_at, cancel_reason, error_message, input_tokens, output_tokens, node_id, degraded)
+		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+		t.ID, t.RequestID, t.SessionID, t.AppID, t.TenantID, t.LogicalModel, t.ActualModel,
+		int(t.Priority), string(t.State), streamInt, t.CreatedAt.Format(time.RFC3339),
+		startedAt, completedAt, t.CancelReason, t.ErrorMessage,
+		t.InputTokens, t.OutputTokens, t.NodeID, degradedInt,
+	)
+	return err
+}
+
+// Get retrieves a task by ID.
+func (s *Store) Get(id string) (*Task, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var (
+		requestID, sessionID, appID, tenantID, logicalModel, actualModel, state  string
+		priority                                                                 int
+		streamInt                                                                int
+		createdAtStr, startedAt, completedAt, cancelReason, errorMessage, nodeID sql.NullString
+		inputTokens, outputTokens, degradedInt                                   int
+	)
+
+	err := s.db.QueryRow(
+		`SELECT request_id, session_id, app_id, tenant_id, logical_model, actual_model, priority, state, stream, created_at, started_at, completed_at, cancel_reason, error_message, input_tokens, output_tokens, node_id, degraded FROM tasks WHERE id = ?`,
+		id,
+	).Scan(&requestID, &sessionID, &appID, &tenantID, &logicalModel, &actualModel, &priority, &state, &streamInt, &createdAtStr, &startedAt, &completedAt, &cancelReason, &errorMessage, &inputTokens, &outputTokens, &nodeID, &degradedInt)
+
+	if err == sql.ErrNoRows {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	t := &Task{
+		ID:           id,
+		RequestID:    requestID,
+		SessionID:    sessionID,
+		AppID:        appID,
+		TenantID:     tenantID,
+		LogicalModel: logicalModel,
+		ActualModel:  actualModel,
+		Priority:     TaskPriority(priority),
+		State:        TaskState(state),
+		Stream:       streamInt == 1,
+		InputTokens:  inputTokens,
+		OutputTokens: outputTokens,
+		NodeID:       nodeID.String,
+		Degraded:     degradedInt == 1,
+		CancelReason: cancelReason.String,
+		ErrorMessage: errorMessage.String,
+		cancelCh:     make(chan struct{}),
+	}
+	t.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr.String)
+	if startedAt.Valid {
+		tt, _ := time.Parse(time.RFC3339, startedAt.String)
+		t.StartedAt = &tt
+	}
+	if completedAt.Valid {
+		tt, _ := time.Parse(time.RFC3339, completedAt.String)
+		t.CompletedAt = &tt
+	}
+
+	return t, nil
+}
+
+// RecoverPendingTasks marks RUNNING/STREAMING tasks as FAILED on startup.
+func (s *Store) RecoverPendingTasks() (int, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	result, err := s.db.Exec(
+		`UPDATE tasks SET state = 'FAILED', error_message = 'gateway restart' WHERE state IN ('RUNNING', 'STREAMING')`)
+	if err != nil {
+		return 0, err
+	}
+	n, _ := result.RowsAffected()
+	return int(n), nil
+}
+
+// Close closes the database connection.
+func (s *Store) Close() error {
+	return s.db.Close()
+}
+
+// Ensure json is imported for future use.
+var _ = json.Marshal
@@ -0,0 +1,190 @@
+package task
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/edgeai/gateway/internal/observability"
+)
+
+// TaskState represents the lifecycle state of a task.
+type TaskState string
+
+const (
+	StateQueued    TaskState = "QUEUED"
+	StateRunning   TaskState = "RUNNING"
+	StateStreaming TaskState = "STREAMING"
+	StateCompleted TaskState = "COMPLETED"
+	StateFailed    TaskState = "FAILED"
+	StateCancelled TaskState = "CANCELLED"
+)
+
+// TaskPriority levels (P0 highest, P4 lowest).
+type TaskPriority int
+
+const (
+	PriorityRealtime   TaskPriority = 0 // P0
+	PriorityHigh       TaskPriority = 1 // P1
+	PriorityNormal     TaskPriority = 2 // P2 (default)
+	PriorityLow        TaskPriority = 3 // P3
+	PriorityBackground TaskPriority = 4 // P4
+)
+
+// Task represents an inference task in the system.
+type Task struct {
+	ID           string
+	RequestID    string
+	SessionID    string
+	AppID        string
+	TenantID     string
+	LogicalModel string
+	ActualModel  string
+	Priority     TaskPriority
+	State        TaskState
+	Stream       bool
+	CreatedAt    time.Time
+	StartedAt    *time.Time
+	CompletedAt  *time.Time
+	CancelReason string
+	ErrorMessage string
+	InputTokens  int
+	OutputTokens int
+	NodeID       string
+	Degraded     bool
+	cancelCh     chan struct{}
+	cancelOnce   sync.Once
+	mu           sync.RWMutex
+}
+
+// NewTask creates a new task in QUEUED state.
+func NewTask(id, requestID, appID, tenantID, logicalModel string, priority TaskPriority, stream bool) *Task {
+	return &Task{
+		ID:           id,
+		RequestID:    requestID,
+		AppID:        appID,
+		TenantID:     tenantID,
+		LogicalModel: logicalModel,
+		Priority:     priority,
+		State:        StateQueued,
+		Stream:       stream,
+		CreatedAt:    time.Now(),
+		cancelCh:     make(chan struct{}),
+	}
+}
+
+// AllowedTransitions defines valid state transitions.
+var allowedTransitions = map[TaskState][]TaskState{
+	StateQueued:    {StateRunning, StateFailed, StateCancelled},
+	StateRunning:   {StateStreaming, StateCompleted, StateFailed, StateCancelled},
+	StateStreaming: {StateCompleted, StateFailed, StateCancelled},
+	StateCompleted: {},
+	StateFailed:    {},
+	StateCancelled: {},
+}
+
+// Transition changes the task state if the transition is valid.
+func (t *Task) Transition(to TaskState) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	allowed, ok := allowedTransitions[t.State]
+	if !ok {
+		return fmt.Errorf("unknown current state: %s", t.State)
+	}
+
+	valid := false
+	for _, s := range allowed {
+		if s == to {
+			valid = true
+			break
+		}
+	}
+	if !valid {
+		return fmt.Errorf("invalid transition: %s -> %s", t.State, to)
+	}
+
+	from := t.State
+	t.State = to
+	now := time.Now()
+
+	switch to {
+	case StateRunning:
+		t.StartedAt = &now
+	case StateCompleted, StateFailed, StateCancelled:
+		t.CompletedAt = &now
+	}
+
+	_ = from
+	return nil
+}
+
+// Cancel signals task cancellation and transitions to CANCELLED if possible.
+func (t *Task) Cancel(reason string) error {
+	t.cancelOnce.Do(func() {
+		close(t.cancelCh)
+	})
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	if t.State == StateCompleted || t.State == StateFailed || t.State == StateCancelled {
+		return errors.New("task already in terminal state")
+	}
+
+	t.CancelReason = reason
+	t.State = StateCancelled
+	now := time.Now()
+	t.CompletedAt = &now
+	return nil
+}
+
+// Cancelled returns a channel that's closed when the task is cancelled.
+func (t *Task) Cancelled() <-chan struct{} {
+	return t.cancelCh
+}
+
+// IsCancelled returns true if the task has been cancelled.
+func (t *Task) IsCancelled() bool {
+	select {
+	case <-t.cancelCh:
+		return true
+	default:
+		return false
+	}
+}
+
+// GetState returns the current state (thread-safe).
+func (t *Task) GetState() TaskState {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.State
+}
+
+// IsTerminal returns true if the task is in a terminal state.
+func (t *Task) IsTerminal() bool {
+	s := t.GetState()
+	return s == StateCompleted || s == StateFailed || s == StateCancelled
+}
+
+// StateMachineLogger logs state transitions.
+type StateMachineLogger struct {
+	logger *observability.Logger
+}
+
+func NewStateMachineLogger(logger *observability.Logger) *StateMachineLogger {
+	return &StateMachineLogger{logger: logger}
+}
+
+// LogTransition logs a state transition.
+func (sml *StateMachineLogger) LogTransition(task *Task, from, to TaskState, reason string) {
+	sml.logger.Info("task state transition",
+		observability.F().
+			Event("state_transition").
+			TaskID(task.ID).
+			Set("from_state", string(from)).
+			Set("to_state", string(to)).
+			Reason(reason))
+	_ = from // used in log field above
+}
@@ -0,0 +1,146 @@
+package task
+
+import (
+	"testing"
+	"time"
+)
+
+func TestNewTask(t *testing.T) {
+	task := NewTask("task-1", "req-1", "app-1", "tenant-1", "general-chat", PriorityNormal, false)
+	if task.ID != "task-1" {
+		t.Errorf("expected ID task-1, got %s", task.ID)
+	}
+	if task.State != StateQueued {
+		t.Errorf("expected state QUEUED, got %s", task.State)
+	}
+	if task.Priority != PriorityNormal {
+		t.Errorf("expected priority P2, got %d", task.Priority)
+	}
+}
+
+func TestValidTransitions(t *testing.T) {
+	tests := []struct {
+		from TaskState
+		to   TaskState
+		ok   bool
+	}{
+		{StateQueued, StateRunning, true},
+		{StateQueued, StateFailed, true},
+		{StateQueued, StateCancelled, true},
+		{StateQueued, StateCompleted, false},
+		{StateRunning, StateStreaming, true},
+		{StateRunning, StateCompleted, true},
+		{StateRunning, StateFailed, true},
+		{StateRunning, StateCancelled, true},
+		{StateRunning, StateQueued, false},
+		{StateStreaming, StateCompleted, true},
+		{StateStreaming, StateFailed, true},
+		{StateStreaming, StateCancelled, true},
+		{StateStreaming, StateRunning, false},
+		{StateCompleted, StateRunning, false},
+		{StateFailed, StateCompleted, false},
+		{StateCancelled, StateRunning, false},
+	}
+
+	for _, tt := range tests {
+		task := &Task{State: tt.from, cancelCh: make(chan struct{})}
+		err := task.Transition(tt.to)
+		if tt.ok && err != nil {
+			t.Errorf("expected %s -> %s to succeed, got error: %v", tt.from, tt.to, err)
+		}
+		if !tt.ok && err == nil {
+			t.Errorf("expected %s -> %s to fail, but it succeeded", tt.from, tt.to)
+		}
+	}
+}
+
+func TestTaskCancel(t *testing.T) {
+	task := NewTask("task-1", "req-1", "app-1", "tenant-1", "general-chat", PriorityNormal, false)
+
+	if task.IsCancelled() {
+		t.Error("task should not be cancelled initially")
+	}
+
+	err := task.Cancel("client_disconnect")
+	if err != nil {
+		t.Errorf("cancel failed: %v", err)
+	}
+
+	if !task.IsCancelled() {
+		t.Error("task should be cancelled after Cancel()")
+	}
+
+	if task.GetState() != StateCancelled {
+		t.Errorf("expected state CANCELLED, got %s", task.GetState())
+	}
+
+	if task.CancelReason != "client_disconnect" {
+		t.Errorf("expected cancel reason 'client_disconnect', got %s", task.CancelReason)
+	}
+
+	// Cancel again should fail
+	err = task.Cancel("second_attempt")
+	if err == nil {
+		t.Error("expected error on double cancel")
+	}
+}
+
+func TestTaskCancelledChannel(t *testing.T) {
+	task := NewTask("task-1", "req-1", "app-1", "tenant-1", "general-chat", PriorityNormal, false)
+
+	select {
+	case <-task.Cancelled():
+		t.Error("channel should not be closed before cancel")
+	default:
+	}
+
+	task.Cancel("test")
+
+	select {
+	case <-task.Cancelled():
+		// expected
+	case <-time.After(100 * time.Millisecond):
+		t.Error("channel should be closed after cancel")
+	}
+}
+
+func TestIsTerminal(t *testing.T) {
+	tests := []struct {
+		state    TaskState
+		terminal bool
+	}{
+		{StateQueued, false},
+		{StateRunning, false},
+		{StateStreaming, false},
+		{StateCompleted, true},
+		{StateFailed, true},
+		{StateCancelled, true},
+	}
+
+	for _, tt := range tests {
+		task := &Task{State: tt.state}
+		if task.IsTerminal() != tt.terminal {
+			t.Errorf("expected IsTerminal()=%v for state %s, got %v", tt.terminal, tt.state, task.IsTerminal())
+		}
+	}
+}
+
+func TestTransitionSetsTimestamps(t *testing.T) {
+	task := &Task{State: StateQueued, cancelCh: make(chan struct{})}
+
+	err := task.Transition(StateRunning)
+	if err != nil {
+		t.Fatalf("transition to RUNNING failed: %v", err)
+	}
+	if task.StartedAt == nil {
+		t.Error("expected StartedAt to be set after transition to RUNNING")
+	}
+
+	err = task.Transition(StateCompleted)
+	if err != nil {
+		t.Fatalf("transition to COMPLETED failed: %v", err)
+	}
+	if task.CompletedAt == nil {
+		t.Error("expected CompletedAt to be set after transition to COMPLETED")
+	}
+}