[🍒][CDAP-21043] Retry on SocketTimeoutException in MessagingProgramStatePublisher.publish

cdap-app-fabric/pom.xml

@@ -149,6 +149,7 @@
     <dependency>
       <groupId>junit</groupId>
       <artifactId>junit</artifactId>
+      <version>4.13.2</version>
       <scope>test</scope>
     </dependency>
     <dependency>

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/program/MessagingProgramStatePublisher.java

@@ -42,6 +42,7 @@
 import io.cdap.cdap.proto.id.ProgramRunId;
 import io.cdap.cdap.proto.id.TopicId;
 import java.io.IOException;
+import java.net.SocketTimeoutException;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -142,9 +143,7 @@ public void publish(Notification.Type notificationType, Map<String, String> prop
             .build());
         LOG.trace("Published program status notification: {}", programStatusNotification);
         done = true;
-      } catch (IOException | AccessException e) {
-        throw Throwables.propagate(e);
-      } catch (TopicNotFoundException | ServiceUnavailableException e) {
+      } catch (TopicNotFoundException | ServiceUnavailableException | SocketTimeoutException e) {
         // These exceptions are retry-able due to TMS not completely started
         if (startTime < 0) {
           startTime = System.currentTimeMillis();
@@ -164,6 +163,8 @@ public void publish(Notification.Type notificationType, Map<String, String> prop
           Thread.currentThread().interrupt();
           done = true;
         }
+      } catch (AccessException | IOException e) {
+        throw Throwables.propagate(e);
       }
     }
   }

cdap-app-fabric/src/test/java/io/cdap/cdap/internal/app/program/MessagingProgramStatePublisherTest.java

@@ -29,6 +29,7 @@
 import io.cdap.cdap.proto.id.ApplicationId;
 import io.cdap.cdap.proto.id.ProgramRunId;
 import java.io.IOException;
+import java.net.SocketTimeoutException;
 import java.util.Map;
 import org.junit.Assert;
 import org.junit.Rule;
@@ -95,4 +96,65 @@ public void testMultipleTopics_noRun() throws TopicNotFoundException, IOExceptio
     StoreRequest storeRequest = storeRequestCaptor.getValue();
     Assert.assertEquals("programstatusevent0", storeRequest.getTopicId().getTopic());
   }
+
+  @Test
+  public void testPublishSuccessOnRetryableException() throws TopicNotFoundException, IOException {
+    CConfiguration cConf = CConfiguration.create();
+    cConf.setInt(Constants.AppFabric.PROGRAM_STATUS_EVENT_NUM_PARTITIONS, 1);
+    cConf.setInt("system.program.state.retry.policy.max.retries", 3);
+
+    MessagingService messagingService = Mockito.mock(MessagingService.class);
+    MessagingProgramStatePublisher publisher = new MessagingProgramStatePublisher(cConf, messagingService);
+    Mockito.when(messagingService.publish(Mockito.any()))
+           .thenThrow(new SocketTimeoutException())
+           .thenReturn(null);
+
+    publisher.publish(Notification.Type.PROGRAM_STATUS, ImmutableMap.of());
+
+    Mockito.verify(messagingService, Mockito.times(2)).publish(storeRequestCaptor.capture());
+    StoreRequest storeRequest = storeRequestCaptor.getValue();
+    Assert.assertEquals("programstatusevent", storeRequest.getTopicId().getTopic());
+  }
+
+  @Test
+  public void testPublishThrowsOnRetryExhausted() throws TopicNotFoundException, IOException {
+    CConfiguration cConf = CConfiguration.create();
+    cConf.setInt(Constants.AppFabric.PROGRAM_STATUS_EVENT_NUM_PARTITIONS, 1);
+    cConf.setInt("system.program.state.retry.policy.max.retries", 3);
+
+    MessagingService messagingService = Mockito.mock(MessagingService.class);
+    MessagingProgramStatePublisher publisher = new MessagingProgramStatePublisher(cConf, messagingService);
+    Mockito.when(messagingService.publish(Mockito.any()))
+        .thenThrow(new SocketTimeoutException());
+
+    RuntimeException outerException = Assert.assertThrows(RuntimeException.class,
+        () -> publisher.publish(Notification.Type.PROGRAM_STATUS, ImmutableMap.of()));
+    Assert.assertNotNull(outerException.getCause());
+    Assert.assertTrue(outerException.getCause() instanceof SocketTimeoutException);
+
+    Mockito.verify(messagingService, Mockito.times(4)).publish(storeRequestCaptor.capture());
+    StoreRequest storeRequest = storeRequestCaptor.getValue();
+    Assert.assertEquals("programstatusevent", storeRequest.getTopicId().getTopic());
+  }
+
+  @Test
+  public void testPublishThrowsForNonRetryableException() throws TopicNotFoundException, IOException {
+    CConfiguration cConf = CConfiguration.create();
+    cConf.setInt(Constants.AppFabric.PROGRAM_STATUS_EVENT_NUM_PARTITIONS, 1);
+
+    MessagingService messagingService = Mockito.mock(MessagingService.class);
+    MessagingProgramStatePublisher publisher = new MessagingProgramStatePublisher(cConf,
+        messagingService);
+    Mockito.when(messagingService.publish(Mockito.any()))
+        .thenThrow(new IOException())
+        .thenReturn(null);
+
+    RuntimeException outerException = Assert.assertThrows(RuntimeException.class,
+        () -> publisher.publish(Notification.Type.PROGRAM_STATUS, ImmutableMap.of()));
+    Assert.assertNotNull(outerException.getCause());
+    Assert.assertTrue(outerException.getCause() instanceof IOException);
+    Mockito.verify(messagingService).publish(storeRequestCaptor.capture());
+    StoreRequest storeRequest = storeRequestCaptor.getValue();
+    Assert.assertEquals("programstatusevent", storeRequest.getTopicId().getTopic());
+  }
 }