1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
|
//! Memory quota tracker, core and low-level API
//!
//! # Example
//!
//! ```
//! use std::{collections::VecDeque, sync::{Arc, Mutex}};
//! use tor_rtcompat::{CoarseInstant, CoarseTimeProvider, PreferredRuntime};
//! use tor_memquota::{mtracker, MemoryQuotaTracker, MemoryReclaimedError};
//! use void::{ResultVoidExt, Void};
//!
//! #[derive(Debug)]
//! struct TrackingQueue(Mutex<Result<Inner, MemoryReclaimedError>>);
//! #[derive(Debug)]
//! struct Inner {
//! partn: mtracker::Participation,
//! data: VecDeque<(Box<[u8]>, CoarseInstant)>,
//! }
//!
//! impl TrackingQueue {
//! fn push(&self, now: CoarseInstant, bytes: Box<[u8]>) -> Result<(), MemoryReclaimedError> {
//! let mut inner = self.0.lock().unwrap();
//! let inner = inner.as_mut().map_err(|e| e.clone())?;
//! inner.partn.claim(bytes.len())?;
//! inner.data.push_back((bytes, now));
//! Ok(())
//! }
//! }
//!
//! impl mtracker::IsParticipant for TrackingQueue {
//! fn get_oldest(&self) -> Option<CoarseInstant> {
//! let inner = self.0.lock().unwrap();
//! Some(inner.as_ref().ok()?.data.front()?.1)
//! }
//! fn reclaim(self: Arc<Self>) -> mtracker::ReclaimFuture {
//! let mut inner = self.0.lock().unwrap();
//! *inner = Err(MemoryReclaimedError::new());
//! Box::pin(async { mtracker::Reclaimed::Collapsing })
//! }
//! }
//!
//! let runtime = PreferredRuntime::create().unwrap();
//! let config = tor_memquota::Config::builder().max(1024*1024*1024).build().unwrap();
//! let trk = MemoryQuotaTracker::new(&runtime, config).unwrap();
//! let account = trk.new_account(None).unwrap();
//!
//! let queue: Arc<TrackingQueue> = account.register_participant_with(
//! runtime.now_coarse(),
//! |partn| {
//! Ok::<_, Void>(Arc::new(TrackingQueue(Mutex::new(Ok(Inner {
//! partn,
//! data: VecDeque::new(),
//! })))))
//! },
//! ).unwrap().void_unwrap();
//!
//! queue.push(runtime.now_coarse(), Box::new([0; 24])).unwrap();
//! ```
//
// For key internal documentation about the data structure, see the doc comment for
// `struct State` (down in the middle of the file).
use crate::internal_prelude::*;
mod bookkeeping;
mod reclaim;
mod total_qty_notifier;
#[cfg(test)]
mod test;
use bookkeeping::{BookkeepableQty, ClaimedQty, ParticipQty, TotalQty};
use total_qty_notifier::TotalQtyNotifier;
/// Maximum amount we'll "cache" locally in a [`Participation`]
///
/// ie maximum value of `Participation.cache`.
//
// TODO is this a good amount? should it be configurable?
pub(crate) const MAX_CACHE: Qty = Qty(16384);
/// Target cache size when we seem to be claiming
const TARGET_CACHE_CLAIMING: Qty = Qty(MAX_CACHE.as_usize() * 3 / 4);
/// Target cache size when we seem to be releasing
#[allow(clippy::identity_op)] // consistency
const TARGET_CACHE_RELEASING: Qty = Qty(MAX_CACHE.as_usize() * 1 / 4);
//---------- public data types ----------
/// Memory data tracker
///
/// Instance of the memory quota system.
///
/// Usually found as `Arc<MemoryQuotaTracker>`.
#[derive(Debug)]
pub struct MemoryQuotaTracker {
/// The actual tracker state etc.
state: Mutex<State>,
}
/// Handle onto an Account
///
/// An `Account` is a handle. All clones refer to the same underlying conceptual Account.
///
/// `Account`s are created using [`MemoryQuotaTracker::new_account`].
#[derive(Educe)]
#[educe(Debug)]
pub struct Account {
/// The account ID
aid: refcount::Ref<AId>,
/// The underlying tracker
#[educe(Debug(ignore))]
tracker: Arc<MemoryQuotaTracker>,
}
/// Weak handle onto an Account
///
/// Like [`Account`], but doesn't keep the account alive.
/// Must be upgraded before use.
//
// Doesn't count for ARecord.account_clones
//
// We can't lift out Arc, so that the caller sees `Arc<Account>`,
// because an Account is Arc<MemoryQuotaTracker> plus AId,
// not Arc of something account-specific.
#[derive(Clone, Educe)]
#[educe(Debug)]
pub struct WeakAccount {
/// The account ID
aid: AId,
/// The underlying tracker
#[educe(Debug(ignore))]
tracker: Weak<MemoryQuotaTracker>,
}
/// Handle onto a participant's participation in a tracker
///
/// `Participation` is a handle. All clones are for use by the same conceptual Participant.
/// It doesn't keep the underlying Account alive.
///
/// `Participation`s are created by registering new participants,
/// for example using [`Account::register_participant`].
///
/// Variables of this type are often named `partn`.
#[derive(Debug)]
pub struct Participation {
/// Participant id
pid: refcount::Ref<PId>,
/// Account id
aid: AId,
/// The underlying tracker
tracker: Weak<MemoryQuotaTracker>,
/// Quota we have preemptively claimed for use by this Account
///
/// Has been added to `PRecord.used`,
/// but not yet returned by `Participation::claim`.
///
/// This cache field arranges that most of the time we don't have to hammer a
/// single cache line.
///
/// The value here is bounded by a configured limit.
///
/// Invariants on memory accounting:
///
/// * `Participation.cache < configured limit`
/// * `PRecord.used = Participation.cache + Σ Participation::claim - Σ P'n::release`
/// except if `PRecord` has been deleted
/// (ie when we aren't tracking any more and think the Participant is `Collapsing`).
/// * `Σ PRecord.used = State.total_used`
///
/// Enforcement of these invariants is partially assured by
/// types in [`bookkeeping`].
cache: ClaimedQty,
}
/// Participants provide an impl of the hooks in this trait
///
/// Trait implemented by client of the memtrack API.
///
/// # Panic handling, "unwind safety"
///
/// If these methods panic, the memory tracker will tear down its records of the
/// participant, preventing future allocations.
///
/// But, it's not guaranteed that these methods on `IsParticipant` won't be called again,
/// even if they have already panicked on a previous occasion.
/// Thus the implementations might see "broken invariants"
/// as discussed in the docs for `std::panic::UnwindSafe`.
///
/// Nevertheless we don't make `RefUnwindSafe` a supertrait of `IsParticipant`.
/// That would force the caller to mark *all* their methods unwind-safe,
/// which is unreasonable (and probably undesirable).
///
/// Variables which are `IsParticipant` are often named `particip`.
pub trait IsParticipant: Debug + Send + Sync + 'static {
/// Return the age of the oldest data held by this Participant
///
/// `None` means this Participant holds no data.
///
/// # Performance and reentrancy
///
/// This function runs with the `MemoryQuotaTracker`'s internal global lock held.
/// Therefore:
///
/// * It must be fast.
/// * it *must not* call back into methods from [`tracker`](crate::mtracker).
/// * It *must not* even `Clone` or `Drop` a [`MemoryQuotaTracker`],
/// [`Account`], or [`Participation`].
fn get_oldest(&self) -> Option<CoarseInstant>;
/// Start memory reclamation
///
/// The Participant should start to free all of its memory,
/// and then return `Reclaimed::Collapsing`.
//
// In the future:
//
// Should free *at least* all memory at least as old as discard_...
//
// v1 of the actual implementation might not have `discard_everything_as_old_as`
// and `but_can_stop_discarding_...`,
// and might therefore only support Reclaimed::Collapsing
fn reclaim(
self: Arc<Self>,
// Future:
// discard_everything_as_old_as_this: RoughTime,
// but_can_stop_discarding_after_freeing_this_much: Qty,
) -> ReclaimFuture;
}
/// Future returned by the [`IsParticipant::reclaim`] reclamation request
pub type ReclaimFuture = Pin<Box<dyn Future<Output = Reclaimed> + Send + Sync>>;
/// Outcome of [`IsParticipant::reclaim`]
#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
#[non_exhaustive]
pub enum Reclaimed {
/// Participant is responding to reclamation by collapsing completely.
///
/// All memory will be freed and `release`'d soon (if it hasn't been already).
/// `MemoryQuotaTracker` should forget the Participant and all memory it used, right away.
///
/// Currently this is the only supported behaviour.
Collapsing,
// Future:
// /// Participant has now reclaimed some memory as instructed
// ///
// /// If this is not sufficient, tracker must call reclaim() again.
// /// (We may not want to implement Partial right away but the API
// /// ought to support it so let's think about it now, even if we don't implement it.)
// Partial,
}
//---------- principal data structure ----------
slotmap::new_key_type! {
/// Identifies an Account
///
/// After an account is torn down, the `AId` becomes invalid
/// and attempts to use it will give an error.
///
/// The same `AId` won't be reused for a later Account.
struct AId;
/// Identifies a Participant within an Account
///
/// Ie, PId is scoped within in the context of an account.
///
/// As with `AId`, a `PId` is invalid after the
/// participation is torn down, and is not reused.
struct PId;
}
/// Memory tracker inner, including mutable state
///
/// # Module internal documentation
///
/// ## Data structure
///
/// * [`MemoryQuotaTracker`] contains mutex-protected `State`.
/// * The `State` contains a [`SlotMap`] of account records [`ARecord`].
/// * Each `ARecord` contains a `SlotMap` of participant records [`PRecord`].
///
/// The handles [`Account`], [`WeakAccount`], and [`Participation`],
/// each contain a reference (`Arc`/`Weak`) to the `MemoryQuotaTracker`,
/// and the necessary slotmap keys.
///
/// The `ARecord` and `PRecord` each contain a reference count,
/// which is used to clean up when all the handles are gone.
///
/// The slotmap keys which count for the reference count (ie, strong references)
/// are stored as [`refcount::Ref`],
/// which helps assure correct reference counting.
/// (Bare ids [`AId`] and [`PId`] are weak references.)
///
/// ## Data structure lookup
///
/// Given a reference to the tracker, and some ids, the macro `find_in_tracker!`
/// is used to obtain mutable references to the `ARecord` and (if applicable) `PRecord`.
///
/// ## Bookkeeping
///
/// We use separate types for quantities of memory in various "states",
/// rather than working with raw quantities.
///
/// The types, and the legitimate transactions, are in `bookkeeping`.
///
/// ## Reentrancy (esp. `Drop` and `Clone`)
///
/// When the handle structs are dropped or cloned, they must manipulate the refcount(s).
/// So they must take the lock.
/// Therefore, an `Account` and `Participation` may not be dropped with the lock held!
///
/// Internally, this is actually fairly straightforward:
/// we take handles by reference, and constructors only make them at the last moment on return,
/// so our internal code here, in this module, doesn't have owned handles.
///
/// We also need to worry about reentrantly reentering the tracker code, from user code.
/// The user supplies a `dyn IsParticipant`.
/// The principal methods are from [`IsParticipant`],
/// for which we handle reentrancy in the docs.
/// But we also implicitly invoke its `Drop` impl, which might in turn drop stuff of ours,
/// such as [`Account`]s and [`Participation`]s, whose `Drop` impls need to take our lock.
/// To make sure this isn't done reentrantly, we have a special newtype around it,
/// and defer some of our drops during reclaim.
/// That's in `drop_reentrancy` and `tracker::reclaim::deferred_drop`.
///
/// The `Debug` impl isn't of concern, since we don't call it ourselves.
/// And we don't rely on it being `Clone`, since it's in an `Arc`.
///
/// ## Drop bombs
///
/// With `#[cfg(test)]`, several of our types have "drop bombs":
/// they cause a panic if dropped inappropriately.
/// This is intended to detect bad code paths during testing.
#[derive(Debug, Deref, DerefMut)]
struct State {
/// Global parts of state
///
/// Broken out to allow passing both
/// `&mut Global` and `&mut ARecord`/`&mut PRecord`
/// to some function(s).
#[deref]
#[deref_mut]
global: Global,
/// Accounts
accounts: SlotMap<AId, ARecord>,
}
/// Global parts of `State`
#[derive(Debug)]
struct Global {
/// Total memory used
///
/// Wrapper type for ensuring we wake up the reclaimation task
total_used: TotalQtyNotifier,
/// Configuration
config: Config,
}
/// Account record, within `State.accounts`
#[derive(Debug)]
#[must_use = "don't just drop, call auto_release"]
struct ARecord {
/// Number of clones of `Account`; to know when to tear down the account
refcount: refcount::Count<AId>,
/// Child accounts
children: Vec<AId>,
/// Participants linked to this Account
ps: SlotMap<PId, PRecord>,
}
/// Participant record, within `ARecord.ps`
#[derive(Debug)]
#[must_use = "don't just drop, call auto_release"]
struct PRecord {
/// Number of clones of `Participation`; to know when to tear down the participant
refcount: refcount::Count<PId>,
/// Memory usage of this participant
///
/// Not 100% accurate, can lag, and be (boundedly) an overestimate
used: ParticipQty,
/// The hooks provided by the Participant
particip: drop_reentrancy::ProtectedWeak<dyn IsParticipant>,
}
//#################### IMPLEMENTATION ####################
/// Given a `&Weak<MemoryQuotaTracker>`, find an account and maybe participant
///
/// ### Usage templates
///
/// ```rust,ignore
/// find_in_tracker! {
/// weak_tracker => + tracker, state;
/// aid => arecord;
/// [ pid => precord; ]
/// [ ?Error | ?None ]
/// };
///
/// find_in_tracker! {
/// strong_tracker => state;
/// .. // as above
/// };
/// ```
///
/// ### Input expressions (value arguments to the macro0
///
/// * `weak_tracker: &Weak<MemoryQuotaTracker>` (or equivalent)
/// * `strong_tracker: &MemoryQuotaTracker` (or equivalent)
/// * `aid: AId`
/// * `pid: PId`
///
/// ### Generated bindings (identifier arguments to the macro)
///
/// * `tracker: Arc<MemoryQuotaTracker>`
/// * `state: &mut State` (borrowed from a `MutexGuard<State>` borrowed from `tracker`)
/// * `arecord: &mut ARecord` (mut borrowed from `state.accounts`)
/// * `precord: &mut PRecord` (mut borrowed from `arecord.ps`)
///
/// There is no access to the `MutexGuard` itself.
/// For control of the mutex release point, place `find_in_tracker!` in an enclosing block.
///
/// ### Error handling
///
/// If the tracker, account, or participant, can't be found,
/// the macro returns early from the enclosing scope (using `?`).
///
/// If `Error` is specified, applies `?` to `Err(Error::...)`.
/// If `None` is specified, just returns `None` (by applying `?` to None`).
//
// This has to be a macro because it makes a self-referential set of bindings.
// Input syntax is a bit janky because macro_rules is so bad.
// For an internal macro with ~9 call sites it's not worth making a big parsing contraption.
macro_rules! find_in_tracker { {
// This `+` is needed because otherwise it's LL1-ambiguous and macro_rules can't cope
$tracker_input:expr => $( + $tracker:ident, )? $state:ident;
$aid:expr => $arecord:ident;
$( $pid:expr => $precord:ident; )?
// Either `Error` or None, to be passed to `find_in_tracker_eh!($eh ...: ...)`
// (We need this to be an un-repeated un-optional binding, because
// it is used within some other $( ... )?, and macro_rules gets confused.)
? $eh:tt
} => {
let tracker = &$tracker_input;
$(
let $tracker: Arc<MemoryQuotaTracker> = find_in_tracker_eh!(
$eh TrackerShutdown:
tracker.upgrade()
);
let tracker = &$tracker;
)?
let mut state: MutexGuard<State> = find_in_tracker_eh!(
$eh TrackerCorrupted:
tracker.state.lock().ok()
);
let $state: &mut State = &mut *state;
let aid: AId = $aid;
let $arecord: &mut ARecord = find_in_tracker_eh!(
$eh AccountClosed:
$state.accounts.get_mut(aid)
);
$(
let pid: PId = $pid;
let $precord: &mut PRecord = find_in_tracker_eh!(
$eh ParticipantShutdown:
$arecord.ps.get_mut(pid)
);
)?
} }
/// Error handling helper for `find_in_tracker`
macro_rules! find_in_tracker_eh {
{ None $variant:ident: $result:expr } => { $result? };
{ Error $variant:ident: $result:expr } => { $result.ok_or(Error::$variant)? };
}
//========== impls on public types, including public methods and trait impls ==========
//---------- MemoryQuotaTracker ----------
impl MemoryQuotaTracker {
/// Set up a new `MemoryDataTracker`
pub fn new<R: Spawn>(runtime: &R, config: Config) -> Result<Arc<Self>, StartupError> {
let (reclaim_tx, reclaim_rx) = mpsc::channel(0 /* plus num_senders, ie 1 */);
let total_used = TotalQtyNotifier::new_zero(reclaim_tx);
let global = Global { total_used, config };
let accounts = SlotMap::default();
let state = Mutex::new(State { global, accounts });
let tracker = Arc::new(MemoryQuotaTracker { state });
// We don't provide a separate `launch_background_tasks`, because this task doesn't
// wake up periodically, or, indeed, do anything until the tracker is used.
let for_task = Arc::downgrade(&tracker);
runtime.spawn(reclaim::task(for_task, reclaim_rx))?;
Ok(tracker)
}
/// Make a new `Account`
///
/// To actually record memory usage, a Participant must be added.
//
// Right now, parent can't be changed after construction of an Account,
// so circular accounts are impossible.
// But, we might choose to support that in the future.
// Circular parent relationships might need just a little care
// in the reclamation loop (to avoid infinitely looping),
// but aren't inherently unsupportable.
#[allow(clippy::redundant_closure_call)] // We have IEFEs for good reaons
pub fn new_account(self: &Arc<Self>, parent: Option<&Account>) -> crate::Result<Account> {
let mut state = self.lock()?;
let parent_aid_good = parent
.map(|parent| {
// Find and check the requested parent's Accountid
let parent_aid = *parent.aid;
let parent_arecord = state
.accounts
.get_mut(parent_aid)
.ok_or(Error::AccountClosed)?;
// Can we insert the new child without reallocating?
if !parent_arecord.children.spare_capacity_mut().is_empty() {
return Ok(parent_aid);
}
// No. Well, let's do some garbage collection.
// (Otherwise .children might grow without bound as accounts come and go)
//
// We would like to scan the accounts array while mutating this account.
// Instead, steal the children array temporarily and put the filtered one back.
// Must be infallible!
//
// The next line can't be in the closure (confuses borrowck)
let mut parent_children = mem::take(&mut parent_arecord.children);
(|| {
parent_children.retain(|child_aid| state.accounts.contains_key(*child_aid));
// Put the filtered list back, so sanity is restored.
state
.accounts
.get_mut(parent_aid)
.expect("parent vanished!")
.children = parent_children;
})();
Ok::<_, Error>(parent_aid)
})
.transpose()?;
// We have resolved the parent AId and prepared to add the new account to its list of
// children. We still hold the lock, so nothing can have changed.
// commitment - infallible IEFE assures that so we don't do half of it
Ok((|| {
let aid = refcount::slotmap_insert(&mut state.accounts, |refcount| ARecord {
refcount,
children: vec![],
ps: SlotMap::default(),
});
if let Some(parent_aid_good) = parent_aid_good {
state
.accounts
.get_mut(parent_aid_good)
.expect("parent vanished!")
.children
.push(*aid);
}
let tracker = self.clone();
Account { aid, tracker } // don't make this fallible, see above.
})())
}
/// Obtain the lock on the state
fn lock(&self) -> Result<MutexGuard<State>, TrackerCorrupted> {
Ok(self.state.lock()?)
}
}
//---------- Account ----------
impl Account {
/// Register a new Participant
///
/// Returns the [`Participation`], which can be used to record memory allocations.
///
/// Often, your implementation of [`IsParticipant`] wants to contain the [`Participation`].
/// If so, use [`register_participant_with`](Account::register_participant_with) instead.
pub fn register_participant(
&self,
particip: Weak<dyn IsParticipant>,
) -> Result<Participation, Error> {
let aid = *self.aid;
find_in_tracker! {
self.tracker => state;
aid => arecord;
?Error
}
let (pid, cache) = refcount::slotmap_try_insert(&mut arecord.ps, |refcount| {
let mut precord = PRecord {
refcount,
used: ParticipQty::ZERO,
particip: drop_reentrancy::ProtectedWeak::new(particip),
};
let cache =
state
.global
.total_used
.claim(&mut precord, MAX_CACHE, &state.global.config)?;
Ok::<_, Error>((precord, cache))
})?;
let tracker = Arc::downgrade(&self.tracker);
Ok(Participation {
tracker,
pid,
aid,
cache,
})
}
/// Set the callbacks for a Participant (identified by its weak ids)
fn set_participant_callbacks(
&self,
aid: AId,
pid: PId,
particip: drop_reentrancy::ProtectedWeak<dyn IsParticipant>,
) -> Result<(), Error> {
find_in_tracker! {
self.tracker => state;
aid => arecord;
pid => precord;
?Error
}
precord.particip = particip;
Ok(())
}
/// Register a new Participant using a constructor
///
/// Passes `constructor` a [`Participation`] for the nascent Participant.
/// Returns the `P: IsParticipant` provided by the constructor.
///
/// For use when your `impl `[`IsParticipant`] wants to own the `Participation`.
///
/// # Re-entrancy guarantees
///
/// The `Participation` *may* be used by `constructor` for claiming memory use,
/// even during construction.
/// `constructor` may also clone the `Participation`, etc.
///
/// Reclamation callbacks (via the `P as IsParticipant` impl) cannot occur
/// until `constructor` returns.
///
/// # Error handling
///
/// Failures can occur before `constructor` is called,
/// or be detected afterwards.
/// If a failure is detected after `constructor` returns,
/// the `Arc<P>` from `constructor` will be dropped
/// (resulting in `P` being dropped, unless `constructor` kept another clone of it).
///
/// `constructor` may also fail (throwing a different error type, `E`),
/// in which case `register_participant_with` returns `Ok(Err(E))`.
///
/// On successful setup of the Participant, returns `Ok(Ok(Arc<P>))`.
pub fn register_participant_with<P: IsParticipant, E>(
&self,
now: CoarseInstant,
constructor: impl FnOnce(Participation) -> Result<Arc<P>, E>,
) -> Result<Result<Arc<P>, E>, Error> {
use std::sync::atomic::{AtomicBool, Ordering};
/// Temporary participant, which stands in during constructon
#[derive(Debug)]
struct TemporaryParticipant {
/// The age, which is right now. We hope this is all fast!
now: CoarseInstant,
/// Did someone call reclaim() ?
collapsing: AtomicBool,
}
impl IsParticipant for TemporaryParticipant {
fn get_oldest(&self) -> Option<CoarseInstant> {
Some(self.now)
}
fn reclaim(self: Arc<Self>) -> ReclaimFuture {
self.collapsing.store(true, Ordering::Release);
Box::pin(async { Reclaimed::Collapsing })
}
}
let temp_particip = Arc::new(TemporaryParticipant {
now,
collapsing: false.into(),
});
let partn = self.register_participant(Arc::downgrade(&temp_particip) as _)?;
let aid = partn.aid;
let pid_weak = *partn.pid;
// We don't hold the state lock here. register_participant took it and released it.
// This is important, because the constructor might call claim!
// (And, also, we don't want the constructor panicking to poison the whole tracker.)
// But it means there can be quite a lot of concurrent excitement,
// including, theoretically, a possible reclaim.
let particip = match constructor(partn) {
Ok(y) => y,
Err(e) => return Ok(Err(e)),
};
let particip = drop_reentrancy::ProtectedArc::new(particip);
// IEFE prevents use from accidentally dropping `particip` until we mean to
let r = (|| {
let weak = {
let weak = particip.downgrade();
// Trait cast, from Weak<P> to Weak<dyn IsParticipant>.
// We can only do this for a primitive, so we must unprotect
// the Weak, converr it, and protect it again.
drop_reentrancy::ProtectedWeak::new(weak.unprotect() as _)
};
self.set_participant_callbacks(aid, pid_weak, weak)?;
if temp_particip.collapsing.load(Ordering::Acquire) {
return Err(Error::ParticipantShutdown);
}
Ok(())
})();
let particip = particip.promise_dropping_is_ok();
r?;
Ok(Ok(particip))
}
/// Obtains a handle for the `MemoryQuotaTracker`
pub fn tracker(&self) -> Arc<MemoryQuotaTracker> {
self.tracker.clone()
}
/// Downgrade to a weak handle for the same Account
pub fn downgrade(&self) -> WeakAccount {
WeakAccount {
aid: *self.aid,
tracker: Arc::downgrade(&self.tracker),
}
}
}
impl Clone for Account {
fn clone(&self) -> Account {
let tracker = self.tracker.clone();
let aid = (|| {
let aid = *self.aid;
find_in_tracker! {
tracker => state;
aid => arecord;
?None
}
let aid = refcount::Ref::new(aid, &mut arecord.refcount).ok()?;
// commitment point
Some(aid)
})()
.unwrap_or_else(|| {
// Either the account has been closed, or our refcount overflowed.
// Return a busted `Account`, which always fails when we try to use it.
//
// If the problem was a refcount overflow, we're technically violating the
// documented behaviour, since the returned `Account` isn't equivalent
// to the original. We could instead choose to tear down the Account;
// that would be legal; but it's a lot of code to marginally change the
// behaviour for a very unlikely situation.
refcount::Ref::null()
});
Account { aid, tracker }
}
}
impl Drop for Account {
fn drop(&mut self) {
(|| {
find_in_tracker! {
self.tracker => state;
*self.aid => arecord;
?None
}
if let Some(refcount::Garbage(mut removed)) =
slotmap_dec_ref!(&mut state.accounts, self.aid.take(), &mut arecord.refcount)
{
// This account is gone. Automatically release everything.
removed.auto_release(state);
}
Some(())
})()
.unwrap_or_else(|| {
// Account has been torn down. Dispose of the strong ref.
// (This has no effect except in cfg(test), when it defuses the drop bombs)
self.aid.take().dispose_container_destroyed();
});
}
}
//---------- WeakAccount ----------
impl WeakAccount {
/// Upgrade to an `Account`, if the account still exists
pub fn upgrade(&self) -> crate::Result<Account> {
let aid = self.aid;
// (we must use a block, and can't use find_in_tracker's upgrade, because borrowck)
let tracker = self.tracker.upgrade().ok_or(Error::TrackerShutdown)?;
let aid = {
find_in_tracker! {
tracker => state;
aid => arecord;
?Error
}
refcount::Ref::new(aid, &mut arecord.refcount)?
// commitment point
};
Ok(Account { aid, tracker })
}
/// Obtains a handle onto the `MemoryQuotaTracker`
///
/// The returned handle is itself weak, and needs to be upgraded before use.
pub fn tracker(&self) -> Weak<MemoryQuotaTracker> {
self.tracker.clone()
}
}
//---------- Participation ----------
impl Participation {
/// Record that some memory has been (or will be) allocated
pub fn claim(&mut self, want: usize) -> crate::Result<()> {
self.claim_qty(Qty(want))
}
/// Record that some memory has been (or will be) allocated (using `Qty`)
pub(crate) fn claim_qty(&mut self, want: Qty) -> crate::Result<()> {
if let Some(got) = self.cache.split_off(want) {
return got.claim_return_to_participant();
}
find_in_tracker! {
self.tracker => + tracker, state;
self.aid => arecord;
*self.pid => precord;
?Error
};
let mut claim = |want| -> Result<ClaimedQty, _> {
state
.global
.total_used
.claim(precord, want, &state.global.config)
};
let got = claim(want)?;
if want <= TARGET_CACHE_CLAIMING {
// While we're here, fill the cache to TARGET_CACHE_CLAIMING.
// Cannot underflow: cache < want (since we failed at `got` earlier
// and we've just checked want <= TARGET_CACHE_CLAIMING.
let want_more_cache = Qty(*TARGET_CACHE_CLAIMING - *self.cache.as_raw());
if let Ok(add_cache) = claim(want_more_cache) {
// On error, just don't do this; presumably the error will show up later
// (we mustn't early exit here, because we've got the claim in our hand).
self.cache.merge_into(add_cache);
}
}
got.claim_return_to_participant()
}
/// Record that some memory has been (or will be) freed by a participant
pub fn release(&mut self, have: usize) // infallible
{
self.release_qty(Qty(have));
}
/// Record that some memory has been (or will be) freed by a participant (using `Qty`)
pub(crate) fn release_qty(&mut self, have: Qty) // infallible
{
let have = ClaimedQty::release_got_from_participant(have);
self.cache.merge_into(have);
if self.cache > MAX_CACHE {
match (|| {
find_in_tracker! {
self.tracker => + tracker, state;
self.aid => arecord;
*self.pid => precord;
?None
}
let return_from_cache = Qty(*self.cache.as_raw() - *TARGET_CACHE_RELEASING);
let from_cache = self.cache.split_off(return_from_cache).expect("impossible");
state.global.total_used.release(precord, from_cache);
Some(())
})() {
Some(()) => {} // we've given our cache back to the tracker
None => {
// account (or whole tracker!) is gone
// throw away the cache so that we don't take this path again for a bit
self.cache.take().dispose_participant_destroyed();
}
}
}
}
/// Obtain a handle onto the account
///
/// The returned handle is weak, and needs to be upgraded before use,
/// since a [`Participation`] doesn't keep its Account alive.
///
/// The returned `WeakAccount` is equivalent to
/// all the other account handles for the same account.
pub fn account(&self) -> WeakAccount {
WeakAccount {
aid: self.aid,
tracker: self.tracker.clone(),
}
}
}
impl Clone for Participation {
fn clone(&self) -> Participation {
let aid = self.aid;
let cache = ClaimedQty::ZERO;
let tracker: Weak<_> = self.tracker.clone();
let pid = (|| {
let pid = *self.pid;
find_in_tracker! {
self.tracker => + tracker_strong, state;
aid => _arecord;
pid => precord;
?None
}
let pid = refcount::Ref::new(pid, &mut precord.refcount).ok()?;
// commitment point
Some(pid)
})()
.unwrap_or_else(|| {
// The account has been closed, the participant torn down, or the refcount
// overflowed. We can a busted `Participation`.
//
// We *haven't* incremented the refcount, so we mustn't return pid as a strong
// reference. We aren't supposed to count towards PRecord.refcount, we we *can*
// return the weak reference aid. (`refcount` type-fu assures this is correct.)
//
// If the problem was refcount overflow, we're technically violating the
// documented behaviour. This is OK; see comment in `<Account as Clone>::clone`.
refcount::Ref::null()
});
Participation {
aid,
pid,
cache,
tracker,
}
}
}
impl Drop for Participation {
fn drop(&mut self) {
(|| {
find_in_tracker! {
self.tracker => + tracker_strong, state;
self.aid => arecord;
*self.pid => precord;
?None
}
// release the cached claim
let from_cache = self.cache.take();
state.global.total_used.release(precord, from_cache);
if let Some(refcount::Garbage(mut removed)) =
slotmap_dec_ref!(&mut arecord.ps, self.pid.take(), &mut precord.refcount)
{
// We might not have called `release` on everything, so we do that here.
removed.auto_release(&mut state.global);
}
Some(())
})()
.unwrap_or_else(|| {
// Account or Participation or tracker destroyed.
// (This has no effect except in cfg(test), when it defuses the drop bombs)
self.pid.take().dispose_container_destroyed();
self.cache.take().dispose_participant_destroyed();
});
}
}
//========== impls on internal types ==========
impl State {
/// Obtain all of the descendants of `parent_aid` according to the Child relation
///
/// The returned `HashSet` includes `parent_aid`, its children,
/// their children, and so on.
///
/// Used in the reclaimation algorithm in [`reclaim`].
fn get_aid_and_children_recursively(&self, parent_aid: AId) -> HashSet<AId> {
let mut out = HashSet::<AId>::new();
let mut queue: Vec<AId> = vec![parent_aid];
while let Some(aid) = queue.pop() {
let Some(arecord) = self.accounts.get(aid) else {
// shouldn't happen but no need to panic
continue;
};
if out.insert(aid) {
queue.extend(arecord.children.iter().cloned());
}
}
out
}
}
impl ARecord {
/// Release all memory that this account's participants claimed
fn auto_release(&mut self, global: &mut Global) {
for (_pid, mut precord) in self.ps.drain() {
precord.auto_release(global);
}
}
}
impl PRecord {
/// Release all memory that this participant claimed
fn auto_release(&mut self, global: &mut Global) {
let for_teardown = self.used.for_participant_teardown();
global.total_used.release(self, for_teardown);
}
}
|