yao/tai/heartbeat.go
Max 83ebe49036 refactor: unify server lifecycle, migrate gRPC client, and clean up sandbox v2
Server lifecycle:
- Introduce service.Service to manage HTTP + gRPC startup/shutdown
- Fix gRPC mutex deadlock in StartServer when port is occupied
- Add GracefulStop with 5s timeout before forced Stop in grpc.go
- Pre-check HTTP and gRPC port availability in cmd/start.go
- Print gRPC server address in startup access-points block

gRPC client refactor:
- Move token manager and client from tai/grpc/ to grpc/client/
- Add backward-compatible aliases in tai/yao.go and tai/token.go
- Update cmd/run.go to import grpc/client directly (no tai dependency)

Sandbox v2 docker migration:
- Delete sandbox/v2/docker/ (moved to tai repo)
- Update sandbox/docker/build.sh hint to point to tai repo
- Clean up .gitignore entries for removed docker directory
- Temporarily disable SandboxV2Test and BenchmarkSandboxV2 in CI
  (docker images need rebuild after tai repo migration)

Tai integration:
- Add direct-mode registration API handlers in tai/api/
- Add heartbeat handler and token management wrappers
- Update tai/registry and tai/tunnel for latest protocol
- Replace yao-grpc references with tai call in docs

Made-with: Cursor
2026-03-07 17:19:19 +08:00

72 lines
1.5 KiB
Go

package tai
import (
"context"
"fmt"
"os"
"os/exec"
"runtime"
"strconv"
"strings"
"time"
)
const defaultHeartbeatInterval = 10 * time.Second
// HeartbeatLoop sends periodic heartbeats to the Yao gRPC server.
// It runs until ctx is cancelled.
func HeartbeatLoop(ctx context.Context, client *YaoClient, sandboxID string) {
interval := defaultHeartbeatInterval
if s := os.Getenv("YAO_HEARTBEAT_INTERVAL"); s != "" {
if d, err := time.ParseDuration(s); err == nil && d > 0 {
interval = d
}
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
cpu, mem := sampleResources()
procs := countUserProcesses()
action, err := client.Heartbeat(ctx, sandboxID, cpu, mem, procs)
if err != nil {
continue
}
if action == "shutdown" {
fmt.Fprintf(os.Stderr, "tai: received shutdown signal\n")
p, _ := os.FindProcess(os.Getpid())
p.Signal(os.Interrupt)
return
}
}
}
}
func countUserProcesses() int32 {
if runtime.GOOS != "linux" {
return 0
}
out, err := exec.Command("sh", "-c", "ps -e --no-headers | wc -l").Output()
if err != nil {
return 0
}
n, _ := strconv.Atoi(strings.TrimSpace(string(out)))
return int32(n)
}
func sampleResources() (cpuPercent int32, memBytes int64) {
if runtime.GOOS != "linux" {
return 0, 0
}
data, err := os.ReadFile("/sys/fs/cgroup/memory.current")
if err == nil {
mem, _ := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64)
memBytes = mem
}
return 0, memBytes
}