How to safely implement reusable scratch memory in C++?

Viewed 298

It is very common that even pure functions require some additional scratch memory for their operations. If the size of this memory is known at compile time, we can allocate this memory on the stack with std::array or a C array. But the size often depends on the input, so we often resort to dynamic allocations on the heap through std::vector. Consider a simple example of building a wrapper around some C api:

void addShapes(std::span<const Shape> shapes) {
    std::vector<CShape> cShapes;
    cShapes.reserve(shapes.size());

    // Convert shapes to a form accepted by the API
    for (const Shape& shape : shapes) {
        cShapes.push_back(static_cast<CShape>(shape));
    }
    cAddShapes(context, cShapes.data(), cShapes.size());
}

Let's say that we call this function repeatedly and that we identify that the overhead of std::vector memory allocations is significant, even with the call to reserve(). So what can we do? We could declare the vector as static to reuse the allocated space between calls, but that comes with several problems. First, it is no longer thread safe, but that can be fixed easily enough by using thread_local instead. Second, the memory doesn't get released until the program or thread terminates. Let's say we are fine with that. And lastly, we have to remember to clear the vector every time, because it's not just the memory that will persist between function calls, but the data as well.

void addShapes(std::span<const Shape> shapes) {
    thread_local std::vector<CShape> cShapes;
    cShapes.clear();

    // Convert shapes to a form accepted by the API
    for (const Shape& shape : shapes) {
        cShapes.push_back(static_cast<CShape>(shape));
    }
    cAddShapes(context, cShapes.data(), cShapes.size());
}

This is the pattern I use whenever I would like to avoid the dynamic allocation on every call. The issue is, I don't think the semantics of this are very apparent if you aren't aware of the pattern. thread_local looks scary, you have to remember to clear the vector and even though the lifetime of the object now extends beyond the scope of the function, it is unsafe to return a reference to it, because another call to the same function would modify it.

My first attempt to make this a bit easier was to define a helper function like this:

template <typename T, typename Cleaner = void (T&)>
T& getScratch(Cleaner cleaner = [] (T& o) { o.clear(); }) {
    thread_local T scratchObj;
    cleaner(scratchObj);
    return scratchObj;
}

void addShapes(std::span<const Shape> shapes) {
    std::vector<CShape>& cShapes = getScratch<std::vector<CShape>>();

    // Convert shapes to a form accepted by the API
    for (const Shape& shape : shapes) {
        cShapes.push_back(static_cast<CShape>(shape));
    }
    cAddShapes(context, cShapes.data(), cShapes.size());
}

But of course, that creates a thread_local variable for each template instantiation of the getScratch function, rather than for each place the function is called. So if we asked for two vectors of the same type at once, we'd get two references to the same vector. Not good.

What would be a good way to implement this sort of a reusable memory safely and cleanly? Are there already existing solutions? Or should we not use thread local storage in this way and just use local allocations despite the performance benefits that reusing them brings: https://quick-bench.com/q/VgkPLveFL_K5wT5wX6NL1MRSE8c ?

2 Answers

To answer my own question, I came up with a solution that builds upon the last example. Rather than keeping only one object for each thread and type, lets keep a free list of them. Upon request, we either reuse an object from the free list or create a new one. The user keeps a RAII-style handle that returns the object into the free list when it leaves the scope. Since we still use thread_local, this is thread safe without any effort. We can wrap all this into a simple class:

template <typename T>
class Scratch {
public:
    template <typename Cleaner = void (T&)>
    explicit Scratch(Cleaner cleaner = [] (T& o) { o.clear(); }) : borrowedObj(acquire()) {
        cleaner(borrowedObj);
    }
    
    T& operator*() {
        return borrowedObj;
    }
    T* operator->() {
        return &borrowedObj;
    }
    
    ~Scratch() {
        release(std::move(borrowedObj));
    }
private:
    static thread_local std::vector<T> freeList;
    T borrowedObj;

    static T acquire() {
        if (!freeList.empty()) {
            T obj = std::move(freeList.back());
            freeList.pop_back();
            return obj;
        } else {
            return T();
        }
    }
    static void release(T&& obj) {
        freeList.push_back(std::move(obj));
    }
};

That can be used simply as:

void addShapes(std::span<const Shape> shapes) {
    Scratch<std::vector<CShape>> cShapes;

    // Convert shapes to a form accepted by the API
    for (const Shape& shape : shapes) {
        cShapes->push_back(static_cast<CShape>(shape));
    }
    cAddShapes(context, cShapes->data(), cShapes->size());
}

You might want to extend this as needed, perhaps add a [] operator for convenience if it's going to be used with containers. You could keep its intended use to be a local object in a function and explicitly make it non-copyable and non-movable, or it could be turned into a general purpose handle like unique_ptr. But beware that the object must be destroyed by the same thread that created it.

In both cases it addresses my issues with a raw thread_local. The clear is implicit and returning a reference to the scratch object or its data is now obviously wrong. It still doesn't automatically free memory, which is what we want after all, but at least it's now easier to implement the functionality to free it on demand as needed.

In general, it should have lower memory usage than the raw thread_local method, too, since allocations of the same type can be reused across different call sites. But there is a scenario in which this behavior will result in a higher memory usage, too. Let's say we have a function that needs a std::vector<int> of size 10000. If we call this function and then ask for a vector of the same type, we will get the one with capacity 10000. If we then call the function again while holding this vector, it will have to create another one, resizing it to 10000 elements, too.

For those reasons I would recommend using it only where you don't expect to see large amounts of data, but rather want to avoid lots of small, but frequent and short-lived allocations.

static to reuse the allocated space between calls, but that comes with several problems. First, it is no longer thread safe, but that can be fixed easily enough by using thread_local instead. Second, the memory doesn't get released until the program or thread terminates.

Exactly. Because only the user of the function knows how and when he wants to call the function and when he wants to do it, only the user of the function should be the one responsible for reusing space if he wants to and for clearing it up, because the user knows if he is going to use it later or not. So add cache object to your function, where you cache the state to speed it up later.

void addShapes(std::span<const Shape> shapes, std::vector<CShape>& cache) {
    cache.reserve(shapes.size());    
    // Convert shapes to a form accepted by the API
    for (const Shape& shape : shapes) {
        cache.push_back(static_cast<CShape>(shape));
    }
    cAddShapes(context, cache.data(), cache.size());
}

Or you could objectify it a bit, like:

class shapes {
    std::vector<CShape> cache;
    void add(std::span<const Shape> shapes) {
        cache.reserve(shapes.size());    
        // Convert shapes to a form accepted by the API
        for (const Shape& shape : shapes) {
            cache.push_back(static_cast<CShape>(shape));
        }
       cAddShapes(context, cache.data(), cache.size());
    }
   void clear_cache() {
      cache.clear();
   }
};
Related