Skip to content

Handle disposed Cosmos clients gracefully #7700

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Handle disposed Cosmos clients gracefully #7700

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 60 additions & 7 deletions tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommon.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,34 @@ public static CallTargetState CreateContainerCallStateExt<TTarget, TQueryDefinit
if (container.Database.TryDuckCast<DatabaseNewStruct>(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<DatabaseOldStruct>(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.");
}
}
}

Expand All @@ -60,14 +82,34 @@ public static CallTargetState CreateDatabaseCallStateExt<TTarget, TQueryDefiniti
if (instance.TryDuckCast<DatabaseNewStruct>(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<DatabaseOldStruct>(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);
Expand All @@ -78,7 +120,18 @@ public static CallTargetState CreateClientCallStateExt<TTarget, TQueryDefinition
Uri endpoint = null;
if (instance.TryDuckCast<CosmosClientStruct>(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);
Expand Down
132 changes: 132 additions & 0 deletions ...er/test/Datadog.Trace.Tests/ClrProfiler/AutoInstrumentation/CosmosDb/CosmosCommonTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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");
}
}
}
Loading