fix registry span

Signed-off-by: Artem Glazychev <[email protected]>
networkservicemesh · Sep 25, 2023 · edf767d · edf767d
1 parent fc827f8
commit edf767d
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 56 deletions.
diff --git a/pkg/registry/core/trace/traceconcise/ns_registry.go b/pkg/registry/core/trace/traceconcise/ns_registry.go
@@ -43,8 +43,7 @@ type conciseNetworkServiceRegistryFindClient struct {
 }
 
 func (t *conciseNetworkServiceRegistryFindClient) Recv() (*registry.NetworkServiceResponse, error) {
-	ctx, finish := withLog(t.Context(), methodNameRecv)
-	defer finish()
+	ctx := t.Context()
 
 	s := streamcontext.NetworkServiceRegistryFindClient(ctx, t.NetworkServiceRegistry_FindClient)
 	rv, err := s.Recv()
@@ -76,12 +75,11 @@ func (t *conciseNetworkServiceRegistryFindClient) Recv() (*registry.NetworkServi
 }
 
 func (t *conciseNetworkServiceRegistryClient) Register(ctx context.Context, in *registry.NetworkService, opts ...grpc.CallOption) (*registry.NetworkService, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameRegister)
-	defer finish()
-
 	ctx, tail := nsClientRegisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameRegister)
+		defer finish()
 		logObject(ctx, "ns-client-register", in)
 	}
 	ctx = withNsClientRegisterTail(ctx, t.traced)
@@ -104,15 +102,15 @@ func (t *conciseNetworkServiceRegistryClient) Register(ctx context.Context, in *
 }
 
 func (t *conciseNetworkServiceRegistryClient) Find(ctx context.Context, in *registry.NetworkServiceQuery, opts ...grpc.CallOption) (registry.NetworkServiceRegistry_FindClient, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameFind)
-	defer finish()
 	conciseFindClient := &conciseNetworkServiceRegistryFindClient{
 		operation: typeutils.GetFuncName(t.traced, methodNameRecv),
 	}
 
 	ctx, tail := nsClientFindTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameFind)
+		defer finish()
 		logObject(ctx, "ns-client-find", in)
 		ctx = withNsClientFindHead(ctx, conciseFindClient)
 	}
@@ -139,12 +137,11 @@ func (t *conciseNetworkServiceRegistryClient) Find(ctx context.Context, in *regi
 }
 
 func (t *conciseNetworkServiceRegistryClient) Unregister(ctx context.Context, in *registry.NetworkService, opts ...grpc.CallOption) (*empty.Empty, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameUnregister)
-	defer finish()
-
 	ctx, tail := nsClientUnregisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameUnregister)
+		defer finish()
 		logObject(ctx, "ns-client-unregister", in)
 	}
 	ctx = withNsClientUnregisterTail(ctx, t.traced)
@@ -177,12 +174,11 @@ type conciseNetworkServiceRegistryServer struct {
 }
 
 func (t *conciseNetworkServiceRegistryServer) Register(ctx context.Context, in *registry.NetworkService) (*registry.NetworkService, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameRegister)
-	defer finish()
-
 	ctx, tail := nsServerRegisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameRegister)
+		defer finish()
 		logObject(ctx, "ns-server-register", in)
 	}
 	ctx = withNsServerRegisterTail(ctx, t.traced)
@@ -205,14 +201,15 @@ func (t *conciseNetworkServiceRegistryServer) Register(ctx context.Context, in *
 }
 
 func (t *conciseNetworkServiceRegistryServer) Find(in *registry.NetworkServiceQuery, s registry.NetworkServiceRegistry_FindServer) error {
-	ctx, finish := withLog(s.Context(), methodNameFind)
-	defer finish()
 	conciseFindServer := &conciseNetworkServiceRegistryFindServer{
 		operation: typeutils.GetFuncName(t.traced, methodNameSend),
 	}
 
-	ctx, tail := nsServerFindTail(ctx)
+	ctx, tail := nsServerFindTail(s.Context())
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameFind)
+		defer finish()
 		logObject(ctx, "ns-server-find", in)
 		ctx = withNsServerFindHead(ctx, conciseFindServer)
 	}
@@ -239,12 +236,11 @@ func (t *conciseNetworkServiceRegistryServer) Find(in *registry.NetworkServiceQu
 }
 
 func (t *conciseNetworkServiceRegistryServer) Unregister(ctx context.Context, in *registry.NetworkService) (*empty.Empty, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameUnregister)
-	defer finish()
-
 	ctx, tail := nsServerUnregisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameUnregister)
+		defer finish()
 		logObject(ctx, "ns-server-unregister", in)
 	}
 	ctx = withNsServerUnregisterTail(ctx, t.traced)
