Heal rework

pkg/networkservice/chains/client/client.go

@@ -1,5 +1,7 @@
 // Copyright (c) 2020-2021 Cisco Systems, Inc.
 //
+// Copyright (c) 2021 Doc.ai and/or its affiliates.
+//
 // SPDX-License-Identifier: Apache-2.0
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,9 +22,9 @@ package client
 import (
 	"context"
 
+	"github.com/google/uuid"
 	"google.golang.org/grpc"
 
-	"github.com/google/uuid"
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
 
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/heal"
@@ -37,18 +39,18 @@ import (
 
 type clientOptions struct {
 	name                    string
-	onHeal                  *networkservice.NetworkServiceClient
 	additionalFunctionality []networkservice.NetworkServiceClient
 	authorizeClient         networkservice.NetworkServiceClient
+	registerClientFunc      heal.RegisterClientFunc
 }
 
 // Option modifies default client chain values.
 type Option func(c *clientOptions)
 
-// WithHeal sets heal for the client.
-func WithHeal(onHeal *networkservice.NetworkServiceClient) Option {
+// WithRegisterHealClientFunc sets server's register client method
+func WithRegisterHealClientFunc(registerClientFunc heal.RegisterClientFunc) Option {
 	return Option(func(c *clientOptions) {
-		c.onHeal = onHeal
+		c.registerClientFunc = registerClientFunc
 	})
 }
 
@@ -84,7 +86,6 @@ func NewClient(ctx context.Context, cc grpc.ClientConnInterface, clientOpts ...O
 	var opts = &clientOptions{
 		name:            "client-" + uuid.New().String(),
 		authorizeClient: null.NewClient(),
-		onHeal:          &rv,
 	}
 	for _, opt := range clientOpts {
 		opt(opts)
@@ -94,7 +95,7 @@ func NewClient(ctx context.Context, cc grpc.ClientConnInterface, clientOpts ...O
 			append([]networkservice.NetworkServiceClient{
 				updatepath.NewClient(opts.name),
 				serialize.NewClient(),
-				heal.NewClient(ctx, networkservice.NewMonitorConnectionClient(cc), opts.onHeal),
+				heal.NewClient(ctx, networkservice.NewMonitorConnectionClient(cc), opts.registerClientFunc),
 				refresh.NewClient(ctx),
 				metadata.NewClient(),
 			}, opts.additionalFunctionality...),

pkg/networkservice/chains/nsmgr/server.go

@@ -34,6 +34,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/discover"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/excludedprefixes"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/filtermechanisms"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/common/heal"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/interpose"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/mechanisms/recvfd"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/mechanisms/sendfd"
@@ -111,6 +112,8 @@ func NewServer(ctx context.Context, nsmRegistration *registryapi.NetworkServiceE
 
 	nsClient := registryadapter.NetworkServiceServerToClient(nsRegistry)
 
+	healServer, registerHealClient := heal.NewServer(ctx, addressof.NetworkServiceClient(adapters.NewServerToClient(rv)))
+
 	// Construct Endpoint
 	rv.Endpoint = endpoint.NewServer(ctx, tokenGenerator,
 		endpoint.WithName(nsmRegistration.Name),
@@ -122,10 +125,11 @@ func NewServer(ctx context.Context, nsmRegistration *registryapi.NetworkServiceE
 			recvfd.NewServer(), // Receive any files passed
 			interpose.NewServer(&interposeRegistryServer),
 			filtermechanisms.NewServer(&urlsRegistryServer),
+			healServer,
 			connect.NewServer(ctx,
 				client.NewClientFactory(
 					client.WithName(nsmRegistration.Name),
-					client.WithHeal(addressof.NetworkServiceClient(adapters.NewServerToClient(rv))),
+					client.WithRegisterHealClientFunc(registerHealClient),
 					client.WithAdditionalFunctionality(
 						recvfd.NewClient(),
 						sendfd.NewClient(),