From e4ea68a10486cc8cc7b902269c3698761b3bc7ed Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Mon, 16 Oct 2023 06:31:21 -0700
Subject: [PATCH] Register TopN status in plugin's writables (#100874)

```
"node_failures": [
  {
    "type": "failed_node_exception",
    "reason": "Failed node [qpdSPb3yQkuDlsI9TH7a2g]",
    "node_id": "qpdSPb3yQkuDlsI9TH7a2g",
    "caused_by": {
      "type": "transport_serialization_exception",
      "reason": "Failed to deserialize response from handler",
      "caused_by": {
        "type": "illegal_argument_exception",
        "reason": "Unknown NamedWriteable [org.elasticsearch.compute.operator.Operator$Status][topn]"
      }
    }
  }
]
```

I hit this error when trying to retrieve ESQL tasks. The issue is that we forget
 to register NamedWritable for the status of TopN.
---
 .../xpack/esql/plugin/EsqlPlugin.java         |  6 ++--
 .../esql/action/NamedWriteablesTests.java     | 29 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/NamedWriteablesTests.java

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java
index 25802894e2832..fba2f00f0b314 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java
@@ -23,7 +23,7 @@
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.lucene.LuceneSourceOperator;
+import org.elasticsearch.compute.lucene.LuceneOperator;
 import org.elasticsearch.compute.lucene.ValuesSourceReaderOperator;
 import org.elasticsearch.compute.operator.AbstractPageMappingOperator;
 import org.elasticsearch.compute.operator.DriverStatus;
@@ -32,6 +32,7 @@
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
 import org.elasticsearch.compute.operator.exchange.ExchangeSinkOperator;
 import org.elasticsearch.compute.operator.exchange.ExchangeSourceOperator;
+import org.elasticsearch.compute.operator.topn.TopNOperatorStatus;
 import org.elasticsearch.env.Environment;
 import org.elasticsearch.env.NodeEnvironment;
 import org.elasticsearch.indices.IndicesService;
@@ -155,7 +156,8 @@ public List<NamedWriteableRegistry.Entry> getNamedWriteables() {
                 ExchangeSinkOperator.Status.ENTRY,
                 ExchangeSourceOperator.Status.ENTRY,
                 LimitOperator.Status.ENTRY,
-                LuceneSourceOperator.Status.ENTRY,
+                LuceneOperator.Status.ENTRY,
+                TopNOperatorStatus.ENTRY,
                 MvExpandOperator.Status.ENTRY,
                 ValuesSourceReaderOperator.Status.ENTRY,
                 SingleValueQuery.ENTRY
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/NamedWriteablesTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/NamedWriteablesTests.java
new file mode 100644
index 0000000000000..f8a036a262701
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/NamedWriteablesTests.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.action;
+
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.compute.operator.topn.TopNOperatorStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class NamedWriteablesTests extends ESTestCase {
+
+    public void testTopNStatus() throws Exception {
+        try (EsqlPlugin plugin = new EsqlPlugin()) {
+            NamedWriteableRegistry registry = new NamedWriteableRegistry(plugin.getNamedWriteables());
+            TopNOperatorStatus origin = new TopNOperatorStatus(randomNonNegativeInt(), randomNonNegativeLong());
+            TopNOperatorStatus copy = (TopNOperatorStatus) copyNamedWriteable(origin, registry, Operator.Status.class);
+            assertThat(copy.occupiedRows(), equalTo(origin.occupiedRows()));
+            assertThat(copy.ramBytesUsed(), equalTo(origin.ramBytesUsed()));
+        }
+    }
+}