toki
b4c6550eab
- Add node store implementation for edge app - Add adapters factory for node app - Update edge and node transport layers - Update domain rules for edge and node - Add bin scripts for edge and node - Update configs and documentation - Add agent-task node_centralized_mgmt directory
104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
package transport
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
toki "git.toki-labs.com/toki/common-proto-socket/go"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
// Handler processes IOP messages received from edge.
|
|
type Handler interface {
|
|
OnRunRequest(ctx context.Context, sess *Session, req *iop.RunRequest) error
|
|
OnCancel(ctx context.Context, sess *Session, req *iop.CancelRequest) error
|
|
}
|
|
|
|
// Session represents the node's persistent connection to edge.
|
|
type Session struct {
|
|
client *toki.TcpClient
|
|
logger *zap.Logger
|
|
mu sync.RWMutex
|
|
handler Handler
|
|
cancelFns sync.Map
|
|
}
|
|
|
|
func newSession(client *toki.TcpClient, logger *zap.Logger) *Session {
|
|
s := &Session{client: client, logger: logger}
|
|
|
|
toki.AddListenerTyped[*iop.RunRequest](&client.Communicator, func(req *iop.RunRequest) {
|
|
go func() {
|
|
s.mu.RLock()
|
|
h := s.handler
|
|
s.mu.RUnlock()
|
|
if h == nil {
|
|
return
|
|
}
|
|
if err := h.OnRunRequest(context.Background(), s, req); err != nil {
|
|
logger.Warn("run request error",
|
|
zap.String("run_id", req.GetRunId()),
|
|
zap.Error(err),
|
|
)
|
|
}
|
|
}()
|
|
})
|
|
|
|
toki.AddListenerTyped[*iop.CancelRequest](&client.Communicator, func(req *iop.CancelRequest) {
|
|
s.mu.RLock()
|
|
h := s.handler
|
|
s.mu.RUnlock()
|
|
if h == nil {
|
|
return
|
|
}
|
|
if err := h.OnCancel(context.Background(), s, req); err != nil {
|
|
logger.Warn("cancel error",
|
|
zap.String("run_id", req.GetRunId()),
|
|
zap.Error(err),
|
|
)
|
|
}
|
|
})
|
|
|
|
client.AddDisconnectListener(func(_ *toki.TcpClient) {
|
|
logger.Info("disconnected from edge")
|
|
})
|
|
|
|
return s
|
|
}
|
|
|
|
// SetHandler attaches the message handler. Called after registration completes.
|
|
func (s *Session) SetHandler(h Handler) {
|
|
s.mu.Lock()
|
|
s.handler = h
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
// Send transmits a proto message to edge.
|
|
func (s *Session) Send(m proto.Message) error {
|
|
return s.client.Send(m)
|
|
}
|
|
|
|
// IsAlive reports whether the connection is active.
|
|
func (s *Session) IsAlive() bool { return s.client.IsAlive() }
|
|
|
|
// Close terminates the connection to edge.
|
|
func (s *Session) Close() error { return s.client.Close() }
|
|
|
|
// RegisterCancel stores a cancel function for a running execution.
|
|
func (s *Session) RegisterCancel(runID string, cancel context.CancelFunc) {
|
|
s.cancelFns.Store(runID, cancel)
|
|
}
|
|
|
|
// DeregisterCancel removes a run's cancel function after completion.
|
|
func (s *Session) DeregisterCancel(runID string) {
|
|
s.cancelFns.Delete(runID)
|
|
}
|
|
|
|
// CancelRun invokes the cancel function for a run if one is registered.
|
|
func (s *Session) CancelRun(runID string) {
|
|
if v, ok := s.cancelFns.Load(runID); ok {
|
|
v.(context.CancelFunc)()
|
|
}
|
|
}
|