Pre-proposal: Type-based allocator selection

[rinon]

TL;DR: Proposing the addition of a trait that indicates the default allocator for a type to the compiler for use in Box, etc. This would allow programmers to use different allocators on a per-type basis, rather than requiring the use of Box::new_in() at every allocation site.

On a personal note, I've been following the alloc-wg work with great interest, and I'd love to help out if I can.

Background

Currently, the interface to allocate using a custom allocator looks like:

let allocator: Allocator = ...;
let value: V = ...;
let boxed_value: Box<V, Allocator> = Box::new_in(value, allocator);

This API allows the programmer to specify a particular allocator for a heap allocation at the allocation site. The allocator itself does not know what type it is allocating, only the type’s size and alignment, so it cannot make allocation decisions based on types. There was a previous pre-proposal to add type information to AllocRef, but it did not get very far, partially due to inherent limitations in the Rust TypeId system.

Issues with Segmented Heaps

For our use case, we want to use multiple allocators and constrain some types to be only allocatable using a specific allocator. This problem consists of two parts: 1) allocating instances of a type in a specific allocator, and 2) preventing allocation of that type with any other allocator.

The first constraint, allocating instances of a type in a specific allocator, is viable although verbose, given the current API. For example, we can create a newtype wrapper around Box that only accepts types that implement a CustomAllocated trait which will be allocated using the custom allocator:

trait CustomAllocated { }

#[repr(transparent)]
struct CustomBox<T: CustomAllocated>(Box<T, CustomAllocator>);
impl CustomAllocated for u32 {}


fn main() { let _shared_box = CustomBox::new(1u32); }

Unfortunately, we don’t know of a way to prevent a type from being allocated using a different allocator from the specific one we want. If allocators had access to type information, we could dispatch to the correct heap segment based on type inside a single allocator. However, as discussed above, there doesn’t seem to be a good way to get type information passed to the allocator. Even if we can’t prevent allocation to the “wrong” allocator, we would at least like to be able to control the default allocator for a type, as this will minimize the possibility of allocating using the wrong allocator.

Use Cases

We know of the following use cases that would benefit from the ability to select allocators based on type information:

Allocation into different sets of physical memory. The memory attached to the CPU has different characteristics (bandwidth, latency, etc) from memory on an accelerator such as a graphics card. Type-based allocation would let the programmer specify that texture objects, for instance, should be allocated in graphics memory, without requiring the programmer to specify the graphics allocator at each allocation site.

