address comments

Signed-off-by: Kaushal Kumar <[email protected]>

Author: kaushalmahi12

server/src/main/java/org/opensearch/wlm/QueryGroupTask.java

@@ -17,7 +17,6 @@
 import org.opensearch.tasks.CancellableTask;
 
 import java.util.Map;
-import java.util.Optional;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -82,12 +81,11 @@ public final String getQueryGroupId() {
      * @param threadContext current threadContext
      */
     public final void setQueryGroupId(final ThreadContext threadContext) {
-        this.queryGroupId = Optional.ofNullable(threadContext)
-            .map(
-                threadContext1 -> Optional.ofNullable(threadContext1.getHeader(QUERY_GROUP_ID_HEADER))
-                    .orElse(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())
-            )
-            .orElse(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get());
+        if (threadContext != null && threadContext.getHeader(QUERY_GROUP_ID_HEADER) != null) {
+            this.queryGroupId = threadContext.getHeader(QUERY_GROUP_ID_HEADER);
+        } else {
+            this.queryGroupId = DEFAULT_QUERY_GROUP_ID_SUPPLIER.get();
+        }
     }
 
     public long getElapsedTime() {

server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java

@@ -8,7 +8,6 @@
 
 package org.opensearch.wlm.stats;
 
-import org.opensearch.Version;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
@@ -124,9 +123,7 @@ public QueryGroupStatsHolder(StreamInput in) throws IOException {
             this.rejections = in.readVLong();
             this.failures = in.readVLong();
             this.totalCancellations = in.readVLong();
-            if (in.getVersion().onOrAfter(Version.V_2_18_0)) {
-                this.shardCompletions = in.readVLong();
-            }
+            this.shardCompletions = in.readVLong();
             this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new);
         }
 
@@ -164,9 +161,7 @@ public static void writeTo(StreamOutput out, QueryGroupStatsHolder statsHolder)
             out.writeVLong(statsHolder.rejections);
             out.writeVLong(statsHolder.failures);
             out.writeVLong(statsHolder.totalCancellations);
-            if (out.getVersion().onOrAfter(Version.V_2_18_0)) {
-                out.writeVLong(statsHolder.shardCompletions);
-            }
+            out.writeVLong(statsHolder.shardCompletions);
             out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo);
         }