diff --git a/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProvider.java b/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProvider.java index 176958fd..308d8826 100644 --- a/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProvider.java +++ b/jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProvider.java @@ -56,13 +56,14 @@ public void close() { channel.shutdown(); try { - channel.awaitTermination(5, TimeUnit.SECONDS); - } catch (InterruptedException e) { - log.error("Failed to shutdown channel within 5 seconds", e); - } finally { - if (!channel.isTerminated()) { + if (!channel.awaitTermination(5, TimeUnit.SECONDS)) { + log.warn("Channel did not terminate within 5 seconds, forcing shutdown"); channel.shutdownNow(); } + } catch (InterruptedException e) { + log.warn("Channel shutdown interrupted, forcing immediate termination", e); + channel.shutdownNow(); + Thread.currentThread().interrupt(); } } } diff --git a/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProviderTest.java b/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProviderTest.java index b017b9e0..88171df3 100644 --- a/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProviderTest.java +++ b/jdbc-core/src/test/java/com/salesforce/datacloud/jdbc/core/JdbcDriverStubProviderTest.java @@ -173,7 +173,7 @@ void shouldNotEnableRetriesWhenDisabled() { void callsManagedChannelCleanup() { val mockChannel = mock(ManagedChannel.class); val mockChannelBuilder = getMockChannelBuilderWithChannel(mockChannel); - when(mockChannel.isTerminated()).thenReturn(false, true); + when(mockChannel.awaitTermination(anyLong(), any(TimeUnit.class))).thenReturn(true); val stubProvider = JdbcDriverStubProvider.of(mockChannelBuilder); stubProvider.close(); @@ -185,10 +185,10 @@ void callsManagedChannelCleanup() { @SneakyThrows @Test - void callsManagedChannelShutdownNow() { + void callsManagedChannelShutdownNowOnTimeout() { val mockChannel = mock(ManagedChannel.class); val mockChannelBuilder = getMockChannelBuilderWithChannel(mockChannel); - when(mockChannel.isTerminated()).thenReturn(false); + when(mockChannel.awaitTermination(anyLong(), any(TimeUnit.class))).thenReturn(false); val stubProvider = JdbcDriverStubProvider.of(mockChannelBuilder); stubProvider.close(); @@ -197,4 +197,25 @@ void callsManagedChannelShutdownNow() { verify(mockChannel).awaitTermination(5, TimeUnit.SECONDS); verify(mockChannel).shutdownNow(); } + + @SneakyThrows + @Test + void handlesInterruptDuringShutdown() { + val mockChannel = mock(ManagedChannel.class); + val mockChannelBuilder = getMockChannelBuilderWithChannel(mockChannel); + when(mockChannel.awaitTermination(anyLong(), any(TimeUnit.class))) + .thenThrow(new InterruptedException("Test interruption")); + + val stubProvider = JdbcDriverStubProvider.of(mockChannelBuilder); + stubProvider.close(); + + verify(mockChannel).shutdown(); + verify(mockChannel).awaitTermination(5, TimeUnit.SECONDS); + verify(mockChannel).shutdownNow(); + + // Verify interrupt status is restored per Java best practices + assert Thread.currentThread().isInterrupted() : "Thread interrupt status should be restored"; + // Clear interrupt status for other tests + Thread.interrupted(); + } }