[DNM] Wrapper allocator PoC

[s0me0ne-unkn0wn]

It's a PoC. Do not merge.

This is the first and minimalistic attempt to implement a bookkeeping global allocator (#6687) without any interthread locks. It works, but I'm unsure if its behavior is sound, so any comments from anyone who has more experience with atomics and memory ordering than me are highly appreciated.

[slumber]

While access to each field of WrapperAllocatorData is atomic, this method (as others too) loads them one by one, and you can expect anything to happen between these loads. For example:

1. Read let allocated = ALLOCATOR_DATA.allocated.load(SeqCst);. Scheduler puts our thread to sleep.
2. Another thread calls alloc and then checkpoint.
3. We wake up and checkpoint goes back in time due to allocated being outdated.

Either these values shouldn't depend on each other and work as plain counters or something, or the whole struct should go under mutex.

I mean, I understand this is our best effort, and thus it depends on what is it we're trying to achieve.

[s0me0ne-unkn0wn]

While access to each field of WrapperAllocatorData is atomic, this method (as others too) loads them one by one, and you can expect anything to happen between these loads.

Yes, that's exactly the concern I'm worried about! And I'm trying to convince myself it is okay. Not sure at all.

It's okay to have a global lock in the checkpoint(). It's definitely not okay to have a global lock in the alloc(). I've omitted a mutex in the checkpoint() for now as it is supposed to be called from a single thread only, but we can add it there, checkpoint is a rare event. alloc() is much more concerning.

Say we have two threads, as in your example, and they're allocating at nearly the same point in time. I'm indexing the values below for clarity.

• Thread1 executes old_alloc1 = allocated.fetch_add(layout_size1). At this point allocated is updated and represents the new size.
• Then the context switches to another thread
• Thread2 executes old_alloc2 = allocated.fetch_add2(layout_size2). At this point, allocated is updated again and represents the true allocated value. Both old_alloc1 and old_alloc2 are set to proper values.
• Thread2 executes peak.fetch_max(old_alloc2 + layout_size2). The peak value is properly updated.
• The context switches back and now thread1 executes peak.fetch_max(old_alloc1 + layout_size1). Here's the tricky part. Although thread1 is not aware of the allocated value has already been updated by another thread, it already has proper old_alloc1, and the peak already stores the proper peak value, including thread2's allocation. What is improper here is the assumption that for thread1 old_alloc1 + layout_size1 is the new allocated value that should be compared with the peak. But do we care? I believe we don't because thread2 has already updated the peak. I mean, by definition, old_alloc1 + layout_size1 < old_alloc2 + layout_size2 because old_alloc2 == old_alloc1 + layout_size1. So we have no chance of missing the new peak value.

Thus it seems we can avoid a global lock inside the alloc() and its siblings, but I'm just afraid of missing something. It's the case where ten eyes are much better than two.

[koute]

Using a spinlock should probably be fine here as the critical section here's going to be very short anyway.

Another more complicated alternative that would help with thread contention would be to make this per thread: put the state in thread local storage and have one spinlock per TLS, and only when grabbing a checkpoint lock them all and collate the data. I don't think it's worth it though.

[alexggh]

This two lines of code could be refactored out as method on WrapperAllocatorData

[alexggh]

A few questions:

1. This will measure only the dynamically allocated memory, it will fail to monitor any unexpected increase in the stack, is that what we want ?
2. How is this different from the bellow get_max_rss_thread and why isn't that enough ?

[mrcnski]

2. How is this different from the bellow get_max_rss_thread and why isn't that enough ?

The issue with that is it is not deterministic enough. (Is mainly being used for gathering metrics right now.) Different kernel configurations/versions may manage memory differently (simple example is some validators may have swap enabled and some not). So they may get different values for the resident memory (how much is actually held in RAM). So some validators may reach the limit and others not.

[s0me0ne-unkn0wn]

1. This will measure only the dynamically allocated memory, it will fail to monitor any unexpected increase in the stack, is that what we want ?

I hope yes, in the preparation phase we're not executing any untrusted code so we just presume that Wasmtime guys know what they're doing and won't abuse the stack usage. We only bother about malicious Wasm code that could force the compiler to allocate a lot of memory.

2. How is this different from the bellow get_max_rss_thread and why isn't that enough ?

A lot of concerns here.

• Not supported on every platform
• Allocated memory is not necessarily resident
• Wasmtime may spawn its own threads and we'd lose their allocations in the per-thread accounting

[mrcnski]

So some validators may reach the limit and others not.

We don't have disputes for preparation, but what can happen is the attacker gets lucky and gets the PVF through pre-checking without hitting the limits, but then the limits are actually hit when preparing for execution causing no-shows (since we don't dispute on preparation errors). I guess we can have a lower, stricter limit for pre-checking, which we should have anyway.

