https://www.boundaryml.com logo
🟡 Memory Leak in BAML Go Bindings/CFFI
# questions
i
Hey folks, I'm posting this here instead of immediately logging it as a GitHub issue, as I am not even sure how to properly describe what (or rather, where) the issue is, but the tl;dr is there appears to be a severe memory leak in the BAML Go bindings/CFFI. I can't even place what exactly is leaking (from my outsider POV, it looks like it's a universal issue with the
raw_ptr_wrapper
abstraction, but the leak only really appears to be blatantly obvious when interacting with `FunctionLog`/`OnTick` callbacks, so I'm not quite sure), but something definitely is. Here's a small reproduction that appears to consistently trigger a memory leak (+ other disturbing memory management-related behavior) on my darwin arm64 machine: https://github.com/invakid404/baml-memory-leak-repro The repo includes a fake OpenAI server to simulate a large streaming response, the logs I applied to the BAML src that made me think this is an actual memory leak and not just the OS being lazy with memory reclamation, and a main.go file that triggers the leak. Formal warning, if you do apply the logs from the included patch, be prepared to be greeted by a lot of logs I have nowhere near enough context about the project to be able to put a pin in it, so I'm hoping you can help me out a bit here. Initially observed on 0.214.0 (we're still stuck on an older version due to multiple panic conditions that I reported as well, hopefully resolved in latest by now), confirmed to still be reproducible on latest canary. (PS: the only reason I went on this witch hunt is because one of our Go services that interacts with BAML was and is going up in memory slowly but surely until it OOMs, so this is very much an issue in practice as well, although it happens much more slowly with real LLM requests; the extreme SSE chunk sizes in the fake server are simply there to make the issue obvious)
v
thanks for reporting this!
let me confirm once i repro this!
hmm, so i'm seeing it stay at:
Copy code
[Periodic GC] Alloc: 0 MB, Sys: 19 MB, HeapInuse: 2 MB, HeapReleased: 8 MB
i
wait, this is insanely weird
if i do
go run main.go
, memory stays stable for me as well
if i instead do
go build -o baml-test && ./baml-test
, it appears to be a different story
what..? :D
v
well then..
i
oh nvm i get it
i think the logs in main.go don't really show anything
v
huh
v
when i run bmal-test, i still get stable memory
i
here is what i get
i
by that i meant i don't think the memory stats log shows the actual process memory
i.e. the log itself is sort of useless :D
v
that said i have a much better way to debug this
oh
i
if you apply my baml patch (also in the repo) you'll see that there's definitely something funny going on and i'm not hallucinating, i promise :D
v
i'm just gonna log every baml objects constructions + destruction
yea i believe ya!
i
some arcs' strong counts shoot up to numbers in the thousands
while attempting to destroy them
v
wow you made my life much easier
ffi's are just hard sadly. luckily once we solve it, it becomes a pretty stable layer that we never have to touch again
i
it seems like it indeed, i was hoping to be able to come up with a fix for this as well, but it was indeed way too complicated for me to figure out in the time i had available :D
v
leak confirmed
i think patch should be quite doable!
there is one more
i'll hopefully identify it asap
but whatever is it is is <<< than the prior one for sure
Copy code
python3 analyze_ffi.py rust.log go.log     
Loaded 299328 events from rust.log
Loaded 296040 events from go.log

Total: 595368 FFI events, 20 callback events, 0 error events

=== FFI Lifecycle Analysis ===

Rust side:
  FFI_CREATE:       2770
  FFI_WRAP_ARC:    95910
  FFI_GIVE_GO:     98680  <- Rust gives ptr to Go
  FFI_GO_RELEASE:  98680  <- Rust destructor called
  FFI_FREE:         3288  <- Arc fully freed

Go side:
  GO_RECEIVE:          98680  <- Go receives object
  GO_DESTRUCTOR_START: 98680  <- Go finalizer starts
  GO_DESTRUCTOR_OK:    98680  <- Go finalizer succeeds
  GO_DESTRUCTOR_ERROR:     0  <- Go finalizer fails

=== NO RUST-SIDE LEAKS ===
All FFI_GIVE_GO have matching FFI_GO_RELEASE.

=== NO GO-SIDE LEAKS ===
All GO_RECEIVE have matching GO_DESTRUCTOR.

=== LEAK COMPARISON ===
Leaked on Rust side only: 0 pointers
Leaked on Go side only:   0 pointers
Leaked on both sides:     0 pointers

Callback map:
  Total ADDs: 10
  Total DELs: 10
  Final map size: 0

=== SUMMARY ===
GO OK: All 98680 objects were properly finalized
RUST OK: All 98680 pointers properly released
Houston we have liftoff
thanks so much for spotting this and your easy repro @invakid404 !
the bug was actually in the go layer as it turns out 🙁 lifetimes weren't correct and we held on to channels too long which caused this.
i
thank you for looking into it!
will be on the lookout for the next release
finally merged!
generally we've got a good fix and now we def don't leak memory. There is still 1 problem: go doesn't provide any guarantees when finalizers are run, so likely we need to provide explicit
.Close()
methods on all FFI handles that devs can then write
defer obj.Close()
with
i
yeah i was thinking of bringing that up, it does sound like a good idea i am personally fine with the memory management feeling "magical", but offering control won't hurt
i'll build from canary locally and test later today, thanks!
v
yea we should be able to add that (but it would be an explicit thing and i think it would make every method on every handle type fallible. but we can do that!
a
fixed in 0.218.0
5 Views