Skip to content

Commit b1a8b39

Browse files
Jade Wangclaude
andcommitted
feat(csharp): implement FeatureFlagCache (WI-3.1)
Implements per-host feature flag caching with reference counting to avoid repeated API calls and rate limiting. Key features: - FeatureFlagContext: Holds cached telemetry enabled state, last fetched timestamp, reference count, and configurable cache duration (default 15 min) - FeatureFlagCache: Singleton managing per-host contexts with thread-safe ConcurrentDictionary storage API: - GetInstance(): Returns the singleton instance - GetOrCreateContext(host): Creates/returns context and increments RefCount - ReleaseContext(host): Decrements RefCount, removes context when zero - IsTelemetryEnabledAsync(): Returns cached value if valid, otherwise fetches Thread safety ensured via ConcurrentDictionary and Interlocked operations. Includes 46 comprehensive unit tests covering all exit criteria. Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 2d85d86 commit b1a8b39

File tree

3 files changed

+1274
-0
lines changed

3 files changed

+1274
-0
lines changed
Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,268 @@
1+
/*
2+
* Copyright (c) 2025 ADBC Drivers Contributors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
using System;
18+
using System.Collections.Concurrent;
19+
using System.Diagnostics;
20+
using System.Net.Http;
21+
using System.Threading;
22+
using System.Threading.Tasks;
23+
24+
namespace AdbcDrivers.Databricks.Telemetry
25+
{
26+
/// <summary>
27+
/// Singleton that manages feature flag cache per host.
28+
/// Prevents rate limiting by caching feature flag responses.
29+
/// </summary>
30+
/// <remarks>
31+
/// This class implements the per-host caching pattern from the JDBC driver:
32+
/// - Feature flags are cached by host to prevent rate limiting
33+
/// - Reference counting tracks number of connections per host
34+
/// - Cache is automatically cleaned up when all connections to a host close
35+
/// - Thread-safe using ConcurrentDictionary
36+
///
37+
/// JDBC Reference: DatabricksDriverFeatureFlagsContextFactory.java
38+
/// </remarks>
39+
internal sealed class FeatureFlagCache
40+
{
41+
private static readonly FeatureFlagCache s_instance = new FeatureFlagCache();
42+
43+
private readonly ConcurrentDictionary<string, FeatureFlagContext> _contexts;
44+
private readonly TimeSpan _defaultCacheDuration;
45+
46+
/// <summary>
47+
/// Gets the singleton instance of the FeatureFlagCache.
48+
/// </summary>
49+
public static FeatureFlagCache GetInstance() => s_instance;
50+
51+
/// <summary>
52+
/// Creates a new FeatureFlagCache with default cache duration (15 minutes).
53+
/// </summary>
54+
internal FeatureFlagCache()
55+
: this(FeatureFlagContext.DefaultCacheDuration)
56+
{
57+
}
58+
59+
/// <summary>
60+
/// Creates a new FeatureFlagCache with the specified default cache duration.
61+
/// </summary>
62+
/// <param name="defaultCacheDuration">The default cache duration for new contexts.</param>
63+
internal FeatureFlagCache(TimeSpan defaultCacheDuration)
64+
{
65+
if (defaultCacheDuration <= TimeSpan.Zero)
66+
{
67+
throw new ArgumentOutOfRangeException(nameof(defaultCacheDuration), "Cache duration must be greater than zero.");
68+
}
69+
70+
_contexts = new ConcurrentDictionary<string, FeatureFlagContext>(StringComparer.OrdinalIgnoreCase);
71+
_defaultCacheDuration = defaultCacheDuration;
72+
}
73+
74+
/// <summary>
75+
/// Gets or creates a feature flag context for the host.
76+
/// Increments reference count.
77+
/// </summary>
78+
/// <param name="host">The host (Databricks workspace URL) to get or create a context for.</param>
79+
/// <returns>The feature flag context for the host.</returns>
80+
/// <exception cref="ArgumentException">Thrown when host is null or whitespace.</exception>
81+
public FeatureFlagContext GetOrCreateContext(string host)
82+
{
83+
if (string.IsNullOrWhiteSpace(host))
84+
{
85+
throw new ArgumentException("Host cannot be null or whitespace.", nameof(host));
86+
}
87+
88+
var context = _contexts.GetOrAdd(host, _ => new FeatureFlagContext(_defaultCacheDuration));
89+
context.IncrementRefCount();
90+
91+
Debug.WriteLine($"[TRACE] FeatureFlagCache: GetOrCreateContext for host '{host}', RefCount={context.RefCount}");
92+
93+
return context;
94+
}
95+
96+
/// <summary>
97+
/// Decrements reference count for the host.
98+
/// Removes context when ref count reaches zero.
99+
/// </summary>
100+
/// <param name="host">The host to release the context for.</param>
101+
/// <remarks>
102+
/// This method is thread-safe. If the reference count reaches zero,
103+
/// the context is removed from the cache. If multiple threads try to
104+
/// release the same context simultaneously, only one will successfully
105+
/// remove it.
106+
/// </remarks>
107+
public void ReleaseContext(string host)
108+
{
109+
if (string.IsNullOrWhiteSpace(host))
110+
{
111+
return;
112+
}
113+
114+
if (_contexts.TryGetValue(host, out var context))
115+
{
116+
var newRefCount = context.DecrementRefCount();
117+
Debug.WriteLine($"[TRACE] FeatureFlagCache: ReleaseContext for host '{host}', RefCount={newRefCount}");
118+
119+
if (newRefCount <= 0)
120+
{
121+
// Try to remove the context. Use TryRemove with the specific value
122+
// to avoid race conditions where a new connection added a reference.
123+
if (context.RefCount <= 0)
124+
{
125+
// Note: We check RefCount again because another thread might have
126+
// incremented it between our check and the removal attempt.
127+
#if NET5_0_OR_GREATER
128+
_contexts.TryRemove(new System.Collections.Generic.KeyValuePair<string, FeatureFlagContext>(host, context));
129+
#else
130+
// For netstandard2.0, we need to be more careful about the removal
131+
// to avoid race conditions.
132+
if (_contexts.TryGetValue(host, out var currentContext) && currentContext == context && currentContext.RefCount <= 0)
133+
{
134+
((System.Collections.Generic.IDictionary<string, FeatureFlagContext>)_contexts).Remove(new System.Collections.Generic.KeyValuePair<string, FeatureFlagContext>(host, context));
135+
}
136+
#endif
137+
Debug.WriteLine($"[TRACE] FeatureFlagCache: Removed context for host '{host}'");
138+
}
139+
}
140+
}
141+
}
142+
143+
/// <summary>
144+
/// Checks if telemetry is enabled for the host.
145+
/// Uses cached value if available and not expired.
146+
/// </summary>
147+
/// <param name="host">The host to check telemetry status for.</param>
148+
/// <param name="featureFlagFetcher">Function to fetch the feature flag from the server.</param>
149+
/// <param name="ct">Cancellation token.</param>
150+
/// <returns>True if telemetry is enabled, false otherwise.</returns>
151+
/// <remarks>
152+
/// This method:
153+
/// 1. Returns the cached value if available and not expired
154+
/// 2. Otherwise fetches the feature flag using the provided fetcher
155+
/// 3. Caches the result for future calls
156+
///
157+
/// All exceptions from the fetcher are caught and logged at TRACE level.
158+
/// On error, returns false (telemetry disabled) as a safe default.
159+
/// </remarks>
160+
public async Task<bool> IsTelemetryEnabledAsync(
161+
string host,
162+
Func<CancellationToken, Task<bool>> featureFlagFetcher,
163+
CancellationToken ct = default)
164+
{
165+
if (string.IsNullOrWhiteSpace(host))
166+
{
167+
return false;
168+
}
169+
170+
if (featureFlagFetcher == null)
171+
{
172+
return false;
173+
}
174+
175+
try
176+
{
177+
if (!_contexts.TryGetValue(host, out var context))
178+
{
179+
// No context for this host, return false
180+
return false;
181+
}
182+
183+
// Check if we have a valid cached value
184+
if (context.TryGetCachedValue(out bool cachedValue))
185+
{
186+
Debug.WriteLine($"[TRACE] FeatureFlagCache: Using cached value for host '{host}': {cachedValue}");
187+
return cachedValue;
188+
}
189+
190+
// Cache miss or expired - fetch from server
191+
Debug.WriteLine($"[TRACE] FeatureFlagCache: Cache miss for host '{host}', fetching from server");
192+
var enabled = await featureFlagFetcher(ct).ConfigureAwait(false);
193+
194+
// Update the cache
195+
context.SetTelemetryEnabled(enabled);
196+
Debug.WriteLine($"[TRACE] FeatureFlagCache: Updated cache for host '{host}': {enabled}");
197+
198+
return enabled;
199+
}
200+
catch (OperationCanceledException)
201+
{
202+
// Don't swallow cancellation
203+
throw;
204+
}
205+
catch (Exception ex)
206+
{
207+
// Swallow all other exceptions per telemetry requirement
208+
// Log at TRACE level to avoid customer anxiety
209+
Debug.WriteLine($"[TRACE] FeatureFlagCache: Error fetching feature flag for host '{host}': {ex.Message}");
210+
return false;
211+
}
212+
}
213+
214+
/// <summary>
215+
/// Gets the number of hosts currently cached.
216+
/// </summary>
217+
internal int CachedHostCount => _contexts.Count;
218+
219+
/// <summary>
220+
/// Checks if a context exists for the specified host.
221+
/// </summary>
222+
/// <param name="host">The host to check.</param>
223+
/// <returns>True if a context exists, false otherwise.</returns>
224+
internal bool HasContext(string host)
225+
{
226+
if (string.IsNullOrWhiteSpace(host))
227+
{
228+
return false;
229+
}
230+
231+
return _contexts.ContainsKey(host);
232+
}
233+
234+
/// <summary>
235+
/// Gets the context for the specified host, if it exists.
236+
/// Does not create a new context or modify reference count.
237+
/// </summary>
238+
/// <param name="host">The host to get the context for.</param>
239+
/// <param name="context">The context if found, null otherwise.</param>
240+
/// <returns>True if the context was found, false otherwise.</returns>
241+
internal bool TryGetContext(string host, out FeatureFlagContext? context)
242+
{
243+
context = null;
244+
245+
if (string.IsNullOrWhiteSpace(host))
246+
{
247+
return false;
248+
}
249+
250+
if (_contexts.TryGetValue(host, out var foundContext))
251+
{
252+
context = foundContext;
253+
return true;
254+
}
255+
256+
return false;
257+
}
258+
259+
/// <summary>
260+
/// Clears all cached contexts.
261+
/// This is primarily for testing purposes.
262+
/// </summary>
263+
internal void Clear()
264+
{
265+
_contexts.Clear();
266+
}
267+
}
268+
}

0 commit comments

Comments
 (0)