tcjlj/OpenFlare
1
0
Fork 0
mirror of https://github.com/Rain-kl/OpenFlare.git synced 2026-10-03 23:06:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement heartbeat and sync services with error handling in Runner

Browse Source
This commit is contained in:
ryan
2026-03-10 14:44:05 +08:00
parent 7be2da0c19
commit f396c8c74e
3 changed files with 449 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files
atsf_agent/internal/agent/runner.go
+39 -10
View File
@@ -2,20 +2,29 @@ package agent
import (
"context"
"log"
"time"
"atsflare-agent/internal/config"
"atsflare-agent/internal/heartbeat"
"atsflare-agent/internal/protocol"
"atsflare-agent/internal/state"
syncservice "atsflare-agent/internal/sync"
)
type HeartbeatService interface {
Register(ctx context.Context, payload protocol.NodePayload) error
Heartbeat(ctx context.Context, payload protocol.NodePayload) error
}
type SyncService interface {
SyncOnStartup(ctx context.Context) error
SyncOnce(ctx context.Context) error
}
type Runner struct {
Config *config.Config
StateStore *state.Store
HeartbeatService *heartbeat.Service
SyncService *syncservice.Service
HeartbeatService HeartbeatService
SyncService SyncService
}
func (r *Runner) Run(ctx context.Context) error {
@@ -23,11 +32,15 @@ func (r *Runner) Run(ctx context.Context) error {
if err != nil {
return err
}
if err = r.SyncService.SyncOnStartup(ctx); err != nil {
return err
}
if err = r.HeartbeatService.Register(ctx, r.nodePayload(nodeID)); err != nil {
return err
log.Printf("agent register failed: %v", err)
}
if err = r.SyncService.SyncOnStartup(ctx); err != nil {
r.recordSyncError(err)
log.Printf("agent startup sync failed: %v", err)
}
if err = r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID)); err != nil {
log.Printf("agent startup heartbeat failed: %v", err)
}
heartbeatTicker := time.NewTicker(r.Config.HeartbeatInterval)
@@ -41,16 +54,32 @@ func (r *Runner) Run(ctx context.Context) error {
return ctx.Err()
case <-heartbeatTicker.C:
if err = r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID)); err != nil {
return err
log.Printf("agent heartbeat failed: %v", err)
}
case <-syncTicker.C:
if err = r.SyncService.SyncOnce(ctx); err != nil {
return err
r.recordSyncError(err)
log.Printf("agent sync failed: %v", err)
}
}
}
}
func (r *Runner) recordSyncError(err error) {
if err == nil || r.StateStore == nil {
return
}
snapshot, loadErr := r.StateStore.Load()
if loadErr != nil {
log.Printf("load state before recording sync error failed: %v", loadErr)
return
}
snapshot.LastError = err.Error()
if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
log.Printf("save state after sync error failed: %v", saveErr)
}
}
func (r *Runner) nodePayload(nodeID string) protocol.NodePayload {
snapshot, _ := r.StateStore.Load()
return protocol.NodePayload{
atsf_agent/internal/agent/runner_test.go
+172
View File
@@ -0,0 +1,172 @@
package agent
import (
"context"
"errors"
"path/filepath"
"sync"
"testing"
"time"
"atsflare-agent/internal/config"
"atsflare-agent/internal/protocol"
"atsflare-agent/internal/state"
)
type fakeHeartbeatService struct {
mu sync.Mutex
registerCalls int
heartbeatCalls int
registerErr error
heartbeatErrs []error
onHeartbeat func(int)
}
func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.NodePayload) error {
f.mu.Lock()
defer f.mu.Unlock()
f.registerCalls++
return f.registerErr
}
func (f *fakeHeartbeatService) Heartbeat(ctx context.Context, payload protocol.NodePayload) error {
f.mu.Lock()
f.heartbeatCalls++
callIndex := f.heartbeatCalls
var err error
if len(f.heartbeatErrs) >= callIndex {
err = f.heartbeatErrs[callIndex-1]
}
onHeartbeat := f.onHeartbeat
f.mu.Unlock()
if onHeartbeat != nil {
onHeartbeat(callIndex)
}
return err
}
type fakeSyncService struct {
mu sync.Mutex
startupErr error
syncOnceErr error
startupCalls int
syncOnceCalls int
onSyncOnceCall func(int)
}
func (f *fakeSyncService) SyncOnStartup(ctx context.Context) error {
f.mu.Lock()
defer f.mu.Unlock()
f.startupCalls++
return f.startupErr
}
func (f *fakeSyncService) SyncOnce(ctx context.Context) error {
f.mu.Lock()
f.syncOnceCalls++
callIndex := f.syncOnceCalls
callback := f.onSyncOnceCall
f.mu.Unlock()
if callback != nil {
callback(callIndex)
}
return f.syncOnceErr
}
func TestRunnerKeepsHeartbeatWhenStartupSyncFails(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
heartbeatService := &fakeHeartbeatService{
onHeartbeat: func(callCount int) {
if callCount >= 2 {
cancel()
}
},
}
syncService := &fakeSyncService{
startupErr: errors.New("当前没有激活版本，保持当前 Nginx 配置"),
}
runner := &Runner{
Config: &config.Config{
NodeName: "edge-01",
NodeIP: "10.0.0.8",
AgentVersion: "0.1.0",
NginxVersion: "1.25.5",
HeartbeatInterval: 10 * time.Millisecond,
SyncInterval: 20 * time.Millisecond,
},
StateStore: stateStore,
HeartbeatService: heartbeatService,
SyncService: syncService,
}
err := runner.Run(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation, got %v", err)
}
if heartbeatService.registerCalls != 1 {
t.Fatalf("expected 1 register call, got %d", heartbeatService.registerCalls)
}
if heartbeatService.heartbeatCalls < 2 {
t.Fatalf("expected heartbeat loop to continue, got %d heartbeat calls", heartbeatService.heartbeatCalls)
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.LastError != "当前没有激活版本，保持当前 Nginx 配置" {
t.Fatalf("expected startup sync error to be recorded, got %q", snapshot.LastError)
}
}
func TestRunnerDoesNotExitOnHeartbeatOrSyncError(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
heartbeatService := &fakeHeartbeatService{
registerErr: errors.New("register timeout"),
heartbeatErrs: []error{errors.New("heartbeat timeout")},
}
syncService := &fakeSyncService{
syncOnceErr: errors.New("nginx reload failed"),
onSyncOnceCall: func(callCount int) {
if callCount >= 1 {
cancel()
}
},
}
runner := &Runner{
Config: &config.Config{
NodeName: "edge-01",
NodeIP: "10.0.0.8",
AgentVersion: "0.1.0",
NginxVersion: "1.25.5",
HeartbeatInterval: 10 * time.Millisecond,
SyncInterval: 10 * time.Millisecond,
},
StateStore: stateStore,
HeartbeatService: heartbeatService,
SyncService: syncService,
}
err := runner.Run(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation, got %v", err)
}
if heartbeatService.registerCalls != 1 {
t.Fatalf("expected register attempt, got %d", heartbeatService.registerCalls)
}
if syncService.syncOnceCalls == 0 {
t.Fatal("expected sync loop to continue after heartbeat/register errors")
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.LastError != "nginx reload failed" {
t.Fatalf("expected sync error to be recorded, got %q", snapshot.LastError)
}
}