fix(csharp/src/Drivers/Databricks): Correct DatabricksCompositeReader and StatusPoller to Stop/Dispose Appropriately

[toddmeng-db]

Motivation

The following cases are not properly stopping or disposing the status poller:

1. If the DatabricksCompositeReader is explicitly disposed by the user
2. CloudFetchReader is done returning results
3. Edge case terminal operation status (timedout_state, unknown_state)

In addition:

• When DatabricksOperationStatusPoller.Dispose(), it may cancel the GetOperationStatusRequest in the client. If the input buffer has data and cancellation is triggered, it leaves the TCLI client with unconsumed/unsent data in the buffer, breaking subsequent requests (fixed in this PR)

Fixes

DatabricksOperationStatusPollerLogic is now more appropriately managed by DatabricksCompositeReader (moved out of BaseDatabricksReader) to handle all cases where null results (indicating completion) are returned.

Disposing DatabricksCompositeReader appropriately disposes the activeReader and statusPoller

TODO

Follow-up PR - when statement is disposed, it should also dispose the reader (the poller is currently stopped when operationhandle is set to null, but this should also happen explicitly)

Need add some unit testing (follow up pr: #3243)

[jadewang-db]

can you confirm, even without this fix, the polling will stop after statement being disposed, right? if not, we need fix there also

[toddmeng-db]

From testing:
Small nit but I think we need to avoid this here, since it's possible that DownloadQueue is full, then exception handling would be stuck. Should I modify the Exception handling below, or was there a reason why it it like this? (line 262) @jadewang-db

catch (Exception ex)
            {
                try
                {
                    _downloadQueue.Add(EndOfResultsGuard.Instance, CancellationToken.None);
                 }
}

Alternatively, we can create a new CancellationToken with Timeout for this attempt

                    CancellationToken GetOperationStatusTimeoutToken = ApacheUtility.GetCancellationToken(_requestTimeoutSeconds, ApacheUtility.TimeUnit.Seconds);

[jadewang-db]

thanks, a cancellation token looks good

[toddmeng-db]

Oh just saw this comment, let me implement

[toddmeng-db]

Actually looks like TryAdd is better suited here

[toddmeng-db]

We need to use a timeout token here, instead of cancelling when canceltoken is triggered; if an interrupt is triggered prematurely, the TCLI client may still have unsent/unconsumed results in the buffers, affecting subsequent calls with that client (which is any future call in the same Session)

[jadewang-db]

are you able to repro this? should we do this to all the thrift rpc calls in the driver?

[toddmeng-db]

I think it is because in THTTPTransport (used by SparkHttpConnection -> DatabricksHttpconnection), a new Stream is created when the request is flushed. If cancellation happens before this, that stream doesn't get discarded:
https://github.com/apache/thrift/blob/master/lib/netstd/Thrift/Transport/Client/THttpTransport.cs#L281

Yes, during testing, got some errors. In the proxy logs, I remember seeing requests sent out with both GetOperationStatus and CloseOperationStatus (in the same request) while testing another PR

I think we are safe in HiveServer2Statement, but we might need to adjust CancellationToken in DatabricksReader, CloudFetchResultFetcher, and DatabricksCompositeReader

[toddmeng-db]

Actually, I think this depends a bit on how CancellationToken could be used by PBI, too
@CurtHagenlocher will mashup ever trigger cancellationTokens passed into IArrowStreamReader.ReadNextBatchAsync? Do we need to ensure that the connection still remains usable for subsequent statements?

[toddmeng-db]

At least for now, I think we can operate this way:

1. If the user cancels the token passed in to ReadNextBatchAsync, we should not to break the client
2. Dispose() should not break the client either

[CurtHagenlocher]

@CurtHagenlocher will mashup ever trigger cancellationTokens passed into IArrowStreamReader.ReadNextBatchAsync? Do we need to ensure that the connection still remains usable for subsequent statements?

This is currently unimplemented but we'll need to implement it before GA for parity with the ODBC implementation. What is probably most important for cancellation is query execution, and unless we manage to push forward the proposed ADBC 1.1 API, currently the only way to cancel a running query is to call AdbcStatement.Cancel. There is currently no implementation of this method for any of the C#-implemented drivers :(.

[CurtHagenlocher]

From a Power BI perspective, the most important use of cancellation is for Direct Query because users can generate a lot of queries simply by clicking around in a visual and in-progress queries will need to be cancelled if their output is no longer needed. DQ output tends to be relatively small, so being able to cancel in the middle of reading the output is arguably less important than being able to cancel before the results start coming back.

[jackyhu-db]

should you respect the connection parameter DatabricksParameters.CloudFetchTimeoutMinutes instead of the default value?

[toddmeng-db]

Do you mean I shouldn't create a new constant here?

[jackyhu-db]

no, what I meant is if you should check the value of the connection parameter CloudFetchTimeoutMinutes (adbc.databricks.cloudfetch.timeout_minutes) which can be set by the client and customer.

[toddmeng-db]

Oh got it, that makes sense, it should be a configurable parameter. To be consistent with the rest of HiveServer2Statement, I'm just using the QueryTimeout parameter (which is what other FetchResultsRequest uses)

I have some changes in a follow-up PR that will make this change easier to do for DatabricksReader, will leave this as a TODO

[CurtHagenlocher]

Thanks! The linter error needs to be fixed and I made a few small low-priority suggestions.