pg_allocator feature causes segfault

[bluejekyll]

I discovered that with a new plugin I’m working on the pg_allocator feature is causing a segfault in the library when loaded and run in Postgres.

I haven’t got more details at-the-moment. But what’s strange is it looks like the segfault is occurring before the allocator is called (based on attempts to log in the GlobalAlloc functions).

The work around is to disable the pg_allocator feature.

[idubrov]

I think, pg_allocator feature is perhaps unsound.

Here is an example that demonstrates the issue which I was talking about (after your talk):

use pg_extern_attr::pg_extern;
use pg_extend::pg_magic;
use std::sync::Mutex;
use lazy_static::lazy_static;

lazy_static! {
    // Smuggle allocated memory in a lazy static!
    static ref COUNTER: Mutex<Box<i32>> = Default::default();
}

pg_magic!(version: pg_sys::PG_VERSION_NUM);

#[pg_extern]
fn counter() -> i32 {
    let mut guard = COUNTER.lock().unwrap();
    **guard += 1;
    **guard
}

The issue here is that I'm sneaking allocated memory (in the form of a Box) in a global state, which gets freed by Postgres after invocation (or transaction?) completes. However, the connections in Postgres are potentially long-living, serving multiple transactions, which is a problem here.

The very first invocation would allocate that Box and subsequent calls will access that memory. However, after the first invocation completes, Postgres would release all the memory, rendering that memory invalid. The next invocation of that function will access memory which was "freed" before (use-after-free case).

The results I am getting:

postgres=# select pg_counter();
 pg_counter 
------------
          1
(1 row)

postgres=# select pg_counter();
 pg_counter 
------------
          2
(1 row)

postgres=# select pg_counter();
 pg_counter 
------------
         51
(1 row)

postgres=# select pg_counter();
 pg_counter 
------------
      12598
(1 row)

Since the code is fully safe Rust not doing anything out of ordinary (shared state is not a good thing by itself, but this could be some cache), I think, this allocator is unsound for Rust code.

Given that Rust code (in a theoretical vacuum) should be good about freeing the memory it uses, what are the benefits of using Postgres allocator?

One option to make it safe(-ish), I think, would be to have a global counter in the PgAllocator: bump some atomic when allocating, decrease when deallocating. When leaving the Rust context, check with the allocator if counter is not 0 and panic if so. This will detect the scenario I've written, but still will not allow allocating across function calls.

The naive check, though, is prone to race conditions: if Rust code starts threads, counter might be zero when we check but code might still allocate just before it leaves to Postgres. Not sure if it is worth chasing, but maybe the solution would be to have function like fn disarm() on allocator that will disable allocations/deallocations first (by setting some global atomic boolean?), then check the counter. Then, on entry from the Postgres world, code will call fn arm() on the allocator to enable allocations.

[bluejekyll]

I agree that there are issues. And yes, it's not threadsafe, though Postgres doesn't want you to use threads for anything, so it would be accurate to say spawning threads and using any postgres libraries would be considered unsafe. Not sure what we should do about that at-the-moment.

Now that you bring up these points, I hadn't considered the fact that it's not guarded, and multiple threads could easily be stomping on each other. This could absolutely be the issue. A mutex should be used in this context.

The only reason to use pg_allocator is to "play by the rules" of Postgres. I haven't built anything complex enough yet to discover if the default allocator would cause any issues, and given that Rust would be safely managing all of it it should be fine, except possibly where pointers are being handed back to postgres. Maybe we could come up with a strategy like you propose with arm/disarm at the FFI boundary? It might be complex to get that correct though. In the simple case, arm/disarm could be done when entering and leaving the FFI wrapper. And we'd have to find different points to do the same elsewhere...

[eeeebbbbrrrr]

I might can shed some light on this. And if y'all already know this, then just please ignore...

Postgres' palloc function allocates memory in whatever the CurrentMemoryContext is for the backend. In general, all memory contexts are children of TopTransactionContext, which is created new at the start of every transaction and free'd when the transaction finishes (COMMITs or ABORTs). There is a TopMemoryContext which is never free'd by the backend (it's akin to just using malloc())

(It looks to me like pg-extend-rs sets up Rust to always use palloc()? If that's the case, extensions built on this will always be using CurrentMemoryContext for all allocations? EDIT: I guess that's if the pg_alloc module is included?)

It's probably also worth noting that CurrentMemoryContext is not necessarily always pointing to TopTransactionContext. All sorts of things in Postgres can be assigned short-lived memory contexts, including function calls.

https://github.com/postgres/postgres/blob/master/src/backend/utils/mmgr/README

In the case of @idubrov's example, the static COUNTER will be allocated whenever the backend first loads the extension .so, using whatever CurrentMemoryContext is at that point in time.

Subsequent calls to select pg_counter(); are then basically doing use-after-free for the COUNTER. Which is, obviously, why the value starts to change wildly. I think if @idubrov's Postgres was configured with --enable-cassert it would have segfaulted on the second use.

For @bluejekyll's problem of things segfaulting (seemingly) before the allocator is even called is likely due to the same sort of thing. Rust likely allocated some memory, using palloc() when the .so was loaded and tried to use it later, long after it had been free'd by Postgres.

Anyways, I wonder if it makes sense for the allocator function in this library to require a MemoryContext argument, and then just use Postgres' MemoryContextAlloc() so the desired memory context can be used?

