include more integration tests

Author: simaoribeiro

src/Cassandra.IntegrationTests/OpenTelemetry/OpenTelemetryTests.cs

@@ -20,7 +20,11 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Cassandra.Data.Linq;
+using Cassandra.IntegrationTests.Policies.Util;
+using Cassandra.IntegrationTests.SimulacronAPI;
 using Cassandra.IntegrationTests.TestBase;
+using Cassandra.IntegrationTests.TestClusterManagement.Simulacron;
 using Cassandra.Mapping;
 using Cassandra.OpenTelemetry;
 using Cassandra.Tests;
@@ -42,10 +46,13 @@ public class OpenTelemetryTests : SharedClusterTest
 
         private readonly List<Activity> _exportedActivities = new List<Activity>();
 
+        private readonly ActivitySource InternalActivitySource = new ActivitySource("testeActivitySource");
+
         public OpenTelemetryTests()
         {
             Sdk.CreateTracerProviderBuilder()
               .AddSource(OpenTelemetrySourceName)
+              .AddSource(InternalActivitySource.Name)
               .AddInMemoryExporter(_exportedActivities)
               .Build();
         }
@@ -191,16 +198,7 @@ public async Task AddOpenTelemetry_MapperAndMapperAsync_SessionRequestIsParentOf
 
             session.ChangeKeyspace(keyspace);
 
-            string createTableQuery = @"
-            CREATE TABLE IF NOT EXISTS song (
-                Artist text,
-                Title text,
-                Id uuid,
-                ReleaseDate timestamp,
-                PRIMARY KEY (id)
-            )";
-
-            session.Execute(createTableQuery);
+            CreateSongTable(session);
 
             var mapper = new Mapper(session);
 
@@ -259,6 +257,174 @@ await mapper.InsertIfNotExistsAsync(songOne, testProfile, true, null)
             ValidateNodeActivityAttributes(syncNodeActivity);
         }
 
