Handle "Unexpected response type for remote procedure call: Close" on query stream opening

[farost]

Usage and product changes

Fix a rare InternalError returned by mistake when a client sends a query request while the transaction is being closed. Now, an expected "The transaction is closed and no further operation is allowed." error is returned instead.

Additionally, wait for specific transaction responses in rollback, commit, and query to solidify the protocol and ensure that the server acts as expected.

Implementation

We faced this error when running a huge number of queries against a server that eventually went to sleep. Presumably, this state led to the GRPC stream interruption, and it "gracefully" stopped without errors. However, the driver's code did not expect that it could happen right after the stream opening request, while it should be considered a normal situation.

The source code where Close responses are generated:

async fn listen_loop(
    mut grpc_source: Streaming<transaction::Server>,
    collector: ResponseCollector,
    shutdown_sink: UnboundedSender<()>,
) {
    loop {
        match grpc_source.next().await {
            Some(Ok(message)) => collector.collect(message).await,
            Some(Err(status)) => break collector.close_with_error(status.into()).await,
            None => break collector.close().await, // <---- here, the stream has ended
        }
    }
    shutdown_sink.send(()).ok();
}

async fn close(self) {
  self.is_open.store(false);
  let mut listeners = std::mem::take(&mut *self.callbacks.write().unwrap());
  for (_, listener) in listeners.drain() {
      listener.finish(Ok(TransactionResponse::Close)); // <--- here, the only `Close` response is sent
  }
  // ....
}

Additionally, we explicitly wait for request's response type or stream closed marker with their correct processing for each other transaction request in transaction_stream. For this, the commit behavior of the server has been made more explicit: typedb/typedb#7484

[dmitrii-ubskii]

Is this the only place that needs to handle Close? What happens in any of the one-shot requests (like rollback) if the stream is closed between sending the request and receiving a response?

[farost]

It's a fair point. It looks like we currently ignore it for single requests, and I thought that it's normal. But, while it's fine for rollbacks (the next request will fail anyway), it might not be good for commits, as there is a chance that the client will think that everything is good:

    pub(crate) fn commit(self: Pin<Box<Self>>) -> impl Promise<'static, Result> {
        let promise = self.single(TransactionRequest::Commit);
        promisify! {
            let _this = self; // move into the promise so the stream isn't dropped until the promise is resolved
            resolve!(promise).map(|_| ()) // <--- here, `TransactionResponse::Close` becomes `()` just like any other
        }
    }

    pub(crate) fn rollback(&self) -> impl Promise<'_, Result> {
        let promise = self.single(TransactionRequest::Rollback);
        promisify! { resolve!(promise).map(|_| ()) } // <--- here, `TransactionResponse::Close` becomes `()` just like any other
    }

I will add the relevant processing here, too.

[farost]

Or maybe it's actually fine. I will test it more tomorrow.

[farost]

Turns out Commit should explicitly wait for Close without any Commit responses. Weird. Will decide if the server needs an update and then finalize this PR.

[farost]

@dmitrii-ubskii It was a good flag. I additionally covered other transaction requests and fixed an incorrect server's behavior: typedb/typedb#7484