fix: [NPM] close telemetry handler before crashing (#3333)

fix: close telemetry handler before crashing

Signed-off-by: Hunter Gregory <[email protected]>

Author: huntergregory

npm/cmd/start.go

@@ -57,7 +57,10 @@ func newStartNPMCmd() *cobra.Command {
 				KubeConfigPath: viper.GetString(flagKubeConfigPath),
 			}
 
-			return start(*config, flags)
+			// start is blocking, unless there's an error
+			err = start(*config, flags)
+			metrics.Close()
+			return err
 		},
 	}
 

npm/metrics/ai-utils.go

@@ -11,6 +11,8 @@ import (
 	"k8s.io/klog"
 )
 
+const telemetryCloseWaitTimeSeconds = 10
+
 var (
 	th         aitelemetry.TelemetryHandle
 	npmVersion int
@@ -54,6 +56,15 @@ func CreateTelemetryHandle(npmVersionNum int, imageVersion, aiMetadata string) e
 	return nil
 }
 
+// Close cleans up the telemetry handle, which effectively waits for all telemetry data to be sent
+func Close() {
+	if th == nil {
+		return
+	}
+
+	th.Close(telemetryCloseWaitTimeSeconds)
+}
+
 // SendErrorLogAndMetric sends a metric through AI telemetry and sends a log to the Kusto Messages table
 func SendErrorLogAndMetric(operationID int, format string, args ...interface{}) {
 	// Send error metrics

npm/pkg/dataplane/ipsets/ipsetmanager_linux.go

@@ -422,6 +422,7 @@ func (iMgr *IPSetManager) applyIPSets() error {
 			msg := fmt.Sprintf("exceeded max consecutive failures (%d) when applying ipsets. final error: %s", maxConsecutiveFailures, restoreError.Error())
 			klog.Error(msg)
 			metrics.SendErrorLogAndMetric(util.IpsmID, msg)
+			metrics.Close()
 			panic(msg)
 		}