Async search context and and state management framework

src/main/java/com/amazon/opendistroforelasticsearch/search/async/AsyncSearchId.java

@@ -16,12 +16,6 @@
 package com.amazon.opendistroforelasticsearch.search.async;
 
 import com.amazon.opendistroforelasticsearch.search.async.context.AsyncSearchContextId;
-import org.apache.lucene.store.ByteArrayDataInput;
-import org.elasticsearch.common.bytes.BytesReference;
-import org.elasticsearch.common.io.stream.BytesStreamOutput;
-
-import java.io.IOException;
-import java.util.Base64;
 import java.util.Objects;
 
 /**
@@ -62,47 +56,6 @@ public String toString() {
         return "[" + node + "][" + taskId + "][" + asyncSearchContextId + "]";
     }
 
-    /**
-     * Encodes the {@linkplain AsyncSearchId} in base64 encoding and returns an identifier for the submitted async search
-     *
-     * @param asyncSearchId The object to be encoded
-     * @return The id which is used to access the submitted async search
-     */
-    public static String buildAsyncId(AsyncSearchId asyncSearchId) {
-        try (BytesStreamOutput out = new BytesStreamOutput()) {
-            out.writeString(asyncSearchId.getNode());
-            out.writeLong(asyncSearchId.getTaskId());
-            out.writeString(asyncSearchId.getAsyncSearchContextId().getContextId());
-            out.writeLong(asyncSearchId.getAsyncSearchContextId().getId());
-            return Base64.getUrlEncoder().withoutPadding().encodeToString(BytesReference.toBytes(out.bytes()));
-        } catch (IOException e) {
-            throw new IllegalArgumentException("Cannot build async search id", e);
-        }
-    }
-
-    /**
-     * Attempts to decode a base64 encoded string into an {@linkplain AsyncSearchId} which contains the details pertaining to
-     * an async search being accessed.
-     *
-     * @param asyncSearchId The string to be decoded
-     * @return The parsed {@linkplain AsyncSearchId}
-     */
-    public static AsyncSearchId parseAsyncId(String asyncSearchId) {
-        try {
-            byte[] bytes = Base64.getUrlDecoder().decode(asyncSearchId);
-            ByteArrayDataInput in = new ByteArrayDataInput(bytes);
-            String node = in.readString();
-            long taskId = in.readLong();
-            String contextId = in.readString();
-            long id = in.readLong();
-            if (in.getPosition() != bytes.length) {
-                throw new IllegalArgumentException("Not all bytes were read");
-            }
-            return new AsyncSearchId(node, taskId, new AsyncSearchContextId(contextId, id));
-        } catch (Exception e) {
-            throw new IllegalArgumentException("Cannot parse async search id", e);
-        }
-    }
 
     @Override
     public int hashCode() {

src/main/java/com/amazon/opendistroforelasticsearch/search/async/AsyncSearchIdConverter.java

@@ -0,0 +1,68 @@
+/*
+ *   Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.search.async;
+
+import com.amazon.opendistroforelasticsearch.search.async.context.AsyncSearchContextId;
+import org.apache.lucene.store.ByteArrayDataInput;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+
+import java.io.IOException;
+import java.util.Base64;
+
+public class AsyncSearchIdConverter {
+    /**
+     * Encodes the {@linkplain AsyncSearchId} in base64 encoding and returns an identifier for the submitted async search
+     *
+     * @param asyncSearchId The object to be encoded
+     * @return The id which is used to access the submitted async search
+     */
+    public static String buildAsyncId(AsyncSearchId asyncSearchId) {
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            out.writeString(asyncSearchId.getNode());
+            out.writeLong(asyncSearchId.getTaskId());
+            out.writeString(asyncSearchId.getAsyncSearchContextId().getContextId());
+            out.writeLong(asyncSearchId.getAsyncSearchContextId().getId());
+            return Base64.getUrlEncoder().withoutPadding().encodeToString(BytesReference.toBytes(out.bytes()));
+        } catch (IOException e) {
+            throw new IllegalArgumentException("Cannot build async search id", e);
+        }
+    }
+
+    /**
+     * Attempts to decode a base64 encoded string into an {@linkplain AsyncSearchId} which contains the details pertaining to
+     * an async search being accessed.
+     *
+     * @param asyncSearchId The string to be decoded
+     * @return The parsed {@linkplain AsyncSearchId}
+     */
+    public static AsyncSearchId parseAsyncId(String asyncSearchId) {
+        try {
+            byte[] bytes = Base64.getUrlDecoder().decode(asyncSearchId);
+            ByteArrayDataInput in = new ByteArrayDataInput(bytes);
+            String node = in.readString();
+            long taskId = in.readLong();
+            String contextId = in.readString();
+            long id = in.readLong();
+            if (in.getPosition() != bytes.length) {
+                throw new IllegalArgumentException("Not all bytes were read");
+            }
+            return new AsyncSearchId(node, taskId, new AsyncSearchContextId(contextId, id));
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Cannot parse async search id", e);
+        }
+    }
+}

