diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommon.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommon.cs index 278e09fcde5f..429d935791e9 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommon.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommon.cs @@ -40,12 +40,34 @@ public static CallTargetState CreateContainerCallStateExt(out var databaseNew)) { databaseId = databaseNew.Id; - endpoint = databaseNew.Client.Endpoint; + try + { + endpoint = databaseNew.Client.Endpoint; + } + catch (ObjectDisposedException) + { + // Container/Client disposed; skip endpoint to avoid throwing in instrumentation + } + catch (Exception ex) + { + Log.Warning(ex, "Unable to get Cosmos client endpoint from DatabaseNewStruct."); + } } else if (container.Database.TryDuckCast(out var databaseOld)) { databaseId = databaseOld.Id; - endpoint = databaseOld.ClientContext.Client.Endpoint; + try + { + endpoint = databaseOld.ClientContext.Client.Endpoint; + } + catch (ObjectDisposedException) + { + // Container/Client disposed; skip endpoint to avoid throwing in instrumentation + } + catch (Exception ex) + { + Log.Warning(ex, "Unable to get Cosmos client endpoint from DatabaseOldStruct."); + } } } @@ -60,14 +82,34 @@ public static CallTargetState CreateDatabaseCallStateExt(out var databaseNew)) { databaseId = databaseNew.Id; - var client = databaseNew.Client; - endpoint = client.Endpoint; + try + { + endpoint = databaseNew.Client.Endpoint; + } + catch (ObjectDisposedException) + { + // Client disposed; leave endpoint null + } + catch (Exception ex) + { + Log.Warning(ex, "Unable to get Cosmos client endpoint from DatabaseNewStruct."); + } } else if (instance.TryDuckCast(out var databaseOld)) { databaseId = databaseOld.Id; - var client = databaseOld.ClientContext.Client; - endpoint = client.Endpoint; + try + { + endpoint = databaseOld.ClientContext.Client.Endpoint; + } + catch (ObjectDisposedException) + { + // Client disposed; leave endpoint null + } + catch (Exception ex) + { + Log.Warning(ex, "Unable to get Cosmos client endpoint from DatabaseOldStruct."); + } } return CreateCosmosDbCallState(instance, queryDefinition, null, databaseId, endpoint); @@ -78,7 +120,18 @@ public static CallTargetState CreateClientCallStateExt(out var c)) { - endpoint = c.Endpoint; + try + { + endpoint = c.Endpoint; + } + catch (ObjectDisposedException) + { + // Client disposed; leave endpoint null + } + catch (Exception ex) + { + Log.Warning(ex, "Unable to get Cosmos client endpoint from CosmosClientStruct."); + } } return CreateCosmosDbCallState(instance, queryDefinition, null, null, endpoint); diff --git a/tracer/test/Datadog.Trace.Tests/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommonTests.cs b/tracer/test/Datadog.Trace.Tests/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommonTests.cs new file mode 100644 index 000000000000..c68320851b5b --- /dev/null +++ b/tracer/test/Datadog.Trace.Tests/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommonTests.cs @@ -0,0 +1,132 @@ +// Copyright (c) Datadog +// Licensed under the Apache 2 License. + +using System; +using Datadog.Trace.ClrProfiler.AutoInstrumentation.CosmosDb; +using Datadog.Trace.ClrProfiler.CallTarget; +using FluentAssertions; +using Xunit; + +namespace Datadog.Trace.Tests.ClrProfiler.AutoInstrumentation.CosmosDb +{ +public class CosmosCommonTests +{ + [Fact] + public void CreateContainerCallState_DoesNotThrow_WhenDatabaseNewClientDisposed() + { + var container = new ContainerStub + { + Id = "container", + Database = new DatabaseNewStub { Id = "db" } + }; + + Action act = () => + { + CallTargetState state = CosmosCommon.CreateContainerCallStateExt(container, "SELECT 1"); + state.Scope?.Dispose(); + }; + + act.Should().NotThrow(); + } + + [Fact] + public void CreateContainerCallState_DoesNotThrow_WhenDatabaseOldClientDisposed() + { + var container = new ContainerStub + { + Id = "container", + Database = new DatabaseOldStub { Id = "db" } + }; + + Action act = () => + { + CallTargetState state = CosmosCommon.CreateContainerCallStateExt(container, "SELECT 1"); + state.Scope?.Dispose(); + }; + + act.Should().NotThrow(); + } + + [Fact] + public void CreateDatabaseCallState_DoesNotThrow_WhenClientDisposed_New() + { + var database = new DatabaseNewStub { Id = "db" }; + + Action act = () => + { + CallTargetState state = CosmosCommon.CreateDatabaseCallStateExt(database, "SELECT 1"); + state.Scope?.Dispose(); + }; + + act.Should().NotThrow(); + } + + [Fact] + public void CreateDatabaseCallState_DoesNotThrow_WhenClientDisposed_Old() + { + var database = new DatabaseOldStub { Id = "db" }; + + Action act = () => + { + CallTargetState state = CosmosCommon.CreateDatabaseCallStateExt(database, "SELECT 1"); + state.Scope?.Dispose(); + }; + + act.Should().NotThrow(); + } + + private class ContainerStub + { + public string Id + { + get; + set; + } + + // Must be object to match ContainerStruct.Database + public object Database + { + get; + set; + } + } + + private class DatabaseNewStub + { + public string Id + { + get; + set; + } + + // Simulate disposed client by throwing on getter access + public ThrowingCosmosClient Client => throw new ObjectDisposedException("CosmosClient"); + } + + private class DatabaseOldStub + { + public string Id + { + get; + set; + } + + public CosmosClientContextStub ClientContext + { + get; + } = new CosmosClientContextStub(); + } + + private class CosmosClientContextStub + { + // Simulate disposed client by throwing on getter access + public ThrowingCosmosClient Client => throw new ObjectDisposedException("CosmosClient"); + } + + private class ThrowingCosmosClient + { + // Not expected to be reached in these tests, but provided for completeness + public Uri Endpoint => new Uri("http://localhost"); + } +} +}