toki
fc16a1e05c
- Implement stream.go with sideband passthrough support - Update chat_handler.go with streaming handler changes - Add run_dispatch.go with new dispatch service logic - Add server tests for streaming and sideband functionality - Archive old plan/code-review docs for cloud-G07
447 lines
13 KiB
Go
447 lines
13 KiB
Go
package transport
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"strconv"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
toki "git.toki-labs.com/toki/proto-socket/go"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"iop/apps/edge/internal/configrefresh"
|
|
edgenode "iop/apps/edge/internal/node"
|
|
"iop/packages/go/events"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
const (
|
|
heartbeatIntervalSec = 30
|
|
// heartbeatWaitSec mirrors the node side. See the comment in
|
|
// apps/node/internal/transport/client.go for rationale.
|
|
heartbeatWaitSec = 45
|
|
)
|
|
|
|
const configRefreshTimeoutSec = 30
|
|
|
|
func edgeParserMap() toki.ParserMap {
|
|
return toki.ParserMap{
|
|
toki.TypeNameOf(&iop.RunEvent{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.RunEvent{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.EdgeNodeEvent{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.EdgeNodeEvent{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.ProviderTunnelFrame{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.ProviderTunnelFrame{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.RegisterRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.RegisterRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandResponse{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.NodeConfigRefreshResponse{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeConfigRefreshResponse{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
}
|
|
|
|
// Server wraps proto-socket TcpServer and manages node connections.
|
|
type Server struct {
|
|
tcp *toki.TcpServer
|
|
listen string
|
|
registry *edgenode.Registry
|
|
nodeStoreMu sync.RWMutex
|
|
nodeStore *edgenode.NodeStore
|
|
logger *zap.Logger
|
|
handlerMu sync.RWMutex
|
|
onRunEvent func(*iop.RunEvent)
|
|
onNodeEvent func(*iop.EdgeNodeEvent)
|
|
onTunnelFrame func(*iop.ProviderTunnelFrame)
|
|
stopping atomic.Bool
|
|
HeartbeatInterval int
|
|
HeartbeatWait int
|
|
}
|
|
|
|
func NewServer(listen string, registry *edgenode.Registry, nodeStore *edgenode.NodeStore, logger *zap.Logger) (*Server, error) {
|
|
host, portStr, err := net.SplitHostPort(listen)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
port, err := strconv.Atoi(portStr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
s := &Server{
|
|
listen: listen,
|
|
registry: registry,
|
|
nodeStore: nodeStore,
|
|
logger: logger,
|
|
HeartbeatInterval: heartbeatIntervalSec,
|
|
HeartbeatWait: heartbeatWaitSec,
|
|
}
|
|
s.tcp = toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
|
|
return toki.NewTcpClient(conn, s.HeartbeatInterval, s.HeartbeatWait, edgeParserMap())
|
|
})
|
|
s.tcp.OnClientConnected = s.onNodeConnected
|
|
return s, nil
|
|
}
|
|
|
|
func (s *Server) Start(ctx context.Context) error {
|
|
s.stopping.Store(false)
|
|
if err := s.tcp.Start(ctx); err != nil {
|
|
return err
|
|
}
|
|
s.logger.Info("edge listening for nodes", zap.String("addr", s.listen))
|
|
return nil
|
|
}
|
|
|
|
func (s *Server) SetNodeStore(store *edgenode.NodeStore) {
|
|
s.nodeStoreMu.Lock()
|
|
s.nodeStore = store
|
|
s.nodeStoreMu.Unlock()
|
|
}
|
|
|
|
func (s *Server) nodeStoreSnapshot() *edgenode.NodeStore {
|
|
s.nodeStoreMu.RLock()
|
|
defer s.nodeStoreMu.RUnlock()
|
|
return s.nodeStore
|
|
}
|
|
|
|
func (s *Server) Stop() error {
|
|
s.stopping.Store(true)
|
|
return s.tcp.Stop()
|
|
}
|
|
|
|
func (s *Server) SetRunEventHandler(handler func(*iop.RunEvent)) {
|
|
s.handlerMu.Lock()
|
|
s.onRunEvent = handler
|
|
s.handlerMu.Unlock()
|
|
}
|
|
|
|
func (s *Server) SetNodeEventHandler(handler func(*iop.EdgeNodeEvent)) {
|
|
s.handlerMu.Lock()
|
|
s.onNodeEvent = handler
|
|
s.handlerMu.Unlock()
|
|
}
|
|
|
|
// SetTunnelFrameHandler registers the request-bound ProviderTunnelFrame
|
|
// handler. Tunnel frames carry raw provider passthrough bytes and are routed
|
|
// to a per-request channel by the handler; they must never be published to
|
|
// the run event bus.
|
|
func (s *Server) SetTunnelFrameHandler(handler func(*iop.ProviderTunnelFrame)) {
|
|
s.handlerMu.Lock()
|
|
s.onTunnelFrame = handler
|
|
s.handlerMu.Unlock()
|
|
}
|
|
|
|
func (s *Server) HasTunnelFrameHandler() bool {
|
|
s.handlerMu.RLock()
|
|
defer s.handlerMu.RUnlock()
|
|
return s.onTunnelFrame != nil
|
|
}
|
|
|
|
func (s *Server) HasRunEventHandler() bool {
|
|
s.handlerMu.RLock()
|
|
defer s.handlerMu.RUnlock()
|
|
return s.onRunEvent != nil
|
|
}
|
|
|
|
func (s *Server) HasNodeEventHandler() bool {
|
|
s.handlerMu.RLock()
|
|
defer s.handlerMu.RUnlock()
|
|
return s.onNodeEvent != nil
|
|
}
|
|
|
|
func (s *Server) onNodeConnected(client *toki.TcpClient) {
|
|
s.logger.Info("node connection established")
|
|
|
|
toki.AddListenerTyped[*iop.RunEvent](&client.Communicator, func(e *iop.RunEvent) {
|
|
s.logger.Debug("run event received",
|
|
zap.String("run_id", e.GetRunId()),
|
|
zap.String("type", e.GetType()),
|
|
)
|
|
s.enrichRunEvent(e)
|
|
s.handlerMu.RLock()
|
|
handler := s.onRunEvent
|
|
s.handlerMu.RUnlock()
|
|
if handler != nil {
|
|
handler(e)
|
|
}
|
|
})
|
|
|
|
toki.AddListenerTyped[*iop.ProviderTunnelFrame](&client.Communicator, func(f *iop.ProviderTunnelFrame) {
|
|
s.handlerMu.RLock()
|
|
handler := s.onTunnelFrame
|
|
s.handlerMu.RUnlock()
|
|
if handler == nil {
|
|
s.logger.Warn("provider tunnel frame dropped: no tunnel handler",
|
|
zap.String("run_id", f.GetRunId()),
|
|
zap.String("tunnel_id", f.GetTunnelId()),
|
|
)
|
|
return
|
|
}
|
|
handler(f)
|
|
})
|
|
|
|
toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
|
|
&client.Communicator,
|
|
func(req *iop.RegisterRequest) (*iop.RegisterResponse, error) {
|
|
nodeStore := s.nodeStoreSnapshot()
|
|
if nodeStore == nil {
|
|
return &iop.RegisterResponse{Accepted: false, Reason: "edge node store unavailable"}, nil
|
|
}
|
|
rec, ok := nodeStore.FindByToken(req.GetToken())
|
|
if !ok {
|
|
s.logger.Warn("unknown token", zap.String("token_prefix", safePrefix(req.GetToken())))
|
|
s.emitNodeEvent(events.NewEdgeNodeEvent(
|
|
events.SourceEdge,
|
|
events.TypeNodeRegistrationFailed,
|
|
"",
|
|
"",
|
|
events.ReasonUnknownToken,
|
|
map[string]string{
|
|
events.MetadataFailureReason: events.ReasonUnknownToken,
|
|
events.MetadataTokenPrefix: safePrefix(req.GetToken()),
|
|
},
|
|
))
|
|
return &iop.RegisterResponse{Accepted: false, Reason: "unknown token"}, nil
|
|
}
|
|
|
|
cfg, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
s.logger.Error("build config payload failed",
|
|
zap.String("node_id", rec.ID),
|
|
zap.Error(err),
|
|
)
|
|
return &iop.RegisterResponse{Accepted: false, Reason: "internal config error"}, nil
|
|
}
|
|
|
|
entry := &edgenode.NodeEntry{
|
|
NodeID: rec.ID,
|
|
Alias: rec.Alias,
|
|
AgentKind: rec.AgentKind,
|
|
LifecycleState: edgenode.LifecycleConnected,
|
|
Client: client,
|
|
Index: rec.Index,
|
|
HasIndex: true,
|
|
}
|
|
toki.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(e *iop.EdgeNodeEvent) {
|
|
if e.NodeId == "" {
|
|
e.NodeId = rec.ID
|
|
}
|
|
if e.Alias == "" {
|
|
e.Alias = rec.Alias
|
|
}
|
|
s.emitNodeEvent(e)
|
|
})
|
|
client.AddDisconnectListener(func(_ *toki.TcpClient) {
|
|
transportInfo := client.DisconnectInfo()
|
|
fields := []zap.Field{zap.String("node_id", rec.ID)}
|
|
fields = append(fields, transportDisconnectFields(transportInfo)...)
|
|
s.logger.Info("node unregistered", fields...)
|
|
reason := events.ReasonTransportClosed
|
|
if transportInfo.Reason == "heartbeat_timeout" {
|
|
reason = events.ReasonHeartbeatTimeout
|
|
s.registry.UpdateLifecycle(rec.ID, edgenode.LifecycleFailed)
|
|
} else if s.stopping.Load() {
|
|
reason = events.ReasonEdgeShutdown
|
|
}
|
|
|
|
meta := transportDisconnectMetadata(transportInfo)
|
|
if meta == nil {
|
|
meta = make(map[string]string)
|
|
}
|
|
if reason == events.ReasonHeartbeatTimeout {
|
|
meta[events.MetadataFailureReason] = events.ReasonHeartbeatTimeout
|
|
meta[events.MetadataLifecycleState] = edgenode.LifecycleFailed
|
|
} else {
|
|
meta[events.MetadataLifecycleState] = edgenode.LifecycleFailed
|
|
}
|
|
meta[events.MetadataAgentKind] = rec.AgentKind
|
|
|
|
s.emitNodeEvent(events.NewEdgeNodeEvent(
|
|
events.SourceEdge,
|
|
events.TypeNodeDisconnected,
|
|
rec.ID,
|
|
rec.Alias,
|
|
reason,
|
|
meta,
|
|
))
|
|
s.registry.UnregisterIfClient(rec.ID, client)
|
|
})
|
|
if !s.registry.RegisterIfAbsent(entry) {
|
|
reason := "node already connected"
|
|
s.logger.Warn("duplicate registration rejected",
|
|
zap.String("node_id", rec.ID),
|
|
zap.String("agent_kind", rec.AgentKind),
|
|
)
|
|
s.emitNodeEvent(events.NewEdgeNodeEvent(
|
|
events.SourceEdge,
|
|
events.TypeNodeRegistrationFailed,
|
|
rec.ID,
|
|
rec.Alias,
|
|
events.ReasonDuplicateConnection,
|
|
map[string]string{
|
|
events.MetadataFailureReason: events.ReasonDuplicateConnection,
|
|
events.MetadataAgentKind: rec.AgentKind,
|
|
},
|
|
))
|
|
return &iop.RegisterResponse{Accepted: false, Reason: reason}, nil
|
|
}
|
|
s.logger.Info("node registered",
|
|
zap.String("node_id", rec.ID),
|
|
zap.String("alias", rec.Alias),
|
|
)
|
|
s.emitNodeEvent(events.NewEdgeNodeEvent(
|
|
events.SourceEdge,
|
|
events.TypeNodeConnected,
|
|
rec.ID,
|
|
rec.Alias,
|
|
events.ReasonRegistered,
|
|
map[string]string{
|
|
events.MetadataAgentKind: rec.AgentKind,
|
|
events.MetadataLifecycleState: edgenode.LifecycleConnected,
|
|
},
|
|
))
|
|
|
|
return &iop.RegisterResponse{
|
|
Accepted: true,
|
|
NodeId: rec.ID,
|
|
Alias: rec.Alias,
|
|
Config: cfg,
|
|
}, nil
|
|
},
|
|
)
|
|
}
|
|
|
|
// PushConfigRefresh sends a NodeConfigRefreshRequest to every configured node
|
|
// and collects per-node results. Nodes that are not in the live registry or
|
|
// have no alive client are recorded as skipped. Connected nodes receive a
|
|
// node-specific payload built from the refreshed NodeStore record. Transport
|
|
// failures are recorded as failed and do not roll back the local config commit.
|
|
func (s *Server) PushConfigRefresh(ctx context.Context, req *iop.NodeConfigRefreshRequest) []configrefresh.NodeResult {
|
|
ns := s.nodeStoreSnapshot()
|
|
var records []*edgenode.NodeRecord
|
|
if ns != nil {
|
|
records = ns.All()
|
|
}
|
|
results := make([]configrefresh.NodeResult, 0, len(records))
|
|
timeout := time.Duration(configRefreshTimeoutSec) * time.Second
|
|
for _, rec := range records {
|
|
nr := configrefresh.NodeResult{
|
|
NodeID: rec.ID,
|
|
Alias: rec.Alias,
|
|
}
|
|
entry, ok := s.registry.Get(rec.ID)
|
|
if !ok || entry.Client == nil || !entry.Client.IsAlive() {
|
|
nr.Status = "skipped"
|
|
results = append(results, nr)
|
|
continue
|
|
}
|
|
nodeReq := &iop.NodeConfigRefreshRequest{
|
|
RequestId: req.GetRequestId(),
|
|
ChangedPaths: req.GetChangedPaths(),
|
|
}
|
|
if payload, err := edgenode.BuildConfigPayload(rec); err == nil {
|
|
nodeReq.Config = payload
|
|
}
|
|
resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
|
|
&entry.Client.Communicator,
|
|
nodeReq,
|
|
timeout,
|
|
)
|
|
if err != nil {
|
|
nr.Status = "failed"
|
|
nr.Error = err.Error()
|
|
results = append(results, nr)
|
|
continue
|
|
}
|
|
switch resp.GetStatus() {
|
|
case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED:
|
|
nr.Status = "applied"
|
|
case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_RESTART_REQUIRED:
|
|
nr.Status = "restart_required"
|
|
nr.RestartRequiredPaths = resp.GetRestartRequiredPaths()
|
|
case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED:
|
|
nr.Status = "failed"
|
|
nr.Error = resp.GetError()
|
|
case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_SKIPPED:
|
|
nr.Status = "skipped"
|
|
default:
|
|
nr.Status = "skipped"
|
|
}
|
|
results = append(results, nr)
|
|
}
|
|
return results
|
|
}
|
|
|
|
func (s *Server) enrichRunEvent(event *iop.RunEvent) {
|
|
if event == nil || event.GetNodeAlias() != "" || s.registry == nil {
|
|
return
|
|
}
|
|
nodeID := event.GetNodeId()
|
|
if nodeID == "" {
|
|
return
|
|
}
|
|
if entry, ok := s.registry.Get(nodeID); ok && entry.Alias != "" {
|
|
event.NodeAlias = entry.Alias
|
|
}
|
|
}
|
|
|
|
func (s *Server) emitNodeEvent(event *iop.EdgeNodeEvent) {
|
|
s.logger.Debug("node event received",
|
|
zap.String("node_id", event.GetNodeId()),
|
|
zap.String("type", event.GetType()),
|
|
zap.String("reason", event.GetReason()),
|
|
)
|
|
s.handlerMu.RLock()
|
|
handler := s.onNodeEvent
|
|
s.handlerMu.RUnlock()
|
|
if handler != nil {
|
|
handler(event)
|
|
}
|
|
}
|
|
|
|
func safePrefix(s string) string {
|
|
if len(s) > 8 {
|
|
return s[:8] + "..."
|
|
}
|
|
return s
|
|
}
|
|
|
|
func transportDisconnectMetadata(info toki.DisconnectInfo) map[string]string {
|
|
metadata := make(map[string]string, 2)
|
|
if info.Reason != "" {
|
|
metadata[events.MetadataTransportCloseReason] = info.Reason
|
|
}
|
|
if info.Error != "" {
|
|
metadata[events.MetadataTransportCloseError] = info.Error
|
|
}
|
|
if len(metadata) == 0 {
|
|
return nil
|
|
}
|
|
return metadata
|
|
}
|
|
|
|
func transportDisconnectFields(info toki.DisconnectInfo) []zap.Field {
|
|
fields := make([]zap.Field, 0, 2)
|
|
if info.Reason != "" {
|
|
fields = append(fields, zap.String("transport_close_reason", info.Reason))
|
|
}
|
|
if info.Error != "" {
|
|
fields = append(fields, zap.String("transport_close_error", info.Error))
|
|
}
|
|
return fields
|
|
}
|