fix other levels

Signed-off-by: Artem Glazychev <[email protected]>

pkg/networkservice/core/trace/client.go

@@ -27,19 +27,21 @@ import (
 
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
 
-	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/trace/traceconcise"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/trace/traceverbose"
 )
 
 type traceClient struct {
+	traced networkservice.NetworkServiceClient
+
 	verbose networkservice.NetworkServiceClient
 	concise networkservice.NetworkServiceClient
 }
 
 // NewNetworkServiceClient - wraps tracing around the supplied networkservice.NetworkServiceClient
 func NewNetworkServiceClient(traced networkservice.NetworkServiceClient) networkservice.NetworkServiceClient {
 	return &traceClient{
+		traced:  traced,
 		verbose: traceverbose.NewNetworkServiceClient(traced),
 		concise: traceconcise.NewNetworkServiceClient(traced),
 	}
@@ -52,7 +54,7 @@ func (t *traceClient) Request(ctx context.Context, request *networkservice.Netwo
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Request(ctx, request, opts...)
 	}
-	return next.Client(ctx).Request(ctx, request, opts...)
+	return t.traced.Request(ctx, request, opts...)
 }
 
 func (t *traceClient) Close(ctx context.Context, conn *networkservice.Connection, opts ...grpc.CallOption) (*empty.Empty, error) {
@@ -62,5 +64,5 @@ func (t *traceClient) Close(ctx context.Context, conn *networkservice.Connection
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Close(ctx, conn, opts...)
 	}
-	return next.Client(ctx).Close(ctx, conn, opts...)
+	return t.traced.Close(ctx, conn, opts...)
 }

pkg/networkservice/core/trace/server.go

@@ -27,18 +27,19 @@ import (
 	"github.com/golang/protobuf/ptypes/empty"
 
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
-
-	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
 )
 
 type traceServer struct {
+	traced networkservice.NetworkServiceServer
+
 	verbose networkservice.NetworkServiceServer
 	concise networkservice.NetworkServiceServer
 }
 
 // NewNetworkServiceServer - wraps tracing around the supplied traced
 func NewNetworkServiceServer(traced networkservice.NetworkServiceServer) networkservice.NetworkServiceServer {
 	return &traceServer{
+		traced:  traced,
 		verbose: traceverbose.NewNetworkServiceServer(traced),
 		concise: traceconcise.NewNetworkServiceServer(traced),
 	}
@@ -51,7 +52,7 @@ func (t *traceServer) Request(ctx context.Context, request *networkservice.Netwo
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Request(ctx, request)
 	}
-	return next.Server(ctx).Request(ctx, request)
+	return t.traced.Request(ctx, request)
 }
 
 func (t *traceServer) Close(ctx context.Context, conn *networkservice.Connection) (*empty.Empty, error) {
@@ -61,5 +62,5 @@ func (t *traceServer) Close(ctx context.Context, conn *networkservice.Connection
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Close(ctx, conn)
 	}
-	return next.Server(ctx).Close(ctx, conn)
+	return t.traced.Close(ctx, conn)
 }

pkg/registry/core/trace/ns_registry.go

