Skip to content

Commit 1583fd4

Browse files
samikshya-dbclaude
andcommitted
[PECOBLR-1146] Implement feature flag cache with reference counting
Implements per-host feature flag caching with reference counting for telemetry control. Key features: - Per-host singleton cache to prevent rate limiting - Reference counting tied to connection lifecycle - 15-minute TTL with automatic refresh - Thread-safe concurrent access using RWMutex - HTTP integration with Databricks feature flag API - Fallback to cached value on fetch errors Comprehensive test coverage includes: - Reference counting increment/decrement - Cache expiration and refresh logic - Concurrent access safety - HTTP fetch success/failure scenarios - Multiple host management 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
1 parent 56b8a73 commit 1583fd4

File tree

3 files changed

+643
-1
lines changed

3 files changed

+643
-1
lines changed

telemetry/DESIGN.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1743,7 +1743,7 @@ func BenchmarkInterceptor_Disabled(b *testing.B) {
17431743
- [x] Add unit tests for configuration and tags
17441744

17451745
### Phase 2: Per-Host Management
1746-
- [ ] Implement `featureflag.go` with caching and reference counting
1746+
- [x] Implement `featureflag.go` with caching and reference counting ✅ COMPLETED (PECOBLR-1146)
17471747
- [ ] Implement `manager.go` for client management
17481748
- [ ] Implement `circuitbreaker.go` with state machine
17491749
- [ ] Add unit tests for all components

telemetry/featureflag.go

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
package telemetry
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"net/http"
8+
"strings"
9+
"sync"
10+
"time"
11+
)
12+
13+
// featureFlagCache manages feature flag state per host with reference counting.
14+
// This prevents rate limiting by caching feature flag responses.
15+
type featureFlagCache struct {
16+
mu sync.RWMutex
17+
contexts map[string]*featureFlagContext
18+
}
19+
20+
// featureFlagContext holds feature flag state and reference count for a host.
21+
type featureFlagContext struct {
22+
enabled *bool
23+
lastFetched time.Time
24+
refCount int
25+
cacheDuration time.Duration
26+
}
27+
28+
var (
29+
flagCacheOnce sync.Once
30+
flagCacheInstance *featureFlagCache
31+
)
32+
33+
// getFeatureFlagCache returns the singleton instance.
34+
func getFeatureFlagCache() *featureFlagCache {
35+
flagCacheOnce.Do(func() {
36+
flagCacheInstance = &featureFlagCache{
37+
contexts: make(map[string]*featureFlagContext),
38+
}
39+
})
40+
return flagCacheInstance
41+
}
42+
43+
// getOrCreateContext gets or creates a feature flag context for the host.
44+
// Increments reference count.
45+
func (c *featureFlagCache) getOrCreateContext(host string) *featureFlagContext {
46+
c.mu.Lock()
47+
defer c.mu.Unlock()
48+
49+
ctx, exists := c.contexts[host]
50+
if !exists {
51+
ctx = &featureFlagContext{
52+
cacheDuration: 15 * time.Minute,
53+
}
54+
c.contexts[host] = ctx
55+
}
56+
ctx.refCount++
57+
return ctx
58+
}
59+
60+
// releaseContext decrements reference count for the host.
61+
// Removes context when ref count reaches zero.
62+
func (c *featureFlagCache) releaseContext(host string) {
63+
c.mu.Lock()
64+
defer c.mu.Unlock()
65+
66+
if ctx, exists := c.contexts[host]; exists {
67+
ctx.refCount--
68+
if ctx.refCount <= 0 {
69+
delete(c.contexts, host)
70+
}
71+
}
72+
}
73+
74+
// isTelemetryEnabled checks if telemetry is enabled for the host.
75+
// Uses cached value if available and not expired.
76+
func (c *featureFlagCache) isTelemetryEnabled(ctx context.Context, host string, httpClient *http.Client) (bool, error) {
77+
c.mu.RLock()
78+
flagCtx, exists := c.contexts[host]
79+
c.mu.RUnlock()
80+
81+
if !exists {
82+
return false, nil
83+
}
84+
85+
// Check if cache is valid
86+
if flagCtx.enabled != nil && time.Since(flagCtx.lastFetched) < flagCtx.cacheDuration {
87+
return *flagCtx.enabled, nil
88+
}
89+
90+
// Fetch fresh value
91+
enabled, err := fetchFeatureFlag(ctx, host, httpClient)
92+
if err != nil {
93+
// Return cached value on error, or false if no cache
94+
if flagCtx.enabled != nil {
95+
return *flagCtx.enabled, nil
96+
}
97+
return false, err
98+
}
99+
100+
// Update cache
101+
c.mu.Lock()
102+
flagCtx.enabled = &enabled
103+
flagCtx.lastFetched = time.Now()
104+
c.mu.Unlock()
105+
106+
return enabled, nil
107+
}
108+
109+
// isExpired returns true if the cache has expired.
110+
func (c *featureFlagContext) isExpired() bool {
111+
return c.enabled == nil || time.Since(c.lastFetched) > c.cacheDuration
112+
}
113+
114+
// fetchFeatureFlag fetches the feature flag value from the Databricks API.
115+
func fetchFeatureFlag(ctx context.Context, host string, httpClient *http.Client) (bool, error) {
116+
// Build endpoint URL, adding https:// if no scheme present
117+
endpoint := host
118+
if !strings.HasPrefix(host, "http://") && !strings.HasPrefix(host, "https://") {
119+
endpoint = "https://" + host
120+
}
121+
endpoint = endpoint + "/api/2.0/feature-flags"
122+
123+
req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil)
124+
if err != nil {
125+
return false, fmt.Errorf("failed to create request: %w", err)
126+
}
127+
128+
// Add query parameters
129+
q := req.URL.Query()
130+
q.Add("flags", "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForAdbc")
131+
req.URL.RawQuery = q.Encode()
132+
133+
resp, err := httpClient.Do(req)
134+
if err != nil {
135+
return false, fmt.Errorf("failed to fetch feature flag: %w", err)
136+
}
137+
defer resp.Body.Close()
138+
139+
if resp.StatusCode != http.StatusOK {
140+
return false, fmt.Errorf("feature flag check failed with status: %d", resp.StatusCode)
141+
}
142+
143+
var result struct {
144+
Flags map[string]bool `json:"flags"`
145+
}
146+
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
147+
return false, fmt.Errorf("failed to decode response: %w", err)
148+
}
149+
150+
return result.Flags["databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForAdbc"], nil
151+
}

0 commit comments

Comments
 (0)