src/main/java/com/amazon/opendistroforelasticsearch/search/async/context/AsyncSearchContext.java

@@ -57,12 +57,12 @@ AsyncSearchProgressListener getAsyncSearchProgressListener() {
     @Nullable
     public AsyncSearchContextListener getContextListener() {return asyncSearchContextListener; }
 
-    public AsyncSearchState getAsyncSearchStage() {
+    public AsyncSearchState getAsyncSearchState() {
         return currentStage;
     }
 
     public boolean isRunning() {
-        return getAsyncSearchStage() == AsyncSearchState.RUNNING;
+        return getAsyncSearchState() == AsyncSearchState.RUNNING;
     }
 
     public AsyncSearchContextId getContextId() {

src/main/java/com/amazon/opendistroforelasticsearch/search/async/context/active/AsyncSearchActiveContext.java

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.search.async.context.active;
 
 import com.amazon.opendistroforelasticsearch.search.async.AsyncSearchId;
+import com.amazon.opendistroforelasticsearch.search.async.AsyncSearchIdConverter;
 import com.amazon.opendistroforelasticsearch.search.async.context.AsyncSearchContext;
 import com.amazon.opendistroforelasticsearch.search.async.context.AsyncSearchContextId;
 import com.amazon.opendistroforelasticsearch.search.async.context.permits.AsyncSearchContextPermits;
@@ -32,6 +33,7 @@
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.threadpool.ThreadPool;
 
+import java.io.Closeable;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.LongSupplier;
@@ -43,7 +45,7 @@
  * The context representing an ongoing search, keeps track of the underlying {@link SearchTask}
  * and {@link SearchProgressActionListener}
  */
-public class AsyncSearchActiveContext extends AsyncSearchContext {
+public class AsyncSearchActiveContext extends AsyncSearchContext implements Closeable {
 
     private static final Logger logger = LogManager.getLogger(AsyncSearchContext.class);
 
@@ -85,7 +87,7 @@ public void setTask(SearchTask searchTask) {
         this.searchTask.set(searchTask);
         this.startTimeMillis = searchTask.getStartTime();
         this.expirationTimeMillis = startTimeMillis + keepAlive.getMillis();
-        this.asyncSearchId.set(AsyncSearchId.buildAsyncId(new AsyncSearchId(nodeId, searchTask.getId(), getContextId())));
+        this.asyncSearchId.set(AsyncSearchIdConverter.buildAsyncId(new AsyncSearchId(nodeId, searchTask.getId(), getContextId())));
     }
 
     public void processSearchFailure(Exception e) {
@@ -147,4 +149,9 @@ public void acquireContextPermit(final ActionListener<Releasable> onPermitAcquir
     public void acquireAllContextPermits(final ActionListener<Releasable> onPermitAcquired, TimeValue timeout, String reason) {
         asyncSearchContextPermits.asyncAcquireAllPermits(onPermitAcquired, timeout, reason);
     }
+
+    @Override
+    public void close() {
+        asyncSearchContextPermits.close();
+    }
 }

src/main/java/com/amazon/opendistroforelasticsearch/search/async/context/permits/AsyncSearchContextPermits.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.common.util.concurrent.RunOnce;
 import org.elasticsearch.threadpool.ThreadPool;
 
+import java.io.Closeable;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -38,15 +39,16 @@
  * before it transitions context to the index. Provides fairness to consumers and throws {@linkplain TimeoutException} after
  * maximum time has elapsed waiting for the in-flight operations block.
  */
-public class AsyncSearchContextPermits {
+public class AsyncSearchContextPermits implements Closeable {
 
     private static final int TOTAL_PERMITS = Integer.MAX_VALUE;
 
-    //visible for testing
     final Semaphore semaphore;
     private final AsyncSearchContextId asyncSearchContextId;
     private volatile String lockDetails;
     private final ThreadPool threadPool;
+    private volatile boolean closed;
+
     private static final Logger logger = LogManager.getLogger(AsyncSearchContextPermits.class);
 
     public AsyncSearchContextPermits(AsyncSearchContextId asyncSearchContextId, ThreadPool threadPool) {
@@ -56,24 +58,27 @@ public AsyncSearchContextPermits(AsyncSearchContextId asyncSearchContextId, Thre
     }
 
     private Releasable acquirePermits(int permits, TimeValue timeout, final String details) throws TimeoutException {
+        RunOnce release = new RunOnce(() -> {});
+        if (closed) {
+            throw new IllegalStateException("trying to acquire permits on closed context ["+ asyncSearchContextId +"]");
+        }
         try {
             if (semaphore.tryAcquire(permits, timeout.getMillis(), TimeUnit.MILLISECONDS)) {
                 this.lockDetails = details;
-                final RunOnce release = new RunOnce(() -> semaphore.release(permits));
+                release = new RunOnce(() -> semaphore.release(permits));
                 return release::run;
             } else {
-                throw new TimeoutException(
-                        "obtaining context lock" + asyncSearchContextId + "timed out after " + timeout.getMillis() + "ms, " +
-                                "previous lock details: [" + lockDetails + "] trying to lock for [" + details + "]");
+                throw new TimeoutException("obtaining context lock" + asyncSearchContextId + "timed out after " +
+                        timeout.getMillis() + "ms, previous lock details: [" + lockDetails + "] trying to lock for [" + details + "]");
             }
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
+            release.run();
             throw new RuntimeException("thread interrupted while trying to obtain context lock", e);
         }
     }
 
-    private void asyncAcquirePermit(
-            int permits, final ActionListener<Releasable> onAcquired, final TimeValue timeout, String reason) {
+    private void asyncAcquirePermit(int permits, final ActionListener<Releasable> onAcquired, final TimeValue timeout, String reason) {
         threadPool.executor(AsyncSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME).execute(new AbstractRunnable() {
             @Override
             public void onFailure(final Exception e) {
@@ -118,4 +123,9 @@ public void asyncAcquirePermit(final ActionListener<Releasable> onAcquired, fina
     public void asyncAcquireAllPermits(final ActionListener<Releasable> onAcquired, final TimeValue timeout, String reason) {
         asyncAcquirePermit(TOTAL_PERMITS, onAcquired, timeout, reason);
     }
+
+    @Override
+    public void close() {
+        closed = true;
+    }
 }

src/main/java/com/amazon/opendistroforelasticsearch/search/async/context/persistence/AsyncSearchPersistenceContext.java

@@ -112,7 +112,7 @@ public Exception getSearchError() {
     }
 
     @Override
-    public AsyncSearchState getAsyncSearchStage() {
+    public AsyncSearchState getAsyncSearchState() {
         return AsyncSearchState.PERSISTED;
     }