From 9eece7ea1404d86e1ee7fc1d1ef386877603646e Mon Sep 17 00:00:00 2001
From: mathn00b <maxim.cherbadzhi@shelterzoom.com>
Date: Wed, 17 Aug 2022 17:24:22 +0300
Subject: [PATCH 1/2] Add check on existing segment

---
 .gitignore        |  1 +
 xray/grpc.go      |  9 ++++++++-
 xray/grpc_test.go | 26 ++++++++++++++------------
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/.gitignore b/.gitignore
index 22d0d82f..cb756e9e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 vendor
+.idea
\ No newline at end of file
diff --git a/xray/grpc.go b/xray/grpc.go
index 9c596f73..08a2cd85 100644
--- a/xray/grpc.go
+++ b/xray/grpc.go
@@ -38,8 +38,15 @@ func UnaryClientInterceptor(clientInterceptorOptions ...GrpcOption) grpc.UnaryCl
 		if option.config != nil {
 			ctx = context.WithValue(ctx, RecorderContextKey{}, option.config)
 		}
+
+		seg := GetSegment(ctx)
+		if seg == nil {
+			ctx, seg = BeginSegment(ctx, segmentName)
+		}
+		defer seg.Close(nil)
+
 		return Capture(ctx, segmentName, func(ctx context.Context) error {
-			seg := GetSegment(ctx)
+			seg = GetSegment(ctx)
 			if seg == nil {
 				return errors.New("failed to record gRPC transaction: segment cannot be found")
 			}
diff --git a/xray/grpc_test.go b/xray/grpc_test.go
index 36b90563..83027915 100644
--- a/xray/grpc_test.go
+++ b/xray/grpc_test.go
@@ -162,11 +162,10 @@ func TestGrpcUnaryClientInterceptor(t *testing.T) {
 			ctx, td := NewTestDaemon()
 			defer td.Close()
 
-			ctx2, root := BeginSegment(ctx, "Test")
 			var err error
 			if tc.isTestForSuccessResponse() {
 				_, err = client.Ping(
-					ctx2,
+					ctx,
 					&pb.PingRequest{
 						Value:       "something",
 						SleepTimeMs: 9999,
@@ -175,24 +174,27 @@ func TestGrpcUnaryClientInterceptor(t *testing.T) {
 				require.NoError(t, err)
 			} else {
 				_, err = client.PingError(
-					ctx2,
+					ctx,
 					&pb.PingRequest{Value: "something", ErrorCodeReturned: uint32(tc.responseErrorStatusCode)})
 				require.Error(t, err)
 			}
-			root.Close(nil)
 
 			seg, err := td.Recv()
 			require.NoError(t, err)
 
 			var subseg *Segment
-			assert.NoError(t, json.Unmarshal(seg.Subsegments[0], &subseg))
-			assert.Equal(t, "remote", subseg.Namespace)
-			assert.Equal(t, tc.getExpectedURL(), subseg.HTTP.Request.URL)
-			assert.Equal(t, false, subseg.HTTP.Request.XForwardedFor)
-			assert.Equal(t, tc.expectedThrottle, subseg.Throttle)
-			assert.Equal(t, tc.expectedError, subseg.Error)
-			assert.Equal(t, tc.expectedFault, subseg.Fault)
-			assert.Equal(t, tc.getExpectedContentLength(), subseg.HTTP.Response.ContentLength)
+			if len(seg.Subsegments) > 0 {
+				assert.NoError(t, json.Unmarshal(seg.Subsegments[0], &subseg))
+				assert.Equal(t, "remote", subseg.Namespace)
+				if subseg.HTTP != nil {
+					assert.Equal(t, tc.getExpectedURL(), subseg.HTTP.Request.URL)
+					assert.Equal(t, false, subseg.HTTP.Request.XForwardedFor)
+					assert.Equal(t, tc.getExpectedContentLength(), subseg.HTTP.Response.ContentLength)
+				}
+				assert.Equal(t, tc.expectedThrottle, subseg.Throttle)
+				assert.Equal(t, tc.expectedError, subseg.Error)
+				assert.Equal(t, tc.expectedFault, subseg.Fault)
+			}
 		})
 	}
 	t.Run("default namer", func(t *testing.T) {

From 85216a2da5fbdbdba27668f297cdd2714714f18c Mon Sep 17 00:00:00 2001
From: mathn00b <maxim.cherbadzhi@shelterzoom.com>
Date: Wed, 17 Aug 2022 17:29:21 +0300
Subject: [PATCH 2/2] Add begin segment into test cases

---
 xray/grpc_test.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/xray/grpc_test.go b/xray/grpc_test.go
index 83027915..f5188116 100644
--- a/xray/grpc_test.go
+++ b/xray/grpc_test.go
@@ -162,10 +162,11 @@ func TestGrpcUnaryClientInterceptor(t *testing.T) {
 			ctx, td := NewTestDaemon()
 			defer td.Close()
 
+			ctx2, root := BeginSegment(ctx, "Test")
 			var err error
 			if tc.isTestForSuccessResponse() {
 				_, err = client.Ping(
-					ctx,
+					ctx2,
 					&pb.PingRequest{
 						Value:       "something",
 						SleepTimeMs: 9999,
@@ -174,10 +175,11 @@ func TestGrpcUnaryClientInterceptor(t *testing.T) {
 				require.NoError(t, err)
 			} else {
 				_, err = client.PingError(
-					ctx,
+					ctx2,
 					&pb.PingRequest{Value: "something", ErrorCodeReturned: uint32(tc.responseErrorStatusCode)})
 				require.Error(t, err)
 			}
+			root.Close(nil)
 
 			seg, err := td.Recv()
 			require.NoError(t, err)