Wasmtime may spawn its own threads and we'd lose their allocations in the per-thread accounting

Good point. We can't use RUSAGE_SELF instead of RUSAGE_THREAD because there's no way to "reset" the max from a previous job.

[mrcnski]

I think the memory ordering can just be Relaxed right now, since the atomics are being used as counters and not to synchronize other memory operations. Relaxed still gives you a total ordering between accesses to the same atomic.

Locks would make this slightly more deterministic but it doesn't seem worth the cost. Say you have a situation like thread1 calling alloc and getting old_alloc, thread2 stepping in and calling dealloc, and then thread1 setting the max peak to something higher than was actually allocated. You can prevent this with a lock but not totally solve it. There is still some indeterminacy where a worker may call the dealloc thread first and not hit the max at all while other workers would hit it. It's an edge case and not sure how feasible it would be to craft an attack around that.

But I think we decided on using cgroups to set the limit though which seems the way to go given the above. We can instead use this wrapper for detecting OOM by setting a flag before allocs and clearing it after.

[koute]

The if's unnecessary if you make the counters isize.

[s0me0ne-unkn0wn]

Yes, you're right, I'm already thinking about making them isize anyway so I can bother less about negative overflows. But I also don't want even try to update peak when it's definitely not necessary.

[Polkadot-Forum]

This pull request has been mentioned on Polkadot Forum. There might be relevant details there:

https://forum.polkadot.network/t/ux-of-distributing-multiple-binaries-take-2/2854/1

[s0me0ne-unkn0wn]

Take 2 is here. I'm trying a completely different approach.

I like the overall idea of what I'm trying to implement here, but I doubt it will make it into production, not in this form at least. Probably it can be improved to make it more safe and less ugly (any ideas are appreciated). If not, maybe it's worth giving a try to the spinlock approach proposed by @koute.

[koute]

All in all, having a Vec into which we gather the allocation deltas seems kinda unnecessary to me, considering the downsides (extra allocation is necessary every time we start tracking, the extra latency spike when it'll have to tally the deltas at the end, and the possibility of a panic since the buffer needs to be big enough or the extra complexity to make it resizable) and as far as I can see the only thing we're gaining is the lack of an exclusive lock during allocations which otherwise would only have to protect two very cheap operations. Seems like the tradeoff is not really worth it?

Here's the equivalent code using a spinlock:

struct AllocationTracker {
    lock: AtomicBool,
    current: AtomicIsize,
    peak: AtomicIsize
}

impl AllocationTracker {
    fn add(&self, delta: isize) {
        loop {
            // Try to acquire the lock.
            if self.lock.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed).is_ok() {
                break;
            }
            // We failed to acquire the lock; wait until it's unlocked.
            //
            // In theory this should result in less coherency traffic as unlike compare_exchange it is a read-only operation,
            // so multiple cores can execute it simultaneously without taking an exclusive lock over the cache line.
            while(self.lock.load(Ordering::Relaxed)) {}
        }
        let current = self.current.fetch_add(delta, Ordering::Relaxed);
        self.peak.fetch_max(current, Ordering::Relaxed);
        self.lock.store(false, Ordering::Release);
    }
}

As you can see it's quite a bit simpler. (:

(With spin-locks there are also some other tricks/variations that one can attempt to speed it up, but it does depend on the specific situation and generally should be benchmarked to see whether it makes a difference. But we don't make that many allocations so they're probably unnecessary.)

[koute]

Won't this will deadlock if start_tracking is called twice? (Assigning a new backlog will trigger a deallocation of the old backlog.)

[s0me0ne-unkn0wn]

Seems like you're right 😕
Nah, the whole thing now seems FUBAR to me. I was thinking about using the nightly allocator_api feature and Vec::with_capacity_in() to allow side-allocations from the main allocator but everything I try to implement looks too hacky. I'll give a try to your spinlock approach, probably, the overhead will be even less than the overhead of allocating and initializing the 800 Mb array and then replaying it.

[s0me0ne-unkn0wn]

Btw, why didn't you use std::hint::spin_loop() inside while in your spinlock example? Was it on purpose?

[koute]

Btw, why didn't you use std::hint::spin_loop() inside while in your spinlock example? Was it on purpose?

No particular reason; I wrote that code from memory and simply forgot. :P Yeah, you should probably add that. (On AMD64 that will generate a pause instruction, which is generally a good idea to use in a spinlock.)

[s0me0ne-unkn0wn]

Looking at your code a bit more... The critical section is under exclusive lock now, do current and peek still have to be Atomics? Would we cut off a bit of the overhead making them simple isizes?

[koute]