@@ -26,19 +26,21 @@ import (
 	"github.com/networkservicemesh/api/pkg/api/registry"
 	"google.golang.org/grpc"
 
-	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/trace/traceconcise"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/trace/traceverbose"
 )
 
 type traceNetworkServiceRegistryClient struct {
+	traced registry.NetworkServiceRegistryClient
+
 	verbose registry.NetworkServiceRegistryClient
 	concise registry.NetworkServiceRegistryClient
 }
 
 // NewNetworkServiceRegistryClient - wraps registry.NetworkServiceRegistryClient with tracing
 func NewNetworkServiceRegistryClient(traced registry.NetworkServiceRegistryClient) registry.NetworkServiceRegistryClient {
 	return &traceNetworkServiceRegistryClient{
+		traced:  traced,
 		verbose: traceverbose.NewNetworkServiceRegistryClient(traced),
 		concise: traceconcise.NewNetworkServiceRegistryClient(traced),
 	}
@@ -51,7 +53,7 @@ func (t *traceNetworkServiceRegistryClient) Register(ctx context.Context, in *re
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Register(ctx, in, opts...)
 	}
-	return next.NetworkServiceRegistryClient(ctx).Register(ctx, in, opts...)
+	return t.traced.Register(ctx, in, opts...)
 }
 func (t *traceNetworkServiceRegistryClient) Find(ctx context.Context, in *registry.NetworkServiceQuery, opts ...grpc.CallOption) (registry.NetworkServiceRegistry_FindClient, error) {
 	switch logrus.GetLevel() {
@@ -60,7 +62,7 @@ func (t *traceNetworkServiceRegistryClient) Find(ctx context.Context, in *regist
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Find(ctx, in, opts...)
 	}
-	return next.NetworkServiceRegistryClient(ctx).Find(ctx, in, opts...)
+	return t.traced.Find(ctx, in, opts...)
 }
 
 func (t *traceNetworkServiceRegistryClient) Unregister(ctx context.Context, in *registry.NetworkService, opts ...grpc.CallOption) (*empty.Empty, error) {
@@ -70,17 +72,20 @@ func (t *traceNetworkServiceRegistryClient) Unregister(ctx context.Context, in *
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Unregister(ctx, in, opts...)
 	}
-	return next.NetworkServiceRegistryClient(ctx).Unregister(ctx, in, opts...)
+	return t.traced.Unregister(ctx, in, opts...)
 }
 
 type traceNetworkServiceRegistryServer struct {
+	traced registry.NetworkServiceRegistryServer
+
 	verbose registry.NetworkServiceRegistryServer
 	concise registry.NetworkServiceRegistryServer
 }
 
 // NewNetworkServiceRegistryServer - wraps registry.NetworkServiceRegistryServer with tracing
 func NewNetworkServiceRegistryServer(traced registry.NetworkServiceRegistryServer) registry.NetworkServiceRegistryServer {
 	return &traceNetworkServiceRegistryServer{
+		traced:  traced,
 		verbose: traceverbose.NewNetworkServiceRegistryServer(traced),
 		concise: traceconcise.NewNetworkServiceRegistryServer(traced),
 	}
@@ -93,7 +98,7 @@ func (t *traceNetworkServiceRegistryServer) Register(ctx context.Context, in *re
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Register(ctx, in)
 	}
-	return next.NetworkServiceRegistryServer(ctx).Register(ctx, in)
+	return t.traced.Register(ctx, in)
 }
 
 func (t *traceNetworkServiceRegistryServer) Find(in *registry.NetworkServiceQuery, s registry.NetworkServiceRegistry_FindServer) error {
@@ -103,7 +108,7 @@ func (t *traceNetworkServiceRegistryServer) Find(in *registry.NetworkServiceQuer
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Find(in, s)
 	}
-	return next.NetworkServiceRegistryServer(s.Context()).Find(in, s)
+	return t.traced.Find(in, s)
 }
 
 func (t *traceNetworkServiceRegistryServer) Unregister(ctx context.Context, in *registry.NetworkService) (*empty.Empty, error) {
@@ -113,5 +118,5 @@ func (t *traceNetworkServiceRegistryServer) Unregister(ctx context.Context, in *
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Unregister(ctx, in)
 	}
-	return next.NetworkServiceRegistryServer(ctx).Unregister(ctx, in)
+	return t.traced.Unregister(ctx, in)
 }

pkg/registry/core/trace/nse_registry.go

@@ -20,24 +20,26 @@ import (
 	"context"
 
 	"github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
 
-	"github.com/networkservicemesh/sdk/pkg/registry/core/next"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/trace/traceconcise"
 	"github.com/networkservicemesh/sdk/pkg/registry/core/trace/traceverbose"
 
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/networkservicemesh/api/pkg/api/registry"
-	"google.golang.org/grpc"
 )
 
 type traceNetworkServiceEndpointRegistryClient struct {
+	traced registry.NetworkServiceEndpointRegistryClient
+
 	verbose registry.NetworkServiceEndpointRegistryClient
 	concise registry.NetworkServiceEndpointRegistryClient
 }
 
 // NewNetworkServiceEndpointRegistryClient - wraps registry.NetworkServiceEndpointRegistryClient with tracing
 func NewNetworkServiceEndpointRegistryClient(traced registry.NetworkServiceEndpointRegistryClient) registry.NetworkServiceEndpointRegistryClient {
 	return &traceNetworkServiceEndpointRegistryClient{
+		traced:  traced,
 		verbose: traceverbose.NewNetworkServiceEndpointRegistryClient(traced),
 		concise: traceconcise.NewNetworkServiceEndpointRegistryClient(traced),
 	}
@@ -50,7 +52,7 @@ func (t *traceNetworkServiceEndpointRegistryClient) Register(ctx context.Context
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Register(ctx, in, opts...)
 	}
-	return next.NetworkServiceEndpointRegistryClient(ctx).Register(ctx, in, opts...)
+	return t.traced.Register(ctx, in, opts...)
 }
 func (t *traceNetworkServiceEndpointRegistryClient) Find(ctx context.Context, in *registry.NetworkServiceEndpointQuery, opts ...grpc.CallOption) (registry.NetworkServiceEndpointRegistry_FindClient, error) {
 	switch logrus.GetLevel() {
@@ -59,7 +61,7 @@ func (t *traceNetworkServiceEndpointRegistryClient) Find(ctx context.Context, in
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Find(ctx, in, opts...)
 	}
-	return next.NetworkServiceEndpointRegistryClient(ctx).Find(ctx, in, opts...)
+	return t.traced.Find(ctx, in, opts...)
 }
 
 func (t *traceNetworkServiceEndpointRegistryClient) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint, opts ...grpc.CallOption) (*empty.Empty, error) {
@@ -69,17 +71,20 @@ func (t *traceNetworkServiceEndpointRegistryClient) Unregister(ctx context.Conte
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Unregister(ctx, in, opts...)
 	}
-	return next.NetworkServiceEndpointRegistryClient(ctx).Unregister(ctx, in, opts...)
+	return t.traced.Unregister(ctx, in, opts...)
 }
 
 type traceNetworkServiceEndpointRegistryServer struct {
+	traced registry.NetworkServiceEndpointRegistryServer
+
 	verbose registry.NetworkServiceEndpointRegistryServer
 	concise registry.NetworkServiceEndpointRegistryServer
 }
 
 // NewNetworkServiceEndpointRegistryServer - wraps registry.NetworkServiceEndpointRegistryServer with tracing
 func NewNetworkServiceEndpointRegistryServer(traced registry.NetworkServiceEndpointRegistryServer) registry.NetworkServiceEndpointRegistryServer {
 	return &traceNetworkServiceEndpointRegistryServer{
+		traced:  traced,
 		verbose: traceverbose.NewNetworkServiceEndpointRegistryServer(traced),
 		concise: traceconcise.NewNetworkServiceEndpointRegistryServer(traced),
 	}
@@ -92,7 +97,7 @@ func (t *traceNetworkServiceEndpointRegistryServer) Register(ctx context.Context
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Register(ctx, in)
 	}
-	return next.NetworkServiceEndpointRegistryServer(ctx).Register(ctx, in)
+	return t.traced.Register(ctx, in)
 }
 
 func (t *traceNetworkServiceEndpointRegistryServer) Find(in *registry.NetworkServiceEndpointQuery, s registry.NetworkServiceEndpointRegistry_FindServer) error {
@@ -102,7 +107,7 @@ func (t *traceNetworkServiceEndpointRegistryServer) Find(in *registry.NetworkSer
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Find(in, s)
 	}
-	return next.NetworkServiceEndpointRegistryServer(s.Context()).Find(in, s)
+	return t.traced.Find(in, s)
 }
 
 func (t *traceNetworkServiceEndpointRegistryServer) Unregister(ctx context.Context, in *registry.NetworkServiceEndpoint) (*empty.Empty, error) {
@@ -112,5 +117,5 @@ func (t *traceNetworkServiceEndpointRegistryServer) Unregister(ctx context.Conte
 	case logrus.InfoLevel, logrus.DebugLevel:
 		return t.concise.Unregister(ctx, in)
 	}
-	return next.NetworkServiceEndpointRegistryServer(ctx).Unregister(ctx, in)
+	return t.traced.Unregister(ctx, in)
 }