Implementing a Custom Global Allocator in Rust

Modern applications often require fine-grained control over memory allocation for performance optimizations or specific memory management strategies. Rust, with its focus on safety and performance, provides the flexibility to define and use custom global allocators. This challenge will test your understanding of Rust's allocator API and memory management concepts.

Problem Description

Your task is to implement a custom global allocator in Rust that tracks the total amount of memory currently allocated. You will then use this allocator to demonstrate its behavior when allocating and deallocating memory within a simple program.

Key Requirements:

1. Implement the GlobalAlloc trait: Define a struct (e.g., TrackingAllocator) that implements the core::alloc::GlobalAlloc trait.
2. Track Allocated Memory: Your allocator must maintain a count of the total bytes currently allocated.
3. Basic Allocation and Deallocation: Implement the alloc and dealloc methods to handle memory requests and returns. For this challenge, you can delegate the actual memory allocation to Rust's default allocator. The focus is on tracking.
4. Thread Safety: Ensure your allocator is thread-safe, as global allocators are accessed by multiple threads.
5. Global Registration: Register your custom allocator as the global allocator using the #[global_allocator] attribute.
6. Demonstrate Usage: Write a main function that allocates and deallocates memory using standard Rust data structures (e.g., Vec, Box) and prints the total allocated memory before and after these operations.

Expected Behavior:

When the program runs, it should:

• Initialize the custom allocator.
• Report an initial total allocated memory (likely 0 or close to it, depending on program startup).
• Perform some allocations (e.g., creating a Vec with some elements, a Box of a large struct).
• Report the updated total allocated memory.
• Perform deallocations (e.g., dropping the Vec and Box).
• Report the final total allocated memory, which should ideally return to a state close to the initial state.

Edge Cases to Consider:

• Zero-sized allocations: How does your allocator handle requests for 0 bytes?
• Large allocations: While not strictly enforced by constraints, consider the implications of very large allocations.
• Concurrent access: Although we are delegating actual allocation, your tracking mechanism must be thread-safe.

Examples

Example 1:

Scenario: Initialize the allocator and print its initial state.

Input: (Implicitly, the program starts)

Output:

Initial allocated memory: 0 bytes

(Note: The exact initial value might vary slightly depending on the Rust runtime, but for this challenge, we aim for a clean starting point.)

Example 2:

Scenario: Allocate a Vec<u8> with 100 bytes, then print the total allocated memory.

Input:

let mut data = Vec::with_capacity(100);
data.extend_from_slice(&[0u8; 100]);
// After this, print total allocated memory

Output:

Allocated 100 bytes.
Total allocated memory: 100 bytes (approximately)

Example 3:

Scenario: Allocate a Box<[u64; 1000]> and a Vec<i32> with 50 elements, then deallocate both.

Input:

let large_box = Box::new([0u64; 1000]); // 1000 * 8 bytes = 8000 bytes
let mut small_vec = Vec::with_capacity(50);
small_vec.extend_from_slice(&[0i32; 50]); // 50 * 4 bytes = 200 bytes

// Print total allocated memory
// ... then `drop(large_box)` and `drop(small_vec)` ...
// Print total allocated memory again

Output:

Allocated large_box (8000 bytes) and small_vec (200 bytes).
Total allocated memory: 8200 bytes (approximately)

Deallocated large_box and small_vec.
Total allocated memory: 0 bytes (approximately)

Constraints

• You must use the core::alloc::GlobalAlloc trait.
• Your allocator must be registered using #[global_allocator].
• The actual memory allocation and deallocation can be handled by delegating to Rust's default allocator (System).
• The tracking of allocated memory should be accurate within the scope of your custom allocator's logic.
• The solution should be written in Rust.
• Performance of the tracking mechanism itself is not critical for this challenge, but the underlying alloc and dealloc calls should be reasonably efficient (which is achieved by delegating to System).

Notes

• The core::alloc::GlobalAlloc trait requires two methods: alloc and dealloc.
• The Layout struct passed to alloc and dealloc contains size and align information.
• For thread-safe tracking, consider using std::sync::atomic types (e.g., AtomicUsize) or a Mutex.
• You will likely need to use the std::alloc::System allocator as a fallback for the actual memory operations.
• Remember that memory overhead from the allocator itself or the underlying system might lead to minor discrepancies in reported total allocated memory. The goal is to correctly track what your allocator is responsible for.
• To see your custom allocator in action, you'll need to compile with #![feature(allocator_api)] and ensure your Cargo.toml is set up for a binary crate.