[logachev]

Ensure when ResultSet is closed, we stop reading arrow stream as well

[gemini-code-assist]

Code Review

This pull request improves resource cleanup, concurrency handling, and exception safety across BigQueryArrowResultSet, BigQueryJsonResultSet, and BigQueryStatement. Key changes include adding idempotency checks and resource cleanup in result set close() methods, replacing uninterruptible queue operations with standard interruptible operations, and introducing a helper method isWorkerShuttingDown() to check thread and executor states. Feedback suggests wrapping the custom resource cleanup in BigQueryArrowResultSet.close() in a try-finally block to guarantee super.close() is always executed, and leveraging the new isWorkerShuttingDown() helper in BigQueryStatement to simplify checks and avoid potential NullPointerExceptions.

[gemini-code-assist]

Wrap the custom resource cleanup in a try-finally block to guarantee that super.close() is always executed, even if an exception is thrown during the cleanup of these resources.

  @Override
  public void close() {
    if (this.isClosed) {
      return;
    }
    LOG.fineTrace("close", () -> String.format("Closing BigqueryArrowResultSet %s.", this));
    this.isClosed = true;
    try {
      if (ownedTask != null) {
        // cancel the producer task when result set is closed
        ownedTask.cancel(true);
      }
      if (this.buffer != null) {
        this.buffer.clear();
      }
      if (this.arrowDeserializer != null) {
        this.arrowDeserializer.close();
      } else if (this.allocator != null) {
        this.allocator.close();
      }
    } finally {
      super.close();
    }
  }

References

1. When managing a collection of closeable resources, ensure they are closed in an exception-safe manner to prevent resource leaks, meaning all opened resources should be closed even if exceptions occur during their creation or closing.

[gemini-code-assist]

Use the newly introduced isWorkerShuttingDown() helper method instead of repeating the Thread.currentThread().isInterrupted() || executor.isShutdown() check. This improves readability and safely handles potential null values for connection and executor.

Suggested change

	if (Thread.currentThread().isInterrupted() || executor.isShutdown()) {
	arrowBatchWrapperBlockingQueue.clear();
	break;
	}
	if (isWorkerShuttingDown()) {
	arrowBatchWrapperBlockingQueue.clear();
	break;
	}

References

1. When using lazily initialized resources (such as ExecutorService), ensure that teardown or close methods perform explicit null checks before invoking methods on them to prevent NullPointerException.

[gemini-code-assist]

Use !isWorkerShuttingDown() instead of !Thread.currentThread().isInterrupted() && !executor.isShutdown() to simplify the condition and safely handle potential null values for connection and executor.

Suggested change

	if (!Thread.currentThread().isInterrupted() && !executor.isShutdown()) {
	streamCompleted = true;
	} else {
	arrowBatchWrapperBlockingQueue.clear();
	}
	if (!isWorkerShuttingDown()) {
	streamCompleted = true;
	} else {
	arrowBatchWrapperBlockingQueue.clear();
	}

References

1. When using lazily initialized resources (such as ExecutorService), ensure that teardown or close methods perform explicit null checks before invoking methods on them to prevent NullPointerException.

[gemini-code-assist]

Use isWorkerShuttingDown() instead of Thread.currentThread().isInterrupted() || executor.isShutdown() to safely check if the worker is shutting down, avoiding potential NullPointerException if executor is null.

Suggested change

	if (Thread.currentThread().isInterrupted() || executor.isShutdown()) {
	rpcResponseQueue.offer(endOfRpcTuple);
	} else {
	try {
	rpcResponseQueue.put(endOfRpcTuple);
	} catch (InterruptedException ie) {
	Thread.currentThread().interrupt();
	}
	}
	if (isWorkerShuttingDown()) {
	rpcResponseQueue.offer(endOfRpcTuple);
	} else {
	try {
	rpcResponseQueue.put(endOfRpcTuple);
	} catch (InterruptedException ie) {
	Thread.currentThread().interrupt();
	}
	}

References

1. When using lazily initialized resources (such as ExecutorService), ensure that teardown or close methods perform explicit null checks before invoking methods on them to prevent NullPointerException.