@@ -278,10 +274,7 @@ type conciseNetworkServiceRegistryFindServer struct {
 }
 
 func (t *conciseNetworkServiceRegistryFindServer) Send(nsResp *registry.NetworkServiceResponse) error {
-	ctx, finish := withLog(t.Context(), methodNameSend)
-	defer finish()
-
-	ctx, tail := nsServerFindTail(ctx)
+	ctx, tail := nsServerFindTail(t.Context())
 	if tail == t {
 		logObject(ctx, "ns-server-send", nsResp)
 	}

diff --git a/pkg/registry/core/trace/traceconcise/nse_registry.go b/pkg/registry/core/trace/traceconcise/nse_registry.go
@@ -42,8 +42,7 @@ type conciseNetworkServiceEndpointRegistryFindClient struct {
 }
 
 func (t *conciseNetworkServiceEndpointRegistryFindClient) Recv() (*registry.NetworkServiceEndpointResponse, error) {
-	ctx, finish := withLog(t.Context(), methodNameRecv)
-	defer finish()
+	ctx := t.Context()
 
 	s := streamcontext.NetworkServiceEndpointRegistryFindClient(ctx, t.NetworkServiceEndpointRegistry_FindClient)
 	rv, err := s.Recv()
@@ -75,12 +74,11 @@ func (t *conciseNetworkServiceEndpointRegistryFindClient) Recv() (*registry.Netw
 }
 
 func (t *conciseNetworkServiceEndpointRegistryClient) Register(ctx context.Context, in *registry.NetworkServiceEndpoint, opts ...grpc.CallOption) (*registry.NetworkServiceEndpoint, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameRegister)
-	defer finish()
-
 	ctx, tail := nseClientRegisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameRegister)
+		defer finish()
 		logObject(ctx, "nse-client-register", in)
 	}
 	ctx = withNseClientRegisterTail(ctx, t.traced)
@@ -102,15 +100,15 @@ func (t *conciseNetworkServiceEndpointRegistryClient) Register(ctx context.Conte
 	return rv, nil
 }
 func (t *conciseNetworkServiceEndpointRegistryClient) Find(ctx context.Context, in *registry.NetworkServiceEndpointQuery, opts ...grpc.CallOption) (registry.NetworkServiceEndpointRegistry_FindClient, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameFind)
-	defer finish()
 	conciseFindClient := &conciseNetworkServiceEndpointRegistryFindClient{
 		operation: typeutils.GetFuncName(t.traced, methodNameRecv),
 	}
 
 	ctx, tail := nseClientFindTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameFind)
+		defer finish()
 		logObject(ctx, "nse-client-find", in)
 		ctx = withNseClientFindHead(ctx, conciseFindClient)
 	}
@@ -137,12 +135,11 @@ func (t *conciseNetworkServiceEndpointRegistryClient) Find(ctx context.Context,
 }
 
 func (t *conciseNetworkServiceEndpointRegistryClient) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint, opts ...grpc.CallOption) (*empty.Empty, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameUnregister)
-	defer finish()
-
 	ctx, tail := nseClientUnregisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameUnregister)
+		defer finish()
 		logObject(ctx, "nse-client-unregister", in)
 	}
 	ctx = withNseClientUnregisterTail(ctx, t.traced)
@@ -175,12 +172,11 @@ type conciseNetworkServiceEndpointRegistryServer struct {
 }
 
 func (t *conciseNetworkServiceEndpointRegistryServer) Register(ctx context.Context, in *registry.NetworkServiceEndpoint) (*registry.NetworkServiceEndpoint, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameRegister)
-	defer finish()
-
 	ctx, tail := nseServerRegisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameRegister)
+		defer finish()
 		logObject(ctx, "nse-server-register", in)
 	}
 	ctx = withNseServerRegisterTail(ctx, t.traced)
@@ -203,14 +199,15 @@ func (t *conciseNetworkServiceEndpointRegistryServer) Register(ctx context.Conte
 }
 
 func (t *conciseNetworkServiceEndpointRegistryServer) Find(in *registry.NetworkServiceEndpointQuery, s registry.NetworkServiceEndpointRegistry_FindServer) error {
-	ctx, finish := withLog(s.Context(), methodNameFind)
-	defer finish()
 	conciseFindServer := &conciseNetworkServiceEndpointRegistryFindServer{
 		operation: typeutils.GetFuncName(t.traced, methodNameSend),
 	}
 
-	ctx, tail := nseServerFindTail(ctx)
+	ctx, tail := nseServerFindTail(s.Context())
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(s.Context(), methodNameFind)
+		defer finish()
 		logObject(ctx, "nse-server-find", in)
 		ctx = withNseServerFindHead(ctx, conciseFindServer)
 	}
@@ -237,12 +234,11 @@ func (t *conciseNetworkServiceEndpointRegistryServer) Find(in *registry.NetworkS
 }
 
 func (t *conciseNetworkServiceEndpointRegistryServer) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint) (*empty.Empty, error) {
-	var finish func()
-	ctx, finish = withLog(ctx, methodNameUnregister)
-	defer finish()
-
 	ctx, tail := nseServerUnregisterTail(ctx)
 	if tail == nil {
+		var finish func()
+		ctx, finish = withLog(ctx, methodNameUnregister)
+		defer finish()
 		logObject(ctx, "nse-server-unregister", in)
 	}
 	ctx = withNseServerUnregisterTail(ctx, t.traced)
@@ -276,10 +272,7 @@ type conciseNetworkServiceEndpointRegistryFindServer struct {
 }
 
 func (t *conciseNetworkServiceEndpointRegistryFindServer) Send(nseResp *registry.NetworkServiceEndpointResponse) error {
-	ctx, finish := withLog(t.Context(), methodNameSend)
-	defer finish()
-
-	ctx, tail := nseServerFindTail(ctx)
+	ctx, tail := nseServerFindTail(t.Context())
 	if tail == t {
 		logObject(ctx, "nse-server-send", nseResp)
 	}