prometheus-net

1
using System.Buffers;
2

3
namespace Prometheus;
4

5
/// <summary>
6
/// Represents a metric whose lifetime is managed by the caller, either via explicit leases or via extend-on-use behavior (implicit leases).
7
/// </summary>
8
/// <remarks>
9
/// Each metric handle maintains a reaper task that occasionally removes metrics that have expired. The reaper is started
10
/// when the first lifetime-managed metric is created and terminates when the last lifetime-managed metric expires.
11
/// This does mean that the metric handle may keep objects alive until expiration, even if the handle itself is no longer used.
12
/// </remarks>
13
internal abstract class ManagedLifetimeMetricHandle<TChild, TMetricInterface>
14
    : IManagedLifetimeMetricHandle<TMetricInterface>, INotifyLeaseEnded
15
    where TChild : ChildBase, TMetricInterface
16
    where TMetricInterface : ICollectorChild
17
{
18
    internal ManagedLifetimeMetricHandle(Collector<TChild> metric, TimeSpan expiresAfter)
19
    {
20
        _reaperFunc = Reaper;
21

22
        _metric = metric;
23
        _expiresAfter = expiresAfter;
24
    }
25

26
    protected readonly Collector<TChild> _metric;
27
    protected readonly TimeSpan _expiresAfter;
28

29
    #region Lease(string[])
30
    public IDisposable AcquireLease(out TMetricInterface metric, params string[] labelValues)
31
    {
32
        var child = _metric.WithLabels(labelValues);
33
        metric = child;
34

35
        return TakeLease(child);
36
    }
37

38
    public RefLease AcquireRefLease(out TMetricInterface metric, params string[] labelValues)
39
    {
40
        var child = _metric.WithLabels(labelValues);
41
        metric = child;
42

43
        return TakeRefLease(child);
44
    }
45

46
    public void WithLease(Action<TMetricInterface> action, params string[] labelValues)
47
    {
48
        var child = _metric.WithLabels(labelValues);
49
        using var lease = TakeRefLease(child);
50

51
        action(child);
52
    }
53

54
    public void WithLease<TArg>(Action<TArg, TMetricInterface> action, TArg arg, params string[] labelValues)
55
    {
56
        var child = _metric.WithLabels(labelValues);
57
        using var lease = TakeRefLease(child);
58

59
        action(arg, child);
60
    }
61

62
    public async Task WithLeaseAsync(Func<TMetricInterface, Task> action, params string[] labelValues)
63
    {
64
        using var lease = AcquireLease(out var metric, labelValues);
65
        await action(metric);
66
    }
67

68
    public TResult WithLease<TResult>(Func<TMetricInterface, TResult> func, params string[] labelValues)
69
    {
70
        using var lease = AcquireLease(out var metric, labelValues);
71
        return func(metric);
72
    }
73

74
    public async Task<TResult> WithLeaseAsync<TResult>(Func<TMetricInterface, Task<TResult>> func, params string[] labelValues)
75
    {
76
        using var lease = AcquireLease(out var metric, labelValues);
77
        return await func(metric);
78
    }
79
    #endregion
80

81
    #region Lease(ReadOnlyMemory<string>)
82
    public IDisposable AcquireLease(out TMetricInterface metric, ReadOnlyMemory<string> labelValues)
83
    {
84
        var child = _metric.WithLabels(labelValues);
85
        metric = child;
86

87
        return TakeLease(child);
88
    }
89

90
    public RefLease AcquireRefLease(out TMetricInterface metric, ReadOnlyMemory<string> labelValues)
91
    {
92
        var child = _metric.WithLabels(labelValues);
93
        metric = child;
94

95
        return TakeRefLease(child);
96
    }
97

98
    public void WithLease(Action<TMetricInterface> action, ReadOnlyMemory<string> labelValues)
99
    {
100
        var child = _metric.WithLabels(labelValues);
101
        using var lease = TakeRefLease(child);
102

103
        action(child);
104
    }
105

106
    public void WithLease<TArg>(Action<TArg, TMetricInterface> action, TArg arg, ReadOnlyMemory<string> labelValues)
107
    {
108
        var child = _metric.WithLabels(labelValues);
109
        using var lease = TakeRefLease(child);
110

111
        action(arg, child);
112
    }
113

114
    public async Task WithLeaseAsync(Func<TMetricInterface, Task> action, ReadOnlyMemory<string> labelValues)
115
    {
116
        using var lease = AcquireLease(out var metric, labelValues);
117
        await action(metric);
118
    }
119

120
    public TResult WithLease<TResult>(Func<TMetricInterface, TResult> func, ReadOnlyMemory<string> labelValues)
121
    {
122
        using var lease = AcquireLease(out var metric, labelValues);
123
        return func(metric);
124
    }
125

126
    public async Task<TResult> WithLeaseAsync<TResult>(Func<TMetricInterface, Task<TResult>> func, ReadOnlyMemory<string> labelValues)
127
    {
128
        using var lease = AcquireLease(out var metric, labelValues);
129
        return await func(metric);
130
    }
131
    #endregion
132

133
    #region Lease(ReadOnlySpan<string>)
134
    public IDisposable AcquireLease(out TMetricInterface metric, ReadOnlySpan<string> labelValues)
135
    {
136
        var child = _metric.WithLabels(labelValues);
137
        metric = child;
138

139
        return TakeLease(child);
140
    }
141

142
    public RefLease AcquireRefLease(out TMetricInterface metric, ReadOnlySpan<string> labelValues)
143
    {
144
        var child = _metric.WithLabels(labelValues);
145
        metric = child;
146

147
        return TakeRefLease(child);
148
    }
149

150
    public void WithLease(Action<TMetricInterface> action, ReadOnlySpan<string> labelValues)
151
    {
152
        var child = _metric.WithLabels(labelValues);
153
        using var lease = TakeRefLease(child);
154

155
        action(child);
156
    }
157

158
    public void WithLease<TArg>(Action<TArg, TMetricInterface> action, TArg arg, ReadOnlySpan<string> labelValues)
159
    {
160
        var child = _metric.WithLabels(labelValues);
161
        using var lease = TakeRefLease(child);
162

163
        action(arg, child);
164
    }
165

166
    public TResult WithLease<TResult>(Func<TMetricInterface, TResult> func, ReadOnlySpan<string> labelValues)
167
    {
168
        using var lease = AcquireLease(out var metric, labelValues);
169
        return func(metric);
170
    }
171
    #endregion
172

173
    public abstract ICollector<TMetricInterface> WithExtendLifetimeOnUse();
174

175
    /// <summary>
176
    /// Internal to allow the delay logic to be replaced in test code, enabling (non-)expiration on demand.
177
    /// </summary>
178
    internal IDelayer Delayer = RealDelayer.Instance;
179

180
    #region Lease tracking
181
    private readonly Dictionary<TChild, ChildLifetimeInfo> _lifetimes = new();
182

183
    // Guards the collection but not the contents.
184
    private readonly ReaderWriterLockSlim _lifetimesLock = new();
185

186
    private bool HasAnyTrackedLifetimes()
187
    {
188
        _lifetimesLock.EnterReadLock();
189

190
        try
191
        {
192
            return _lifetimes.Count != 0;
193
        }
194
        finally
195
        {
196
            _lifetimesLock.ExitReadLock();
197
        }
198
    }
199

200
    /// <summary>
201
    /// For testing only. Sets all keepalive timestamps to a time in the disstant past,
202
    /// which will cause all lifetimes to expire (if they have no leases).
203
    /// </summary>
204
    internal void SetAllKeepaliveTimestampsToDistantPast()
205
    {
206
        // We cannot just zero this because zero is the machine start timestamp, so zero is not necessarily
207
        // far in the past (especially if the machine is a build agent that just started up). 1 year negative should work, though.
208
        var distantPast = -PlatformCompatibilityHelpers.ElapsedToTimeStopwatchTicks(TimeSpan.FromDays(365));
209

210
        _lifetimesLock.EnterReadLock();
211

212
        try
213
        {
214
            foreach (var lifetime in _lifetimes.Values)
215
                Volatile.Write(ref lifetime.KeepaliveTimestamp, distantPast);
216
        }
217
        finally
218
        {
219
            _lifetimesLock.ExitReadLock();
220
        }
221
    }
222

223
    /// <summary>
224
    /// For anomaly analysis during testing only.
225
    /// </summary>
226
    internal void DebugDumpLifetimes()
227
    {
228
        _lifetimesLock.EnterReadLock();
229

230
        try
231
        {
232
            Console.WriteLine($"Dumping {_lifetimes.Count} lifetimes of {_metric}. Reaper status: {Volatile.Read(ref _reaperActiveBool)}.");
233

234
            foreach (var pair in _lifetimes)
235
            {
236
                Console.WriteLine($"{pair.Key} -> {pair.Value}");
237
            }
238
        }
239
        finally
240
        {
241
            _lifetimesLock.ExitReadLock();
242
        }
243
    }
244

245
    private IDisposable TakeLease(TChild child)
246
    {
247
        var lifetime = GetOrCreateLifetimeAndIncrementLeaseCount(child);
248
        EnsureReaperActive();
249

250
        return new Lease(this, child, lifetime);
251
    }
252

253
    private RefLease TakeRefLease(TChild child)
254
    {
255
        var lifetime = GetOrCreateLifetimeAndIncrementLeaseCount(child);
256
        EnsureReaperActive();
257

258
        return new RefLease(this, child, lifetime);
259
    }
260

261
    private ChildLifetimeInfo GetOrCreateLifetimeAndIncrementLeaseCount(TChild child)
262
    {
263
        _lifetimesLock.EnterReadLock();
264

265
        try
266
        {
267
            // Ideally, there already exists a registered lifetime for this metric instance.
268
            if (_lifetimes.TryGetValue(child, out var existing))
269
            {
270
                // Immediately increment it, to reduce the risk of any concurrent activities ending the lifetime.
271
                Interlocked.Increment(ref existing.LeaseCount);
272
                return existing;
273
            }
274
        }
275
        finally
276
        {
277
            _lifetimesLock.ExitReadLock();
278
        }
279

280
        // No lifetime registered yet - we need to take a write lock and register it.
281
        var newLifetime = new ChildLifetimeInfo
282
        {
283
            LeaseCount = 1
284
        };
285

286
        _lifetimesLock.EnterWriteLock();
287

288
        try
289
        {
290
#if NET
291
            // It could be that someone beats us to it! Probably not, though.
292
            if (_lifetimes.TryAdd(child, newLifetime))
293
                return newLifetime;
294

295
            var existing = _lifetimes[child];
296

297
            // Immediately increment it, to reduce the risk of any concurrent activities ending the lifetime.
298
            // Even if something does, it is not the end of the world - the reaper will create a new lifetime when it realizes this happened.
299
            Interlocked.Increment(ref existing.LeaseCount);
300
            return existing;
301
#else
302
            // On .NET Fx we need to do the pessimistic case first because there is no TryAdd().
303
            if (_lifetimes.TryGetValue(child, out var existing))
304
            {
305
                // Immediately increment it, to reduce the risk of any concurrent activities ending the lifetime.
306
                // Even if something does, it is not the end of the world - the reaper will create a new lifetime when it realizes this happened.
307
                Interlocked.Increment(ref existing.LeaseCount);
308
                return existing;
309
            }
310

311
            _lifetimes.Add(child, newLifetime);
312
            return newLifetime;
313
#endif
314
        }
315
        finally
316
        {
317
            _lifetimesLock.ExitWriteLock();
318
        }
319
    }
320

321
    internal void OnLeaseEnded(TChild child, ChildLifetimeInfo lifetime)
322
    {
323
        // Update keepalive timestamp before anything else, to avoid racing.
324
        Volatile.Write(ref lifetime.KeepaliveTimestamp, LowGranularityTimeSource.GetStopwatchTimestamp());
325

326
        // If the lifetime has been ended while we still held a lease, it means there was a race that we lost.
327
        // The metric instance may or may not be still alive. To ensure proper cleanup, we re-register a lifetime
328
        // for the metric instance, which will ensure it gets cleaned up when it expires.
329
        if (Volatile.Read(ref lifetime.Ended))
330
        {
331
            // We just take a new lease and immediately dispose it. We are guaranteed not to loop here because the
332
            // reaper removes lifetimes from the dictionary once ended, so we can never run into the same lifetime again.
333
            TakeRefLease(child).Dispose();
334
        }
335

336
        // Finally, decrement the lease count to relinquish any claim on extending the lifetime.
337
        Interlocked.Decrement(ref lifetime.LeaseCount);
338
    }
339

340
    void INotifyLeaseEnded.OnLeaseEnded(object child, ChildLifetimeInfo lifetime)
341
    {
342
        OnLeaseEnded((TChild)child, lifetime);
343
    }
344

345
    private sealed class Lease(ManagedLifetimeMetricHandle<TChild, TMetricInterface> parent, TChild child, ChildLifetimeInfo lifetime) : IDisposable
346
    {
347
        public void Dispose() => parent.OnLeaseEnded(child, lifetime);
348
    }
349
#endregion
350

351
    #region Reaper
352
    // Whether the reaper is currently active. This is set to true when a metric instance is created and
353
    // reset when the last metric instance expires (after which it may be set again).
354
    // We use atomic operations without locking.
355
    private int _reaperActiveBool = ReaperInactive;
356

357
    private const int ReaperActive = 1;
358
    private const int ReaperInactive = 0;
359

360
    /// <summary>
361
    /// Call this immediately after creating a metric instance that will eventually expire.
362
    /// </summary>
363
    private void EnsureReaperActive()
364
    {
365
        if (Interlocked.CompareExchange(ref _reaperActiveBool, ReaperActive, ReaperInactive) == ReaperActive)
366
        {
367
            // It was already active - nothing for us to do.
368
            return;
369
        }
370

371
        _ = Task.Run(_reaperFunc);
372
    }
373

374
    private async Task Reaper()
375
    {
376
        while (true)
377
        {
378
            var now = LowGranularityTimeSource.GetStopwatchTimestamp();
379

380
            // Will contains the results of pass 1.
381
            TChild[] expiredInstancesBuffer = null!;
382
            int expiredInstanceCount = 0;
383

384
            // Pass 1: holding only a read lock, make a list of metric instances that have expired.
385
            _lifetimesLock.EnterReadLock();
386

387
            try
388
            {
389
                try
390
                {
391
                    expiredInstancesBuffer = ArrayPool<TChild>.Shared.Rent(_lifetimes.Count);
392

393
                    foreach (var pair in _lifetimes)
394
                    {
395
                        if (Volatile.Read(ref pair.Value.LeaseCount) != 0)
396
                            continue; // Not expired.
397

398
                        if (PlatformCompatibilityHelpers.StopwatchGetElapsedTime(Volatile.Read(ref pair.Value.KeepaliveTimestamp), now) < _expiresAfter)
399
                            continue; // Not expired.
400

401
                        // No leases and keepalive has expired - it is an expired instance!
402
                        expiredInstancesBuffer[expiredInstanceCount++] = pair.Key;
403
                    }
404
                }
405
                finally
406
                {
407
                    _lifetimesLock.ExitReadLock();
408
                }
409

410
                // Pass 2: if we have any work to do, take a write lock and remove the expired metric instances,
411
                // assuming our judgement about their expiration remains valid. We process and lock one by one,
412
                // to avoid holding locks for a long duration if many items expire at once - we are not in any rush.
413
                for (var i = 0; i < expiredInstanceCount; i++)
414
                {
415
                    var expiredInstance = expiredInstancesBuffer[i];
416

417
                    _lifetimesLock.EnterWriteLock();
418

419
                    try
420
                    {
421
                        if (!_lifetimes.TryGetValue(expiredInstance, out var lifetime))
422
                            continue; // Already gone, nothing for us to do.
423

424
                        // We need to check again whether the metric instance is still expired, because it may have been
425
                        // renewed by a new lease in the meantime. If it is still expired, we can remove it.
426
                        if (Volatile.Read(ref lifetime.LeaseCount) != 0)
427
                            continue; // Not expired.
428

429
                        if (PlatformCompatibilityHelpers.StopwatchGetElapsedTime(Volatile.Read(ref lifetime.KeepaliveTimestamp), now) < _expiresAfter)
430
                            continue; // Not expired.
431

432
                        // No leases and keepalive has expired - it is an expired instance!
433

434
                        // We mark the old lifetime as ended - if it happened that it got associated with a new lease
435
                        // (which is possible because we do not prevent lease-taking while in this loop), the new lease
436
                        // upon being ended will re-register the lifetime instead of just extending the existing one.
437
                        // We can be certain that any concurrent lifetime-affecting logic is using the same LifetimeInfo
438
                        // instance because the lifetime dictionary remains locked until we are done (by which time this flag is set).
439
                        Volatile.Write(ref lifetime.Ended, true);
440

441
                        _lifetimes.Remove(expiredInstance);
442

443
                        // If we did encounter a race, removing the metric instance here means that some metric value updates
444
                        // may go missing (until the next lease creates a new instance). This is acceptable behavior, to keep the code simple.
445
                        expiredInstance.Remove();
446
                    }
447
                    finally
448
                    {
449
                        _lifetimesLock.ExitWriteLock();
450
                    }
451
                }
452
            }
453
            finally
454
            {
455
                ArrayPool<TChild>.Shared.Return(expiredInstancesBuffer);
456
            }
457

458
            // Check if we need to shut down the reaper or keep going.
459
            _lifetimesLock.EnterReadLock();
460

461
            try
462
            {
463
                if (_lifetimes.Count != 0)
464
                    goto has_more_work;
465
            }
466
            finally
467
            {
468
                _lifetimesLock.ExitReadLock();
469
            }
470

471
            CleanupReaper();
472
            return;
473

474
        has_more_work:
475
            // Work done! Go sleep a bit and come back when something may have expired.
476
            // We do not need to be too aggressive here, as expiration is not a hard schedule guarantee.
477
            await Delayer.Delay(_expiresAfter);
478
        }
479
    }
480

481
    /// <summary>
482
    /// Called when the reaper has noticed that all metric instances have expired and it has no more work to do. 
483
    /// </summary>
484
    private void CleanupReaper()
485
    {
486
        Volatile.Write(ref _reaperActiveBool, ReaperInactive);
487

488
        // The reaper is now gone. However, as we do not use locking here it is possible that someone already
489
        // added metric instances (which saw "oh reaper is still running") before we got here. Let's check - if
490
        // there appear to be metric instances registered, we may need to start the reaper again.
491
        if (HasAnyTrackedLifetimes())
492
            EnsureReaperActive();
493
    }
494

495
    private readonly Func<Task> _reaperFunc;
496
    #endregion
497
}