added listeners to exit connector on user termination

Author: SravanThotakura05

src/main/java/com/solacecoe/connectors/spark/streaming/partitions/SolaceInputPartitionReader.java

@@ -25,6 +25,7 @@
 import org.apache.spark.sql.execution.streaming.MicroBatchExecution;
 import org.apache.spark.sql.execution.streaming.StreamExecution;
 import org.apache.spark.unsafe.types.UTF8String;
+import org.apache.spark.util.TaskFailureListener;
 
 import java.io.BufferedWriter;
 import java.io.File;
@@ -119,6 +120,13 @@ public boolean next() {
             this.solaceBroker.close();
             throw new SolaceSessionException(this.solaceBroker.getException());
         }
+
+        if (TaskContext.get() != null && TaskContext.get().isInterrupted()) {
+            log.info("SolaceSparkConnector - Interrupted while waiting for next message");
+            SolaceConnectionManager.close(this.solaceInputPartition.getId());
+            SolaceMessageTracker.resetId(uniqueId);
+            throw new RuntimeException("Task was interrupted.");
+        }
         solaceMessage = getNextMessage();
         return solaceMessage != null;
     }
@@ -225,16 +233,24 @@ public void close() {
         }
     }
 
+    private void logShutdownMessage(TaskContext context) {
+        log.info("SolaceSparkConnector - Closing connections to Solace as task {} is interrupted or failed", String.join(",", context.getLocalProperty(StreamExecution.QUERY_ID_KEY()),
+                context.getLocalProperty(MicroBatchExecution.BATCH_ID_KEY()),
+                Integer.toString(context.stageId()),
+                Integer.toString(context.partitionId())));
+        SolaceConnectionManager.close(this.solaceInputPartition.getId());
+        SolaceMessageTracker.resetId(uniqueId);
+    }
+
     private void registerTaskListener() {
+        this.taskContext.addTaskFailureListener((context, error) -> {
+            log.error("SolaceSparkConnector - Input Partition {} failed with error", this.solaceInputPartition.getId(), error);
+            logShutdownMessage(context);
+        });
         this.taskContext.addTaskCompletionListener(context -> {
             log.info("SolaceSparkConnector - Task {} state is completed :: {}, failed :: {}, interrupted :: {}", uniqueId, context.isCompleted(), context.isFailed(), context.isInterrupted());
             if(context.isInterrupted() || context.isFailed()) {
-                log.info("SolaceSparkConnector - Closing connections to Solace as task {} is interrupted or failed", String.join(",", context.getLocalProperty(StreamExecution.QUERY_ID_KEY()),
-                        context.getLocalProperty(MicroBatchExecution.BATCH_ID_KEY()),
-                        Integer.toString(context.stageId()),
-                        Integer.toString(context.partitionId())));
-                SolaceConnectionManager.close(this.solaceInputPartition.getId());
-                SolaceMessageTracker.resetId(uniqueId);
+                logShutdownMessage(context);
             } else if (context.isCompleted()) {
                 List<String> ids = SolaceMessageTracker.getIds();
                 try {

src/main/java/com/solacecoe/connectors/spark/streaming/solace/EventListener.java

@@ -76,7 +76,12 @@ public void onReceive(BytesXMLMessage msg) {
                         SolaceSparkPartitionCheckpoint solaceSparkPartitionCheckpoint = this.checkpoints.stream().filter(checkpoint -> checkpoint.getPartitionId().equals(partitionKey)).findFirst().orElse(null);
                         if(solaceSparkPartitionCheckpoint != null) {
                             lastKnownMessageIDs = Arrays.stream(solaceSparkPartitionCheckpoint.getMessageIDs().split(",")).collect(Collectors.toList());
-                            compareMessageIds(lastKnownMessageIDs, messageID, msg);
+                            if(lastKnownMessageIDs.isEmpty()) {
+                                log.warn("SolaceSparkConnector - No checkpoint found. Message will be reprocessed to ensure reliability—any duplicates must be handled by the downstream system.");
+                                this.messages.add(new SolaceMessage(msg));
+                            } else {
+                                compareMessageIds(lastKnownMessageIDs, messageID, msg);
+                            }
                         } else {
                             log.warn("SolaceSparkConnector - No checkpoint found for partition key {}. Message will be reprocessed to ensure reliability—any duplicates must be handled by the downstream system.", partitionKey);
                             this.messages.add(new SolaceMessage(msg));