Also, it's worth noting that none of Postgres' memory allocation functions are thread-safe.

It seems reasonable to me that even tho Postgres doesn't expect threads in the backend, a Rust extension could be threaded so long as the threads were not allocating memory through Postgres -- their memory usage would need to be completely isolated.

It also seems reasonable to me to just not expose palloc/pfree if Rust can otherwise guarantee that anything it allocated will be free'd whenever it returns control back to Postgres? I'm about two-weeks new to Rust, so I dunno?

Anyways, I hope this is helpful. I've been doing Postgres extension/library development in C since PG v8.0, so I'm p familiar with how it works. However, I'm brand-spankin'-new to Rust. And also, I haven't actually used pg-extend-rs yet!

[bluejekyll]

Thank you for this background. It’s great. For threading related things, I was considering serializing access to the allocator (if possible). Do you think this would be safe?

[eeeebbbbrrrr]

So I just had a thought... Just take the approach Postgres does: If you want memory cleaned up whenever a MemoryContext is destroyed, use palloc, MemoryContextAlloc, MemoryContextSwitchTo, etc, otherwise just use malloc -- and let it leak!

IOW, don't set a GlobalAllocator for Rust, but instead expose all of Postgres' allocator functions for extension authors to use.

I was considering serializing access to the allocator (if possible). Do you think this would be safe?

Nope. Even if Rust is properly serialized around palloc and friends, Postgres isn't. So you could have the Postgres backend process thread allocating memory while some backround Rust thread is too.

EDIT: you also have Postgres' signal handlers to worry about, but I'm not exactly sure if they pause threads too?

I actually ran into this with @zombodb's integration with libcurl -- ended up just letting libcurl manage its memory with malloc/free and ensuring I installed a transaction callback (RegisterXactCallback()) to do any dangling cleanup.

And that might be part of the approach you should take too.

I also think pg-extend-rs needs some rule where it can't elog(ERROR)/ereport(ERROR) from within Rust. Postgres uses setjmp/longjmp to manage elog's, and lord knows what that would do to Rust. Dealing with Rust panic!'s is only part of the drama on this front!

[eeeebbbbrrrr]

Nope. Even if Rust is properly serialized around palloc and friends, Postgres isn't. So you could have the Postgres backend process thread allocating memory while some backround Rust thread is too.

I guess maybe there's an exception here...

If you can guarantee that control isn't returned to Postgres until the threads are all finished (ie, the Rust entry point for that function is joining on them), then yes, I think it would be safe (what about signal handler interruptions tho?), but man, it'd be terribly slow to synchronize memory allocation. I wouldn't actually want to use that.

[bluejekyll]

So you could have the Postgres backend process thread allocating memory while some backround Rust thread is too.

That’s an excellent point. I hadn’t considered it.

The really big deal here is transferring data from one context allocated in Rust (malloc/free) to the Postgres context (managed in some form through the pg allocator).

A minimum requirement we have is to figure out a way to allocate with pg for memory it needs to own, and then in Rust to never free memory owned by pg. it’s not an impossible task, but requires some thought on interfaces we’ll expose to do this.

[eeeebbbbrrrr]

The really big deal here is transferring data from one context allocated in Rust (malloc/free) to the Postgres context (managed in some form through the pg allocator).

So it would need to be some sort of Copy thing where the memory is memcpy'd from where ever Rust allocated it into a properly sized palloc'd buffer.

You could do some tricks with transaction callbacks to whole-sale free memory at the end of a transaction, but then you'd be holding Rust-allocated memory until the transaction finished, and not all transactions are single statements.

A minimum requirement we have is to figure out a way to allocate with pg for memory it needs to own, and then in Rust to never free memory owned by pg. it’s not an impossible task, but requires some thought on interfaces we’ll expose to do this.

Yeah, I think that's just exposing all of Postgres' various memory management functions as top-level Rust functions. Basically, whatever's declared in Postgres' memutils.h and palloc.h headers.

And then from there, if you can implement Postgres' HTAB, StringInfo, and List "objects" as unsafe Rust, that's likely 70% of what you'd need from Postgres internals.

[eeeebbbbrrrr]

Also, looks like Postgres has:

/* Registration of memory context reset/delete callbacks */
extern void MemoryContextRegisterResetCallback(MemoryContext context,
					   MemoryContextCallback *cb);

I've never used it, but the callback gets called as the comment says. Maybe that's a thing that'd be useful too?

Have Rust install a callback whenever it notices it allocates memory in a MemoryContext it hasn't seen yet? I dunno. Just thinking out loud.

[eeeebbbbrrrr]

Maybe you could do a thing where you, I dunno what' they're called in Rust, but annotate a function to indicate the MemoryContext you want it to use?

#[memory_context(this = "CurrentMemoryContext")]
pub fn foo() -> i32 {
   return 1;
}

That's a bad example, but what I'm thinking is that somehow pg-extend-rs would set the Rust allocator to use for that function call to be MemoryContextAlloc(<context_name>, size_t)?

If unspecified, the default would be CurrentMemoryContext, but you could specify others like TopTransactionContext, maybe even create a brand new context with a specific parent (via AllocSetContextCreate().

And then anything that Rust allocates that isn't directly called from Postgres (such as the COUNTER initialization in @idubrov's example) happens with Rust's default allocator?