add a possibility for NSE to Unregister itself on termination

Signed-off-by: Nikita Skrynnik <nikita.skrynnik@xored.com>
networkservicemesh · Nov 2, 2023 · e1b61a3 · e1b61a3
1 parent 653cabb
commit e1b61a3
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 7 deletions.
diff --git a/go.mod b/go.mod
@@ -18,6 +18,8 @@ require (
 	gopkg.in/yaml.v2 v2.4.0
 )
 
+replace github.com/networkservicemesh/sdk => ./local/sdk
+
 require (
 	cloud.google.com/go/compute/metadata v0.2.3 // indirect
 	github.com/OneOfOne/xxhash v1.2.8 // indirect

diff --git a/go.sum b/go.sum
@@ -116,8 +116,6 @@ github.com/miekg/dns v1.1.50 h1:DQUfb9uc6smULcREF09Uc+/Gd46YWqJd5DbpPE9xkcA=
 github.com/miekg/dns v1.1.50/go.mod h1:e3IlAVfNqAllflbibAZEWOXOQ+Ynzk/dDozDxY7XnME=
 github.com/networkservicemesh/api v1.11.1-0.20231031152153-730abd666985 h1:uiWxBSamqODTzO/uobmjjqWAhi93+bpIDQ+OpV3uU58=
 github.com/networkservicemesh/api v1.11.1-0.20231031152153-730abd666985/go.mod h1:E2yBac48+mMkMh6ODnsNyah4EE6rI08SMR9n+86Emxs=
-github.com/networkservicemesh/sdk v0.5.1-0.20231031174426-85ca154d11a8 h1:ufRUQ1zCkX+BuXjKSVQWgu56sT/M0da8akEr8MhYgk4=
-github.com/networkservicemesh/sdk v0.5.1-0.20231031174426-85ca154d11a8/go.mod h1:0SBXaWlnQMrsWmaw0DrZ5lBia1LXyA4oE2oagnOuhaY=
 github.com/networkservicemesh/sdk-sriov v0.0.0-20231031175612-0e655752bf94 h1:QkAeft///LEnoFNuMjnVhx6QT5bsojGydiOjaGi21RM=
 github.com/networkservicemesh/sdk-sriov v0.0.0-20231031175612-0e655752bf94/go.mod h1:SMUxTwlXORyeVDXRWImiqtTMXDtUxRrNgnoFOP3a3Rg=
 github.com/open-policy-agent/opa v0.44.0 h1:sEZthsrWBqIN+ShTMJ0Hcz6a3GkYsY4FaB2S/ou2hZk=

diff --git a/local/sdk b/local/sdk
diff --git a/main.go b/main.go
@@ -91,6 +91,7 @@ type Config struct {
 	CidrPrefix             cidr.Groups       `default:"169.254.0.0/16" desc:"List of CIDR Prefix to assign IPv4 and IPv6 addresses from" split_words:"true"`
 	IdleTimeout            time.Duration     `default:"0" desc:"timeout for automatic shutdown when there were no requests for specified time. Set 0 to disable auto-shutdown." split_words:"true"`
 	RegisterService        bool              `default:"true" desc:"if true then registers network service on startup" split_words:"true"`
+	UnregisterItself       bool              `default:"false" desc:"if true then NSE unregister itself when it completes working" split_words:"true"`
 	PBRConfigPath          string            `default:"/etc/policy-based-routing/config.yaml" desc:"Path to policy based routing config file" split_words:"true"`
 	LogLevel               string            `default:"INFO" desc:"Log level" split_words:"true"`
 	OpenTelemetryEndpoint  string            `default:"otel-collector.observability.svc.cluster.local:4317" desc:"OpenTelemetry Collector Endpoint"`
@@ -112,15 +113,17 @@ func main() {
 	// ********************************************************************************
 	// setup context to catch signals
 	// ********************************************************************************
-	ctx, cancel := signal.NotifyContext(
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	signalCtx, _ := signal.NotifyContext(
 		context.Background(),
 		os.Interrupt,
 		// More Linux signals here
 		syscall.SIGHUP,
 		syscall.SIGTERM,
 		syscall.SIGQUIT,
 	)
-	defer cancel()
 
 	// ********************************************************************************
 	// setup logging
@@ -288,7 +291,7 @@ func main() {
 			registrysendfd.NewNetworkServiceEndpointRegistryClient(),
 		),
 		registryclient.WithAuthorizeNSERegistryClient(registryauthorize.NewNetworkServiceEndpointRegistryClient(
-			registryauthorize.WithPolicies(config.RegistryClientPolicies...))),
+			registryauthorize.Any())),
 	)
 	nse := getNseEndpoint(config, listenOn)
 	nse, err = nseRegistryClient.Register(ctx, nse)
@@ -298,12 +301,20 @@ func main() {
 		log.FromContext(ctx).Fatalf("unable to register nse %+v", err)
 	}
 
+	if config.UnregisterItself {
+		defer func() {
+			_, err = nseRegistryClient.Unregister(context.Background(), nse)
+			if err != nil {
+				log.FromContext(ctx).Errorf("nse failed to unregister itself on termination: %s", err.Error())
+			}
+		}()
+	}
 	// ********************************************************************************
 	log.FromContext(ctx).Infof("startup completed in %v", time.Since(starttime))
 	// ********************************************************************************
 
 	// wait for server to exit
-	<-ctx.Done()
+	<-signalCtx.Done()
 }
 
 func getNseEndpoint(config *Config, listenOn fmt.Stringer) *registryapi.NetworkServiceEndpoint {