+        [Category(TestCategory.RealClusterLong)]
+        [Test]
+        [NUnit.Framework.Ignore("LINQ flow seems different than the flow in which OTEL is implemented!")]
+        public void AddOpenTelemetry_Linq_SessionRequestIsParentOfNodeRequest()
+        {
+            var keyspace = TestUtils.GetUniqueKeyspaceName().ToLowerInvariant();
+
+            var cluster = GetNewTemporaryCluster(b => b
+                .AddOpenTelemetryInstrumentation());
+
+            var session = cluster.Connect();
+
+            session.CreateKeyspaceIfNotExists(keyspace);
+
+            session.ChangeKeyspace(keyspace);
+
+            CreateSongTable(session);
+
+            var table = new Table<Song>(Session, new MappingConfiguration().Define(new Map<Song>().TableName("song")));
+
+            var song = new Song
+            {
+                Id = Guid.NewGuid(),
+                Artist = "Led Zeppelin",
+                Title = "Mothership",
+                ReleaseDate = DateTimeOffset.UtcNow
+            };
+
+            // Clear activities to get the Linq one
+            _exportedActivities.Clear();
+
+            table.Insert(song);
+
+            var syncActivities = GetActivities();
+        }
+
+        [Test]
+        public void AddOpenTelemetry_WhenMethodIsInvokedAfterQuery_TraceIdIsTheSameThroughRequest()
+        {
+            var firstMethodName = "FirstMethod";
+            var secondMethodName = "SecondMethod";
+
+            using (var activity = InternalActivitySource.StartActivity(firstMethodName, ActivityKind.Internal))
+            {
+                var cluster = GetNewTemporaryCluster(b => b.AddOpenTelemetryInstrumentation());
+                var session = cluster.Connect();
+
+                var statement = new SimpleStatement("SELECT key FROM system.local");
+                session.Execute(statement);
+
+                SecondMethod(secondMethodName);
+            }
+
+            var firstMethodActivity = GetActivities().First(x => x.DisplayName == firstMethodName);
+            var secondMethodActivity = GetActivities().First(x => x.DisplayName == secondMethodName);
+            var sessionActivity = GetActivities().First(x => x.DisplayName == SessionActivityName);
+
+            Assert.AreEqual(firstMethodActivity.TraceId, sessionActivity.TraceId);
+            Assert.AreEqual(firstMethodActivity.SpanId, sessionActivity.ParentSpanId);
+
+            Assert.AreEqual(firstMethodActivity.TraceId, secondMethodActivity.TraceId);
+            Assert.AreEqual(firstMethodActivity.SpanId, secondMethodActivity.ParentSpanId);
+        }
+
+        [Test]
+        public void AddOpenTelemetry_RetryOnNextHost_ShouldProduceOneErrorAndOneValidSpansForTheSameSessionSpan()
+        {
+            var expectedErrorDescription = "overloaded";
+
+            using (var simulacronCluster = SimulacronCluster.CreateNew(3))
+            {
+                var contactPoint = simulacronCluster.InitialContactPoint;
+                var nodes = simulacronCluster.GetNodes().ToArray();
+                var loadBalancingPolicy = new CustomLoadBalancingPolicy(
+                    nodes.Select(n => n.ContactPoint).ToArray());
+                
+                var builder = ClusterBuilder()
+                                     .AddContactPoint(contactPoint)
+                                     .WithSocketOptions(new SocketOptions()
+                                                        .SetConnectTimeoutMillis(10000)
+                                                        .SetReadTimeoutMillis(5000))
+                                     .WithLoadBalancingPolicy(loadBalancingPolicy)
+                                     .WithRetryPolicy(TryNextHostRetryPolicy.Instance)
+                                     .AddOpenTelemetryInstrumentation();
+
+                using (var cluster = builder.Build())
+                {
+                    var session = (Session)cluster.Connect();
+                    const string cql = "select * from table2";
+
+                    nodes[0].PrimeFluent(
+                        b => b.WhenQuery(cql).
+                               ThenOverloaded(expectedErrorDescription));
+
+                    nodes[1].PrimeFluent(
+                        b => b.WhenQuery(cql).
+                               ThenRowsSuccess(new[] { ("text", DataType.Ascii) }, rows => rows.WithRow("test1").WithRow("test2")));
+
+                    session.Execute(new SimpleStatement(cql).SetConsistencyLevel(ConsistencyLevel.One));
+
+                    var activities = GetActivities();
+                    var sessionActivity = activities.First(x => x.DisplayName.StartsWith(SessionActivityName));
+                    var validNodeActivity = activities.First(x => x.DisplayName.StartsWith(NodeActivityName) && x.Status != ActivityStatusCode.Error);
+                    var errorNodeActivity = activities.First(x => x.DisplayName.StartsWith(NodeActivityName) && x.Status == ActivityStatusCode.Error);
+                    
+                    Assert.AreEqual(sessionActivity.TraceId, validNodeActivity.TraceId);
+                    Assert.AreEqual(sessionActivity.TraceId, errorNodeActivity.TraceId);
+                    Assert.AreEqual(sessionActivity.SpanId, validNodeActivity.ParentSpanId);
+                    Assert.AreEqual(sessionActivity.SpanId, errorNodeActivity.ParentSpanId);
+                    Assert.AreEqual(expectedErrorDescription, errorNodeActivity.StatusDescription);
+                    Assert.True(validNodeActivity.StartTimeUtc > errorNodeActivity.StartTimeUtc);
+                }
+            }
+        }
+
+        [Test]
+        public void AddOpenTelemetry_WithPaginationOnQuery_ShouldMultipleSpansForTheSameTraceId()
+        {
+            var expectedNumberOfSessionActivities = 2;
+
+            using (var activity = InternalActivitySource.StartActivity("Paging", ActivityKind.Internal))
+            {
+                var session = GetNewTemporaryCluster(b => b.AddOpenTelemetryInstrumentation()).Connect();
+
+                session.CreateKeyspaceIfNotExists(KeyspaceName, null, false);
+
+                session.ChangeKeyspace(KeyspaceName);
+
+                CreateSongTable(session);
+
+                for (int i = 0; i < expectedNumberOfSessionActivities; i++)
+                {
+                    session.Execute(string.Format("INSERT INTO {0}.song (Artist, Title, Id, ReleaseDate) VALUES('Pink Floyd', 'The Dark Side Of The Moon', {1}, {2})", KeyspaceName, Guid.NewGuid(), ((DateTimeOffset)DateTime.UtcNow).ToUnixTimeSeconds()));
+                }
+
+                _exportedActivities.Clear();
+
+                var rs = session.Execute(new SimpleStatement(string.Format("SELECT * FROM {0}.song", KeyspaceName)).SetPageSize(1));
+
+                rs.FetchMoreResults();
+
+                var sessionActivities = GetActivities().Where(x => x.DisplayName == SessionActivityName);
+
+                Assert.AreEqual(expectedNumberOfSessionActivities, sessionActivities.Count());
+                Assert.AreEqual(expectedNumberOfSessionActivities, rs.InnerQueueCount);
+
+                var firstActivity = sessionActivities.First();
+                var lastActivity = sessionActivities.Last();
+                
+                Assert.AreEqual(firstActivity.TraceId, lastActivity.TraceId);
+                Assert.AreNotEqual(firstActivity.SpanId, lastActivity.SpanId);
+            }
+        }
+
+        private static void CreateSongTable(ISession session)
+        {
+            string createTableQuery = @"
+            CREATE TABLE IF NOT EXISTS song (
+                Artist text,
+                Title text,
+                Id uuid,
+                ReleaseDate timestamp,
+                PRIMARY KEY (id)
+            )";
+
+            session.Execute(createTableQuery);
+        }
+
         private List<Activity> GetActivities()
         {
             var count = 0;
@@ -317,5 +483,13 @@ private static void ValidateNodeActivityAttributes(Activity activity)
                 Assert.AreEqual(tags.FirstOrDefault(x => x.Key == pair.Key).Value, expectedTags[pair.Key]);
             }
         }