Larger multiprocessor systems have non-uniform memory access times (NUMA) meaning that the time to access a given memory location depends on the distance between the processor and the memory module. In this case, type-based allocation let’s programmers express NUMA-aware optimizations via the type system (See #29). The same goes for systems that can access remote memory (i.e., another host) via remote DMA.

Allocation into segmented memory domains. Currently, memory is generally segmented into kernel memory and per-process memory, subject to memory protection so that each domain cannot access others. These memory access permissions are enforced via the primary and secondary address translation mechanisms in processors with memory management units. However, modern processors support fast switching of access permissions inside a single process via features such as memory protection keys (Intel) and memory domains (ARM). We are using these primitives to put memory belonging to each library in a separate memory domain and use types to express which memory objects may be shared across domains. This has security and reliability benefits when a process contains a mix of safe (Rust) and unsafe (C/C++) code. This use case is our strongest motivator since inadvertent allocation into the wrong domain is far less likely, in our experience, when types can be used to constrain the allocation domain instead of correctly using Box::new_in at each allocation site.

For a concrete example, the RedLeaf research operating system uses memory domains to separate memory allocations belonging to different components (such as device drivers and OS services) for fault tolerance. This design already includes a concept of exchangeable memory which are types that may be moved across memory domains – a concept which can be implemented more cleanly with support for type-based allocation according to the authors of the RedLeaf papers. See https://www.usenix.org/system/files/osdi20-narayanan_vikram.pdf.

Performance optimization. Being able to distinguish allocated and deallocated regions of memory by type allows for increased performance gains, orthogonal to other performance-enhancing allocation techniques such as pooling (this was also noted in the literature). Type-based allocation further aids the construction of object caches, which can significantly reduce allocation overhead in heavy workload scenarios (e.g., in web browsers) and garbage collectors (see the aforementioned pre-proposal).

While region allocation can be approximated by simply using allocation sizes, this is less efficient as different but identically sized structures may have different access patterns. In addition, mixing types of the same size in the same region is less secure against memory corruption resulting in type confusion.

Proposed Solution

We propose defining a new trait, e.g., TypeAllocated, that the compiler will use to select a default allocator for certain types. When a type implements this trait, the compiler will use the associated allocator to allocate boxed objects of this type instead of the default global allocator. By itself, this change will not prevent the type from being allocated via a custom allocator if it is allocated using, e.g., Box::new_in(). However, selecting a default allocator for a type handles the majority of common cases and allows a type’s allocator to be changed in one place without having to change it at every allocation site.

Example:

#![feature(allocator_api)]
use std::alloc::Allocator;
use std::boxed::Box;

trait TypeAllocated {
   type TypeAllocator: Allocator+Default;
}

struct SpecialType(u32);

impl TypeAllocated for SpecialType {
   type TypeAllocator = CustomAllocator;
}

struct CustomAllocator;

impl Allocator for CustomAllocator {
   // ...
}

fn main() {
   // The compiler would see that TypeAllocated was implemented for
   // SpecialType and allocate this box using the provided default
    // allocator
   let _special_box: Box<SpecialType, CustomAllocator> =
          Box::new(SpecialType(1u32));
}

In order to use an allocator in this way, the compiler needs to be able to create an instance of the default allocator when constructing a box. This is handled by requiring Default on the type bound of TypeAllocator in TypeAllocated. The compiler can use <T as TypeAllocated>::TypeAllocator::default() to construct an instance of the allocator

Potential Issues

The type of Box is currently Box<T, A: Allocator>. Thus, changing the implementation of TypeAllocated for a type would change the type signature of Box::new() for that type. That may not be an issue, but changing an impl in one place and breaking types may not be ideal. Unfortunately, that’s also the benefit of this approach, so that will need to be evaluated.

[edit] We can possibly deal with this issue by changing the definition of Box to default to the TypeAllocator allocator and use a blanket implementation of TypeAllocated to remain backwards compatible:

struct Box<T: TypeAllocated, A: AllocRef = <T as TypeAllocated>::TypeAllocator>;

// requires specialization to allow users to override this
impl<T> TypeAllocated for T {
    type TypeAllocator = Global;
}

Implementing TypeAllocated doesn’t prevent a type from being allocated with a non-default allocator. This could be desirable in order to prevent bugs where types are allocated using an unexpected allocator. One option would be to disable Box::new_in if type-based allocation is requested for the object type.

Alternatives

It would be possible to implement this via a macro at call sites and dispatch to the correct Box::new_in() variant. However, this still requires the programmer to opt in to the special allocation at every allocation site, something which we want to avoid. In addition, an allocation-site macro would not allow a crate to control the allocator for its own types in a foreign crate, although this is unlikely to work unless the default allocator in the signature of Box::new() is controlled by TypeAllocated.

I've prototyped a simpler version of this proposal in the compiler, which passes an additional parameter to the global allocator indicating the memory domain for the allocation based on a trait implementation for the type being allocated. This is somewhat similar to #27, but passing a memory domain id instead of the type, thus avoiding some of the issues around TypeId. However, this approach seems strictly less flexible than the approach proposed above, so I don't think it's a great solution.

[TimDiekmann]

I think this proposal (#79) will solve your issues. Basically it replaces the A: Allocator parameter with a storage trait which is typed. You may want to read through the first post and take a look at the documentation of an experiment I made, which was inspired by the linked proposal. Does this solve your issues with the current API? While it tries to solve another problem, your problem can be solved this way, too, I guess.

[rinon]

The basic problem this proposal aims to solve is relieving the programmer of the need to explicitly opt in to an alternative allocator at every allocation site for some set of types. I could be missing something, but I'm not sure how #79 would do that. Afaict #79 would add a generic backing store field to containers, but still requires the use of new_in, and in fact makes that API more complicated.

[RustyYato]

@rinon how would this interact with generic code? Because it seems to mean that anyone who assumes that Box<T> is just a pointer (which is it right now) would be wrong to assume that in the future, because it would really be Box<T, PossiblyCustomAllocator>, or would this just be sugar on Box::new? I don't see how that could be added in a backwards-compatible way. (at a minimum it would make implementing TypeAllocated a breaking change).

edit: if you really needed a easy way to sugar over Box::new, you could instead do MyType::new_box with the custom allocator. No need to change the lang for this

[rinon]

Yes, implementing TypeAllocated for a type would require that Box<T> declarations for that type be updated given the current type of Box. However, that's exactly what I am after - the ability to change the implementation of TypeAllocated and have that take effect for every allocation of objects of that type.

Ideally I'd like the default allocator type of Box<T> to change depending on the T. I just tested an idea that I think could work for this - could we change the declaration of Box to:

struct Box<T: TypeAllocated, A: AllocRef = <T as TypeAllocated>::TypeAllocator>;

// requires specialization to allow users to override this
impl<T> TypeAllocated for T {
    type TypeAllocator = Global;
}

If you really needed a easy way to sugar over Box::new, you could instead do MyType::new_box with the custom allocator.

Yeah, considered that. Two major issues, from my perspective:

1. It would be easy for a programmer to accidentally use a plain Box::new(), which in at least a couple of our use cases is a bug with security implications. It's a lot easier to audit for opting out of the default allocator.
2. Long-term, I would like this to apply to other collections, so that other things like Vec have a configurable default allocator.

[TimDiekmann]

I don't think this will work. Assuming we have specialization for overwriting TypedAllocated, only the crate who defines T can overwrite the implementation because of the orphan rule.

[rinon]

Yes, that would be the idea, that the crate that owns T should know where it needs to be allocated. Do you think that's too much of a limitation? For the use cases I'm looking at, I think it would be ok. The allocation behavior is determined by the type, rather than the code location where that type is being allocated. Of course, this won't always be the case, so in cases where the allocation site needs to control the allocator we always have the current allocator API.

[TimDiekmann]

Do you think that's too much of a limitation?

To be honest, I think it is. Since these structs are defined in the standard library, I don't feel such a limitation should exist. Especially since this is only a specific use case.

I still believe that #79 can solve this problem more elegantly. I guess it will not be possible to change Box::new in a backward compatible way, so for another allocator you will most likely have to define a new method, for example Box::new_in. Whether an allocator or a buffer is given here is not important for the time being, it only has to be typed. The advantage of the typed-buffer approach is that a buffer can be defined everywhere and can be used in any other place. For example a buffer which is optimized for integral values (SIMD for example?), then the trait Buffer<Integral> is implemented. This buffer can be used only with Box<Integral> then. Such a buffer could rely on an allocator, too.

[RustyYato]

@rinon

It would be easy for a programmer to accidentally use a plain Box::new(), which in at least a couple of our use cases is a bug with security implications. It's a lot easier to audit for opting out of the default allocator.

If security is your biggest concern, then I would create a newtype wrapper around Box (or Vec, or others), and use that throughout your codebase. This newtype can have the semantics that you desire, and it would force you to use the correct default allocator when needed.

Second, I don't think we can add a specialization based default allocator, because that can't be retrofitted on, Box::new wouldn't be allowed to create the default allocator (as it does with Global), because it can't know if TypeAllocator was specialized. This would require something like Default for TypeAllocated::Allocator, possibly more.

But in any case, Box<T> is currently guaranteed to have the same layout as NonNull<T>, and there is code in the wild that relies on this guarantee, so I don't think we can change the default allocator significantly.

[rinon]

Do you think that's too much of a limitation?

To be honest, I think it is. Since these structs are defined in the standard library, I don't feel such a limitation should exist. Especially since this is only a specific use case.

Fair, but I think the restriction makes sense. The "owner" of a type (its crate) could declare that it should be allocated with a certain allocator by default, probably guarded by a feature flag. This would push consumers of that type to use the allocator that was selected, rather than a generic other allocator. If you want to do this for a foreign type, just make a local newtype.

I still believe that #79 can solve this problem more elegantly. I guess it will not be possible to change Box::new in a backward compatible way, so for another allocator you will most likely have to define a new method, for example Box::new_in. Whether an allocator or a buffer is given here is not important for the time being, it only has to be typed. The advantage of the typed-buffer approach is that a buffer can be defined everywhere and can be used in any other place. For example a buffer which is optimized for integral values (SIMD for example?), then the trait Buffer<Integral> is implemented. This buffer can be used only with Box<Integral> then. Such a buffer could rely on an allocator, too.

I'm fairly certain it is possible to change Box::new in a backwards compatible way. What #79 proposes, in my understanding, is changing the backing store for things like Vec and Box from a pointer to a generic type that may be a pointer or may be something more interesting, e.g. a union of an inline buffer and a heap pointer. This doesn't affect my use case, as what I would like can already be done by using Box<T, SomeTypeBasedAllocator> for relevant types. #79 would let you use the same allocator everywhere and effectively use different heaps, but it would still require explicitly specifying the buffer (heap) at every allocation site. It is a different syntax for the same thing that we could already do in the case where the backing store is still just a pointer.

If security is your biggest concern, then I would create a newtype wrapper around Box (or Vec, or others), and use that throughout your codebase. This newtype can have the semantics that you desire, and it would force you to use the correct default allocator when needed.

That sort of works (as long as you duplicate the entire Box API with boilerplate) until you need to pass that Box or Vec to another crate, as it is no longer the expected type. On a related note, using collections from other crates is hard because even if they support custom allocators, you have to wrap them in boilerplate.

Second, I don't think we can add a specialization based default allocator, because that can't be retrofitted on, Box::new wouldn't be allowed to create the default allocator (as it does with Global), because it can't know if TypeAllocator was specialized. This would require something like Default for TypeAllocated::Allocator, possibly more.

Why not? I added the Default constraint for TypeAllocated::TypeAllocator for that exact reason. The implementation of Box::new(x) is actually just box x, a compiler builtin. It's not hard to check the trait impls of x in that builtin at that point in the compiler.

But in any case, Box is currently guaranteed to have the same layout as NonNull, and there is code in the wild that relies on this guarantee, so I don't think we can change the default allocator significantly.

I'm not proposing that at all. The layout of Box would not change here, it would still be just a pointer. Changing the allocator doesn't affect that.

[RustyYato]

That sort of works (as long as you duplicate the entire Box API with boilerplate) until you need to pass that Box or Vec to another crate, as it is no longer the expected type. On a related note, using collections from other crates is hard because even if they support custom allocators, you have to wrap them in boilerplate.

Given that custom allocators are still on nightly, I don't think many crates actually support it well (if at all). So... I don't think that's a great argument. You can still add some API for accessing external crates onto your newtype, and that would ease the boilerplate. The only thing big downside I see is adding the various traits that are on Box, (the problem with all newtypes), and I agree that there isn't a good solution to this.

Why not? I added the Default constraint for TypeAllocated::TypeAllocator for that exact reason. The implementation of Box::new(x) is actually just box x, a compiler builtin. It's not hard to check the trait impls of x in that builtin at that point in the compiler.

Ok, I think I just missed that, my bad.

I'm not proposing that at all. The layout of Box would not change here, it would still be just a pointer. Changing the allocator doesn't affect that.

Adding an allocator that is non-zero sized or has an alignment larger than usize will change the layout of Box<T>. Remember, Box is basically just this:

struct Box<T: ?Sized, A: Allocator> {
    ptr: NonNull<T>,
    allocator: A,
}

(ignoring minor details that don't alter the layout)

[rinon]

Given that custom allocators are still on nightly, I don't think many crates actually support it well (if at all). So... I don't think that's a great argument.

What I was hoping for was to be able to change allocation behavior in external crates as long as they assumed the default allocator. It's looking like that's tricky due to the layout issues you point out.

Adding an allocator that is non-zero sized or has an alignment larger than usize will change the layout of Box

Sorry I've been misunderstanding, my bad! I somehow missed that we're actually adding the allocator to the Box structure rather than just tracking the type of the allocator in the type system. Yeah, what I'm proposing would only work for zero-sized allocators then. I don't think we have a good way to constrain that, is there?

[RustyYato]

No, there isn't a way to constrain it to only zero sized types, until we get const generics