From 08e3cbe8c69790648965014ed64d52d0519b40c2 Mon Sep 17 00:00:00 2001
From: Zixuan Liu <nodeces@gmail.com>
Date: Fri, 19 Jan 2024 11:18:02 +0800
Subject: [PATCH] [fix][broker] Fix getMessageById throws 500 (#21919)

Signed-off-by: Zixuan Liu <nodeces@gmail.com>
(cherry picked from commit 6f7b9d9a6147c0de3550a091f5cfba0389d98cec)
---
 .../admin/impl/PersistentTopicsBase.java      |  3 +++
 .../broker/admin/PersistentTopicsTest.java    | 21 ++++++-------------
 .../client/admin/internal/TopicsImpl.java     | 16 +-------------
 3 files changed, 10 insertions(+), 30 deletions(-)

diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/PersistentTopicsBase.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/PersistentTopicsBase.java
index 121b57fc8eea7..758b8048f3d13 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/PersistentTopicsBase.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/PersistentTopicsBase.java
@@ -2829,6 +2829,9 @@ protected CompletableFuture<Response> internalGetMessageById(long ledgerId, long
                         @Override
                         public void readEntryFailed(ManagedLedgerException exception,
                                                     Object ctx) {
+                            if (exception instanceof ManagedLedgerException.LedgerNotExistException) {
+                                throw new RestException(Status.NOT_FOUND, "Message id not found");
+                            }
                             throw new RestException(exception);
                         }
 
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/PersistentTopicsTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/PersistentTopicsTest.java
index 25ad6cab94272..10321cb776101 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/PersistentTopicsTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/PersistentTopicsTest.java
@@ -1358,21 +1358,12 @@ public void testGetMessageById() throws Exception {
         Message<byte[]> message2 = admin.topics().getMessageById(topicName2, id2.getLedgerId(), id2.getEntryId());
         Assert.assertEquals(message2.getData(), data2.getBytes());
 
-        Message<byte[]> message3 = null;
-        try {
-            message3 = admin.topics().getMessageById(topicName2, id1.getLedgerId(), id1.getEntryId());
-            Assert.fail();
-        } catch (Exception e) {
-            Assert.assertNull(message3);
-        }
-
-        Message<byte[]> message4 = null;
-        try {
-            message4 = admin.topics().getMessageById(topicName1, id2.getLedgerId(), id2.getEntryId());
-            Assert.fail();
-        } catch (Exception e) {
-            Assert.assertNull(message4);
-        }
+        Assert.expectThrows(PulsarAdminException.NotFoundException.class, () -> {
+            admin.topics().getMessageById(topicName2, id1.getLedgerId(), id1.getEntryId());
+        });
+        Assert.expectThrows(PulsarAdminException.NotFoundException.class, () -> {
+            admin.topics().getMessageById(topicName1, id2.getLedgerId(), id2.getEntryId());
+        });
     }
 
     @Test
diff --git a/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/TopicsImpl.java b/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/TopicsImpl.java
index e0c64319ea2d9..a89da628372bd 100644
--- a/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/TopicsImpl.java
+++ b/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/TopicsImpl.java
@@ -965,21 +965,7 @@ public CompletableFuture<Void> truncateAsync(String topic) {
 
     @Override
     public CompletableFuture<Message<byte[]>> getMessageByIdAsync(String topic, long ledgerId, long entryId) {
-        CompletableFuture<Message<byte[]>> future = new CompletableFuture<>();
-        getRemoteMessageById(topic, ledgerId, entryId).handle((r, ex) -> {
-            if (ex != null) {
-                if (ex instanceof NotFoundException) {
-                    log.warn("Exception '{}' occurred while trying to get message.", ex.getMessage());
-                    future.complete(r);
-                } else {
-                    future.completeExceptionally(ex);
-                }
-                return null;
-            }
-            future.complete(r);
-            return null;
-        });
-        return future;
+        return getRemoteMessageById(topic, ledgerId, entryId);
     }
 
     private CompletableFuture<Message<byte[]>> getRemoteMessageById(String topic, long ledgerId, long entryId) {