+
+        private void SecondMethod(string activityName)
+        {
+            using (var activity = InternalActivitySource.StartActivity(activityName, ActivityKind.Internal))
+            {
+                activity.AddTag("db.test", "t");
+            }
+        }
     }
 }

src/Cassandra.IntegrationTests/Policies/Tests/RetryPolicyShortTests.cs

@@ -19,6 +19,7 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Cassandra.IntegrationTests.Policies.Util;
 using Cassandra.IntegrationTests.SimulacronAPI;
 using Cassandra.IntegrationTests.SimulacronAPI.PrimeBuilder.Then;
 using Cassandra.IntegrationTests.TestBase;
@@ -263,37 +264,5 @@ public RetryDecision OnRequestError(IStatement statement, Configuration config,
                 return RetryDecision.Retry(null, true);
             }
         }
-
-        private class CustomLoadBalancingPolicy : ILoadBalancingPolicy
-        {
-            private ICluster _cluster;
-            private readonly string[] _hosts;
-
-            public CustomLoadBalancingPolicy(string[] hosts)
-            {
-                _hosts = hosts;
-            }
-
-            public void Initialize(ICluster cluster)
-            {
-                _cluster = cluster;
-            }
-
-            public HostDistance Distance(Host host)
-            {
-                return HostDistance.Local;
-            }
-
-            public IEnumerable<Host> NewQueryPlan(string keyspace, IStatement query)
-            {
-                var queryPlan = new List<Host>();
-                var allHosts = _cluster.AllHosts();
-                foreach (var host in _hosts)
-                {
-                    queryPlan.Add(allHosts.Single(h => h.Address.ToString() == host));
-                }
-                return queryPlan;
-            }
-        }
     }
 }

src/Cassandra.IntegrationTests/Policies/Util/CustomLoadBalancingPolicy.cs

@@ -0,0 +1,37 @@
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Cassandra.IntegrationTests.Policies.Util
+{
+    public class CustomLoadBalancingPolicy : ILoadBalancingPolicy
+    {
+        private ICluster _cluster;
+        private readonly string[] _hosts;
+
+        public CustomLoadBalancingPolicy(string[] hosts)
+        {
+            _hosts = hosts;
+        }
+
+        public void Initialize(ICluster cluster)
+        {
+            _cluster = cluster;
+        }
+
+        public HostDistance Distance(Host host)
+        {
+            return HostDistance.Local;
+        }
+
+        public IEnumerable<Host> NewQueryPlan(string keyspace, IStatement query)
+        {
+            var queryPlan = new List<Host>();
+            var allHosts = _cluster.AllHosts();
+            foreach (var host in _hosts)
+            {
+                queryPlan.Add(allHosts.Single(h => h.Address.ToString() == host));
+            }
+            return queryPlan;
+        }
+    }
+}

src/Cassandra/Requests/RequestHandler.cs

@@ -25,7 +25,6 @@
 using Cassandra.Collections;
 using Cassandra.Connections;
 using Cassandra.ExecutionProfiles;
-using Cassandra.Observers;
 using Cassandra.Observers.Abstractions;
 using Cassandra.Serialization;
 using Cassandra.SessionManagement;

src/Extensions/Cassandra.OpenTelemetry/OpenTelemetryRequestTracker.cs

@@ -134,13 +134,17 @@ public virtual Task OnErrorAsync(RequestTrackingInfo request, Exception ex)
         /// <returns>Completed task./returns>
         public virtual Task OnNodeSuccessAsync(RequestTrackingInfo request, HostTrackingInfo hostInfo)
         {
-            request.Items.TryGetValue($"{otelActivityKey}.{hostInfo.Host.HostId}", out object context);
+            var activityKey = $"{otelActivityKey}.{hostInfo.Host.HostId}";
+
+            request.Items.TryGetValue(activityKey, out object context);
 
             if (context is Activity activity)
             {
                 activity?.Dispose();
             }
 
+            request.Items.TryRemove(activityKey, out _);
+
             return Task.CompletedTask;
         }
 
@@ -154,7 +158,9 @@ public virtual Task OnNodeSuccessAsync(RequestTrackingInfo request, HostTracking
         /// <returns>Completed task./returns>
         public virtual Task OnNodeErrorAsync(RequestTrackingInfo request, HostTrackingInfo hostInfo, Exception ex)
         {
-            request.Items.TryGetValue($"{otelActivityKey}.{hostInfo.Host.HostId}", out object context);
+            var activityKey = $"{otelActivityKey}.{hostInfo.Host.HostId}";
+            
+            request.Items.TryGetValue(activityKey, out object context);
 
             if (!(context is Activity activity))
             {
@@ -166,6 +172,8 @@ public virtual Task OnNodeErrorAsync(RequestTrackingInfo request, HostTrackingIn
 
             activity?.Dispose();
 
+            request.Items.TryRemove(activityKey, out _);
+
             return Task.CompletedTask;
         }