The critical section is under exclusive lock now, do current and peek still have to be Atomics? Would we cut off a bit of the overhead making them simple isizes?

AFAIK they don't. I just used them as I didn't feel like using an UnsafeCell for the example.

[mrcnski]

Repeating some things as my review yesterday got lost. The memory overhead per prepare job right now is too high to be viable, if we really do that many allocations. It is an interesting idea, but it does have the extra complexity of determining what is a safe amount to pre-allocate. Not sure if the latency spike at the end is a problem, we do it at the end of an already blocking operation, but ideally we benchmark the overall start-to-end time and get numbers for the different approaches. The panic can be worked around (just set some bit that we overflowed and return an error in end_tracking). Overall the spinlock strategy does seem like the best approach so far.

[s0me0ne-unkn0wn]

Take 3: use a spinlock. Looks simple enough not to screw everything up.

I wanted to keep tracking but got rid of it for now, as it introduces overhead comparable to the whole tracking logic overhead. But I'm concerned a bit about how things will go in the non-tracking mode. It will keep tracking in vain (which is okay by itself), and the current counter may overflow the isize at some point. Shouldn't be a problem for the release build, but I'm not sure about the debug ones.

[mrcnski]

Looking good, now just need performance numbers. :)

It will keep tracking in vain (which is okay by itself)

Preparation should do the vast majority of allocations in a worker, and tracking will be turned on for that, so this seems fine?

and the current counter may overflow the isize at some point. Shouldn't be a problem for the release build, but I'm not sure about the debug ones.

Can use saturating_add. But we reset the counter for each job right? I think it's only a potential issue on 32-bitness. And eventually we should have a memory limit for preparation that's well below isize::MAX.

[s0me0ne-unkn0wn]

Can use saturating_add. But we reset the counter for each job right?

Silly me 🤦 We're counting the current allocation and the peak, not the total of all the allocations, so it cannot overflow (unless we indeed try to allocate more than isize.MAX at the same time, but then we're in trouble anyway). So everything is fine, no additional care is needed.

[mrcnski]

It's a good point about overflows though. It's not safe for a panic to happen in this code. Only possible in debug mode, but still, I would opt for extra caution here.

[s0me0ne-unkn0wn]

Compiling the Moonbeam runtime in single-threaded mode, 100 iterations, times are in wallclock milliseconds:

features	min	avg	max
none	7238	7449	8269
jemalloc-allocator	7202	7363	7962
tracking-allocator	7138	7322	8676

Compiling the Moonbeam runtime in multi-threaded mode, 100 iterations, times are in wallclock milliseconds:

features	min	avg	max
none	1132	1176	1266
jemalloc-allocator	1145	1203	1320
tracking-allocator	1674	1791	1883

NB: I'm not sure that enabling the jemalloc-allocator feature on the polkadot-node-core-pvf-worker package really enables Jemalloc as the global allocator (but enabling tracking-allocator does that for sure).

Predictably, multi-threaded performance has degraded. However, single-threaded doesn't show any significant deviation in the average value but diverges more between min and max.

I'll do more experiments to ensure I'm comparing with Jemalloc and not with the System allocator. Also, still searching for ways to estimate the overall performance degradation, not only for the compilation but for the whole node (it's relevant to the current situation when worker binaries haven't been separated yet).

Benchmarking code:

	for _ in 0..NITER {
		let cpu_time_start = ProcessTime::now();
		let start = Instant::now();
		let _result = prepare_artifact(pvf.clone(), cpu_time_start);
		let elapsed = start.elapsed().as_millis() as u64;
		if elapsed < min {
			min = elapsed;
		}
		if elapsed > max {
			max = elapsed;
		}
		total += elapsed
	}
	println!("min {} avg {} max {}", min, total / NITER as u64, max);

[koute]

@s0me0ne-unkn0wn Can you also try the following variations and see how it affects performance in the multithreaded case?

1. Move the while self.lock.load(Ordering::Relaxed) loop to the top so that it executes before compare_exchange_weak.
2. Remove the call to std::hint::spin_loop().
3. Do (1) and (2) at the same time.

Maybe we could also try coalescing small allocations in TLS and only update the global counters when they reach a certain threshold. For example as an extreme example, let's say a thread allocates a single byte and then deallocates it, and does this 100 times. Do we actually need to update the counters each time the thread does this, or could we wait, say, until the total delta is e.g. at least ~1KB and only then update the counters? (So in this example it'd not touch the global counter at all, since the total delta is at most equal to 1.)

[paritytech-cicd-pr]

The CI pipeline was cancelled due to failure one of the required jobs.
Job name: test-linux-stable
Logs: https://gitlab.parity.io/parity/mirrors/polkadot/-/jobs/3393419