dfb33681f9
- Add new gRPC endpoint for Heartbeat in the Yao service, enabling communication with the sandbox. - Update Makefile to include a dedicated unit test target for Sandbox V2, ensuring proper testing of new features. - Enhance CI workflows to incorporate Sandbox V2 tests, allowing for dual-mode testing (local and remote) with Docker. - Modify .gitignore to exclude specific Docker files while allowing shell scripts for Sandbox V2. - Update documentation in DESIGN.md to reflect the new architecture and capabilities of the Sandbox V2. These changes enhance the Yao SDK's functionality, providing improved support for sandbox operations and testing.
76 lines
1.9 KiB
Go
76 lines
1.9 KiB
Go
package sandbox
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/yaoapp/kun/log"
|
|
"github.com/yaoapp/yao/grpc/pb"
|
|
)
|
|
|
|
// HeartbeatData holds the latest heartbeat from a sandbox container.
|
|
type HeartbeatData struct {
|
|
SandboxID string
|
|
CPUPercent int32
|
|
MemBytes int64
|
|
RunningProcs int32
|
|
LastSeen time.Time
|
|
}
|
|
|
|
// Handler implements sandbox-related gRPC methods.
|
|
type Handler struct {
|
|
mu sync.RWMutex
|
|
heartbeats map[string]*HeartbeatData
|
|
onBeat func(data *HeartbeatData) string // optional callback; returns action
|
|
}
|
|
|
|
// NewHandler creates a Handler. onBeat is called on each heartbeat and
|
|
// may return "ok" or "shutdown" to signal the container.
|
|
func NewHandler(onBeat func(data *HeartbeatData) string) *Handler {
|
|
return &Handler{
|
|
heartbeats: make(map[string]*HeartbeatData),
|
|
onBeat: onBeat,
|
|
}
|
|
}
|
|
|
|
// Heartbeat handles the Heartbeat RPC from sandbox containers.
|
|
func (h *Handler) Heartbeat(_ context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) {
|
|
data := &HeartbeatData{
|
|
SandboxID: req.SandboxId,
|
|
CPUPercent: req.CpuPercent,
|
|
MemBytes: req.MemBytes,
|
|
RunningProcs: req.RunningProcs,
|
|
LastSeen: time.Now(),
|
|
}
|
|
|
|
h.mu.Lock()
|
|
h.heartbeats[req.SandboxId] = data
|
|
h.mu.Unlock()
|
|
|
|
action := "ok"
|
|
if h.onBeat != nil {
|
|
if a := h.onBeat(data); a != "" {
|
|
action = a
|
|
}
|
|
}
|
|
|
|
log.Trace("sandbox heartbeat: id=%s cpu=%d%% mem=%d procs=%d → %s",
|
|
req.SandboxId, req.CpuPercent, req.MemBytes, req.RunningProcs, action)
|
|
|
|
return &pb.HeartbeatResponse{Action: action}, nil
|
|
}
|
|
|
|
// LastHeartbeat returns the most recent heartbeat for a sandbox, or nil.
|
|
func (h *Handler) LastHeartbeat(sandboxID string) *HeartbeatData {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
return h.heartbeats[sandboxID]
|
|
}
|
|
|
|
// RemoveHeartbeat cleans up heartbeat data for a removed sandbox.
|
|
func (h *Handler) RemoveHeartbeat(sandboxID string) {
|
|
h.mu.Lock()
|
|
delete(h.heartbeats, sandboxID)
|
|
h.mu.Unlock()
|
|
}
|