flightsql.BaseServer does not exhaust channel in case of an error writing StreamChunks in DoGet

[arnoldwakim]

Describe the bug, including details regarding any error messages, version, and platform.

arrow-go/v18.3.1 on linux/x86-64

Hey everyone,

As this is my first time discussing on this repository, kindly forgive me in case of a mistake.

While exploring Arrow Flight SQL servers with DuckDB, I came across an issue.

Let's assume we have a Flight SQL client that sends SQL queries to a FlightSQL server, while the data is being streamed here (sqlite example), if the client receives a non handled SIGINT, the communication between server and client stops (without a cancel request), and the server catches the error here.
However, I still have an arrow array somewhere buffered in memory (server side) and the record reader that has not been released and therefore I still have arrow arrays in the wild waiting to be freed.

When modifying the source code here, while checking for an error adding a for loop that exhaust the channel to release the StreamChunk resolves the issue.

func exhaustChannelAndReleaseStreamChunks(cc <-chan flight.StreamChunk) {
	for chunk := range cc {
		chunk.Data.Release()
	}
}

if err = wr.WriteWithAppMetadata(chunk.Data, chunk.AppMetadata); err != nil {
    exhaustChannelAndReleaseStreamChunks(cc)
    return err
}

Did I misunderstand something?
I can take care of the PR if need be.
I've been testing that manually. Do you think I can test the behaviour rigorously?

Thank you for your time.

Component(s)

flight/flightsql

[zeroshade]

That's a good catch!

Rather than adding an explicit function like that, we can handle this with a defer:

defer func() {
        for chunk := range cc {
                if chunk.Data != nil {
                      chunk.Data.Release()
               }
        }
}()

wr := flight.NewRecordWriter(stream, ipc.WithSchema(sc))
defer wr.Close()

for chunk := range cc {
     ....
}

that should ensure we always drain it regardless of which condition occurs.

I can take care of the PR if need be.

That would be fantastic! I should automatically get tagged for review, but feel free to manually tag me if it doesn't do that.

I've been testing that manually. Do you think I can test the behaviour rigorously?

It should hopefully not be too difficult to use the testServer framework in server_test.go to manually create a client that makes a DoGet request and then closes the connection such that a CheckedAllocator in the server could detect that some records weren't released (failing without the above fix). If you need help setting it up reach out on the PR and I'll try to find some time to help put something together.