Import changes.

  - e9266a2f7bca606f2c5c73d4529fa59d7676f927

GitOrigin-RevId: e9266a2f7bca606f2c5c73d4529fa59d7676f927

Author: Aalyria Technologies, Inc

agent/internal/agentcli/agentcli.go

@@ -32,6 +32,7 @@ import (
 	"os"
 	"os/signal"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/jonboulle/clockwork"
@@ -44,6 +45,7 @@ import (
 	otelsdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.34.0"
 	oteltracenoop "go.opentelemetry.io/otel/trace/noop"
+
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/backoff"
@@ -57,6 +59,7 @@ import (
 	agent "aalyria.com/spacetime/agent"
 	"aalyria.com/spacetime/agent/enactment"
 	enact_extproc "aalyria.com/spacetime/agent/enactment/extproc"
+
 	"aalyria.com/spacetime/agent/internal/configpb"
 	"aalyria.com/spacetime/agent/internal/protofmt"
 	"aalyria.com/spacetime/agent/internal/snmp"
@@ -175,7 +178,7 @@ func (ac AgentConf) Run(ctx context.Context, args []string) (err error) {
 	}
 	defer shutdownTracer()
 
-	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
+	ctx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM)
 	defer stop()
 
 	g, ctx := errgroup.WithContext(ctx)
@@ -197,7 +200,17 @@ func (ac AgentConf) Run(ctx context.Context, args []string) (err error) {
 		}
 		return nil
 	})
-	return g.Wait()
+
+	// Wait for all goroutines to complete
+	err = g.Wait()
+	// Treat context cancellation (from Ctrl+C or SIGTERM) as normal shutdown
+	// TODO: Check whether this is intended behaviour
+	if errors.Is(err, context.Canceled) {
+		log := zerolog.Ctx(ctx)
+		log.Info().Msg("agent shutdown complete")
+		return nil
+	}
+	return err
 }
 
 const defaultMinConnectTimeout = 20 * time.Second
@@ -492,6 +505,7 @@ enactmentSwitch:
 
 	case *configpb.SdnAgent_EnactmentDriver_Snmp:
 		fmt.Print("TODO: SdnAgent_EnactmentDriver_Snmp")
+
 	case *configpb.SdnAgent_EnactmentDriver_Dynamic:
 		dialOpts, err := getDialOpts(ctx, node.EnactmentDriver.GetConnectionParams(), clock)
 		if err != nil {

agent/internal/configpb/config.proto

@@ -221,6 +221,7 @@ message SnmpObject {
 
 // An SDN agent, controlling any number of commandable elements.
 message SdnAgent {
+
   // An external command is a delegate process that handles implementing some
   // aspect of the SBI contract. Exit codes within the range of the canonical
   // gRPC status codes will be transformed (alongside any messages written to
@@ -276,6 +277,7 @@ message SdnAgent {
   // E2 enactment driver configuration (currently empty, for future use).
   message E2Enactment {
   }
+  
 
   // An EnactmentDriver is responsible for processing incoming enactments for
   // a given node.
@@ -293,6 +295,7 @@ message SdnAgent {
       // SNMPv3 for enactments
       SnmpEnactment snmp = 5;
 
+
       // Use an agent-specific driver to process enactments.
       google.protobuf.Any dynamic = 4;
 
@@ -327,7 +330,7 @@ message SdnAgent {
     // How often the metrics will be gathered and reported.
     google.protobuf.Duration collection_period = 3;
 
-    // The port on which the metrics will be scrapable as Prometheus metrics
+    // The port on which the metrics will be scrapable as Prometheus metrics.
     int32 prometheus_port = 4;
   }
 
@@ -363,6 +366,7 @@ message SdnAgent {
       // Use SNMPv3 GET calls to gather telemetry.
       SnmpTelemetry snmp = 5;
 
+
       // Use an agent-specific driver to gather telemetry.
       google.protobuf.Any dynamic = 4;
 

agent/node_controller.go

@@ -80,6 +80,11 @@ func (a *Agent) newNodeController(node *node, done func()) (*nodeController, err
 		}
 		nc.closers = append(nc.closers, telemetryConn.Close)
 
+		// If the telemetry driver implements io.Closer, register it for cleanup
+		if closer, ok := node.td.(interface{ Close() error }); ok {
+			nc.closers = append(nc.closers, closer.Close)
+		}
+
 		telemetryClient := telemetrypb.NewTelemetryClient(telemetryConn)
 
 		ts := nc.newTelemetryService(telemetryClient, node.td)

agent/telemetry/periodic_driver.go

@@ -84,3 +84,12 @@ func (pd *PeriodicDriver) Run(
 		}
 	}
 }
+
+// Close shuts down the PeriodicDriver. If the underlying ReportGenerator
+// implements Close() error, it will be called for cleanup.
+func (pd *PeriodicDriver) Close() error {
+	if closer, ok := pd.ReportGenerator.(interface{ Close() error }); ok {
+		return closer.Close()
+	}
+	return nil
+}