From fa6e769624563e864c65b48a51b01ef1ece32854 Mon Sep 17 00:00:00 2001 From: Joseph Anttila Hall Date: Mon, 11 Mar 2024 21:33:03 -0700 Subject: [PATCH] Agent drain: implement first half. - Agent sends drain signal to Server. - Server only logs it. --- cmd/agent/app/server.go | 37 +++++++++++++++++++++++++++++++------ pkg/agent/client.go | 15 +++++++++++++++ pkg/agent/clientset.go | 5 ++++- pkg/server/server.go | 2 ++ tests/framework/agent.go | 26 ++++++++++++++++++++++++-- 5 files changed, 76 insertions(+), 9 deletions(-) diff --git a/cmd/agent/app/server.go b/cmd/agent/app/server.go index 8cc3cfd0e..7de11c052 100644 --- a/cmd/agent/app/server.go +++ b/cmd/agent/app/server.go @@ -24,10 +24,13 @@ import ( "net" "net/http" "net/http/pprof" + "os" + "os/signal" "runtime" runpprof "runtime/pprof" "strconv" "strings" + "syscall" "time" "github.com/prometheus/client_golang/prometheus/promhttp" @@ -49,8 +52,8 @@ func NewAgentCommand(a *Agent, o *options.GrpcProxyAgentOptions) *cobra.Command Use: "agent", Long: `A gRPC agent, Connects to the proxy and then allows traffic to be forwarded to it.`, RunE: func(cmd *cobra.Command, args []string) error { - stopCh := make(chan struct{}) - return a.Run(o, stopCh) + drainCh, stopCh := SetupSignalHandler() + return a.Run(o, drainCh, stopCh) }, } @@ -64,13 +67,13 @@ type Agent struct { cs *agent.ClientSet } -func (a *Agent) Run(o *options.GrpcProxyAgentOptions, stopCh <-chan struct{}) error { +func (a *Agent) Run(o *options.GrpcProxyAgentOptions, drainCh, stopCh <-chan struct{}) error { o.Print() if err := o.Validate(); err != nil { return fmt.Errorf("failed to validate agent options with %v", err) } - cs, err := a.runProxyConnection(o, stopCh) + cs, err := a.runProxyConnection(o, drainCh, stopCh) if err != nil { return fmt.Errorf("failed to run proxy connection with %v", err) } @@ -92,7 +95,29 @@ func (a *Agent) Run(o *options.GrpcProxyAgentOptions, stopCh <-chan struct{}) er return nil } -func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, stopCh <-chan struct{}) (*agent.ClientSet, error) { +var shutdownSignals = []os.Signal{os.Interrupt, syscall.SIGTERM} + +func SetupSignalHandler() (drainCh, stopCh <-chan struct{}) { + drain := make(chan struct{}) + stop := make(chan struct{}) + c := make(chan os.Signal, 2) + signal.Notify(c, shutdownSignals...) + labels := runpprof.Labels( + "core", "signalHandler", + ) + go runpprof.Do(context.Background(), labels, func(context.Context) { handleSignals(c, drain, stop) }) + + return drain, stop +} + +func handleSignals(signalCh chan os.Signal, drainCh, stopCh chan struct{}) { + <-signalCh + close(drainCh) + <-signalCh + close(stopCh) +} + +func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, drainCh, stopCh <-chan struct{}) (*agent.ClientSet, error) { var tlsConfig *tls.Config var err error if tlsConfig, err = util.GetClientTLSConfig(o.CaCert, o.AgentCert, o.AgentKey, o.ProxyServerHost, o.AlpnProtos); err != nil { @@ -106,7 +131,7 @@ func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, stopCh <-ch }), } cc := o.ClientSetConfig(dialOptions...) - cs := cc.NewAgentClientSet(stopCh) + cs := cc.NewAgentClientSet(drainCh, stopCh) cs.Serve() return cs, nil diff --git a/pkg/agent/client.go b/pkg/agent/client.go index 16b1df244..f40464ea1 100644 --- a/pkg/agent/client.go +++ b/pkg/agent/client.go @@ -137,7 +137,10 @@ type Client struct { address string opts []grpc.DialOption conn *grpc.ClientConn + + drainCh <-chan struct{} stopCh chan struct{} + // locks sendLock sync.Mutex recvLock sync.Mutex @@ -158,6 +161,7 @@ func newAgentClient(address, agentID, agentIdentifiers string, cs *ClientSet, op agentIdentifiers: agentIdentifiers, opts: opts, probeInterval: cs.probeInterval, + drainCh: cs.drainCh, stopCh: make(chan struct{}), serviceAccountTokenPath: cs.serviceAccountTokenPath, connManager: newConnectionManager(), @@ -325,6 +329,17 @@ func (a *Client) Serve() { case <-a.stopCh: klog.V(2).InfoS("stop agent client.") return + case <-a.drainCh: + klog.V(2).InfoS("drain agent client.") + drainPkt := &client.Packet{ + Type: client.PacketType_DRAIN, + Payload: &client.Packet_Drain{ + Drain: &client.Drain{}, + }, + } + if err := a.Send(drainPkt); err != nil { + klog.ErrorS(err, "drain failure", "") + } default: } diff --git a/pkg/agent/clientset.go b/pkg/agent/clientset.go index c5adcf4f7..6a2510a84 100644 --- a/pkg/agent/clientset.go +++ b/pkg/agent/clientset.go @@ -52,6 +52,8 @@ type ClientSet struct { dialOptions []grpc.DialOption // file path contains service account token serviceAccountTokenPath string + // channel to signal that the agent is pending termination. + drainCh <-chan struct{} // channel to signal shutting down the client set. Primarily for test. stopCh <-chan struct{} @@ -141,7 +143,7 @@ type ClientSetConfig struct { SyncForever bool } -func (cc *ClientSetConfig) NewAgentClientSet(stopCh <-chan struct{}) *ClientSet { +func (cc *ClientSetConfig) NewAgentClientSet(drainCh, stopCh <-chan struct{}) *ClientSet { return &ClientSet{ clients: make(map[string]*Client), agentID: cc.AgentID, @@ -154,6 +156,7 @@ func (cc *ClientSetConfig) NewAgentClientSet(stopCh <-chan struct{}) *ClientSet serviceAccountTokenPath: cc.ServiceAccountTokenPath, warnOnChannelLimit: cc.WarnOnChannelLimit, syncForever: cc.SyncForever, + drainCh: drainCh, stopCh: stopCh, } } diff --git a/pkg/server/server.go b/pkg/server/server.go index d65d5c4be..b3f883c93 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -990,6 +990,8 @@ func (s *ProxyServer) serveRecvBackend(backend Backend, agentID string, recvCh < klog.V(5).InfoS("CLOSE_RSP sent to frontend", "connectionID", resp.ConnectID) } + case client.PacketType_DRAIN: + klog.V(2).InfoS("agent is draining", "agentID", agentID) default: klog.V(5).InfoS("Ignoring unrecognized packet from backend", "packet", pkt, "agentID", agentID) } diff --git a/tests/framework/agent.go b/tests/framework/agent.go index 15b9ae2e2..6d446fca6 100644 --- a/tests/framework/agent.go +++ b/tests/framework/agent.go @@ -27,6 +27,7 @@ import ( "path/filepath" "strconv" "sync" + "syscall" "testing" "time" @@ -53,6 +54,7 @@ type AgentRunner interface { type Agent interface { GetConnectedServerCount() (int, error) Ready() bool + Drain() Stop() Metrics() metricstest.AgentTester } @@ -66,9 +68,10 @@ func (*InProcessAgentRunner) Start(t testing.TB, opts AgentOpts) (Agent, error) } ctx, cancel := context.WithCancel(context.Background()) + drainCh := make(chan struct{}) stopCh := make(chan struct{}) go func() { - if err := a.Run(o, stopCh); err != nil { + if err := a.Run(o, drainCh, stopCh); err != nil { log.Printf("ERROR running agent: %v", err) cancel() } @@ -84,6 +87,7 @@ func (*InProcessAgentRunner) Start(t testing.TB, opts AgentOpts) (Agent, error) pa := &inProcessAgent{ client: a.ClientSet(), + drainCh: drainCh, stopCh: stopCh, healthAddr: healthAddr, } @@ -94,12 +98,21 @@ func (*InProcessAgentRunner) Start(t testing.TB, opts AgentOpts) (Agent, error) type inProcessAgent struct { client *agent.ClientSet + drainOnce sync.Once + drainCh chan struct{} + stopOnce sync.Once stopCh chan struct{} healthAddr string } +func (a *inProcessAgent) Drain() { + a.drainOnce.Do(func() { + close(a.drainCh) + }) +} + func (a *inProcessAgent) Stop() { a.stopOnce.Do(func() { close(a.stopCh) @@ -160,7 +173,16 @@ type externalAgent struct { cmd *exec.Cmd metrics *metricstest.Tester - stopOnce sync.Once + drainOnce sync.Once + stopOnce sync.Once +} + +func (a *externalAgent) Drain() { + a.drainOnce.Do(func() { + if err := a.cmd.Process.Signal(syscall.SIGTERM); err != nil { + log.Fatalf("Error draining agent process: %v", err) + } + }) } func (a *externalAgent) Stop() {