Repository: libsql/libsql
File: libsql-server/src/connection/write_proxy.rs (Rust)

Code before:

        let cb_context = Arc::new(parking_lot::Mutex::new((response_cb, arg)));
        while let Some(resp) = self.response_stream.next().await {
            match resp {
                Ok(resp) => {
                    // there was an interuption, and we moved to the next query
                    if resp.request_id > request_id {
                        return Err(Error::PrimaryStreamInterupted);
                    }

                    // we can ignore response for previously interupted requests
                    if resp.request_id < request_id {
                        continue;
                    }

                    let should_continue = tokio::task::spawn_blocking({
                        let cb_context = cb_context.clone();
                        move || -> Result<bool> {
                            let mut ctx = cb_context.lock();
                            let (ref mut response_cb, ref mut arg) = *ctx;
                            if !response_cb(resp.response.ok_or(Error::PrimaryStreamMisuse)?, arg)?
                            {
                                Ok(false)
                            } else {
                                Ok(true)
                            }
                        }
                    })
                    .await
                    .unwrap()?;

                    if !should_continue {
                        break;
                    }
                }
                Err(e) => {
                    tracing::error!("received an error from connection stream: {e}");
                    return Err(Error::PrimaryStreamDisconnect);
                }
            }
        }

        let Ok(mutex) = Arc::try_unwrap(cb_context) else {
            panic!("failed to get ownership on callback context")
        };
        let (_, arg) = mutex.into_inner();

        Ok(arg)
    }

    async fn execute<B: QueryResultBuilder>(

Review comment (suggestion):
```suggestion
                            response_cb(resp.response.ok_or(Error::PrimaryStreamMisuse)?, arg)
```

Code after:

        let cb_context = Arc::new(parking_lot::Mutex::new((response_cb, arg)));
        while let Some(resp) = self.response_stream.next().await {
            match resp {
                Ok(resp) => {
                    // there was an interuption, and we moved to the next query
                    if resp.request_id > request_id {
                        return Err(Error::PrimaryStreamInterupted);
                    }

                    // we can ignore response for previously interupted requests
                    if resp.request_id < request_id {
                        continue;
                    }

                    let should_continue = tokio::task::spawn_blocking({
                        let cb_context = cb_context.clone();
                        move || -> Result<bool> {
                            let mut ctx = cb_context.lock();
                            let (ref mut response_cb, ref mut arg) = *ctx;
                            response_cb(resp.response.ok_or(Error::PrimaryStreamMisuse)?, arg)
                        }
                    })
                    .await
                    .unwrap()?;

                    if !should_continue {
                        break;
                    }
                }
                Err(e) => {
                    tracing::error!("received an error from connection stream: {e}");
                    return Err(Error::PrimaryStreamDisconnect);
                }
            }
        }

        let Ok(mutex) = Arc::try_unwrap(cb_context) else {
            panic!("failed to get ownership on callback context")
        };
        let (_, arg) = mutex.into_inner();

        Ok(arg)
    }

    async fn execute<B: QueryResultBuilder>(
        &mut self,
        program: Program,
        builder: B,
    ) -> crate::Result<(B, TxnStatus, Option<FrameNo>)> {
        let mut txn_status = TxnStatus::Invalid;