1 files changed, 257 insertions, 0 deletions
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ReferenceCountingRequestHandler.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ReferenceCountingRequestHandler.java
new file mode 100644
index 00000000000..f2bf5b56d5c
--- /dev/null
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ReferenceCountingRequestHandler.java
@@ -0,0 +1,257 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.jdisc.http.server.jetty;
+
+import com.yahoo.jdisc.Request;
+import com.yahoo.jdisc.ResourceReference;
+import com.yahoo.jdisc.Response;
+import com.yahoo.jdisc.SharedResource;
+import com.yahoo.jdisc.handler.CompletionHandler;
+import com.yahoo.jdisc.handler.ContentChannel;
+import com.yahoo.jdisc.handler.NullContent;
+import com.yahoo.jdisc.handler.RequestHandler;
+import com.yahoo.jdisc.handler.ResponseHandler;
+
+import java.nio.ByteBuffer;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * This class wraps a request handler and does reference counting on the request for every object that depends on the
+ * request, such as the response handler, content channels and completion handlers. This ensures that requests (and
+ * hence the current container) will be referenced until the end of the request handling - even with async handling in
+ * non-framework threads - without requiring the application to handle this tedious work.
+ *
+ * @author bakksjo
+ */
+@SuppressWarnings("try")
+class ReferenceCountingRequestHandler implements RequestHandler {
+
+    private static final Logger log = Logger.getLogger(ReferenceCountingRequestHandler.class.getName());
+
+    final RequestHandler delegate;
+
+    ReferenceCountingRequestHandler(RequestHandler delegate) {
+        Objects.requireNonNull(delegate, "delegate");
+        this.delegate = delegate;
+    }
+
+    @Override
+    public ContentChannel handleRequest(Request request, ResponseHandler responseHandler) {
+        try (final ResourceReference requestReference = request.refer()) {
+            ContentChannel contentChannel;
+            final ReferenceCountingResponseHandler referenceCountingResponseHandler
+                    = new ReferenceCountingResponseHandler(request, new NullContentResponseHandler(responseHandler));
+            try {
+                contentChannel = delegate.handleRequest(request, referenceCountingResponseHandler);
+                Objects.requireNonNull(contentChannel, "contentChannel");
+            } catch (Throwable t) {
+                try {
+                    // The response handler might never be invoked, due to the exception thrown from handleRequest().
+                    referenceCountingResponseHandler.unrefer();
+                } catch (Throwable thrownFromUnrefer) {
+                    log.log(Level.WARNING, "Unexpected problem", thrownFromUnrefer);
+                }
+                throw t;
+            }
+            return new ReferenceCountingContentChannel(request, contentChannel);
+        }
+    }
+
+    @Override
+    public void handleTimeout(Request request, ResponseHandler responseHandler) {
+        delegate.handleTimeout(request, new NullContentResponseHandler(responseHandler));
+    }
+
+    @Override
+    public ResourceReference refer() {
+        return delegate.refer();
+    }
+
+    @Override
+    public void release() {
+        delegate.release();
+    }
+
+    @Override
+    public String toString() {
+        return delegate.toString();
+    }
+
+    private static class ReferenceCountingResponseHandler implements ResponseHandler {
+
+        final SharedResource request;
+        final ResourceReference requestReference;
+        final ResponseHandler delegate;
+        final AtomicBoolean closed = new AtomicBoolean(false);
+
+        ReferenceCountingResponseHandler(SharedResource request, ResponseHandler delegate) {
+            Objects.requireNonNull(request, "request");
+            Objects.requireNonNull(delegate, "delegate");
+            this.request = request;
+            this.delegate = delegate;
+            this.requestReference = request.refer();
+        }
+
+        @Override
+        public ContentChannel handleResponse(Response response) {
+            if (closed.getAndSet(true)) {
+                throw new IllegalStateException(delegate + " is already called.");
+            }
+            try (final ResourceReference ref = requestReference) {
+                ContentChannel contentChannel = delegate.handleResponse(response);
+                Objects.requireNonNull(contentChannel, "contentChannel");
+                return new ReferenceCountingContentChannel(request, contentChannel);
+            }
+        }
+
+        @Override
+        public String toString() {
+            return delegate.toString();
+        }
+
+        /**
+         * Close the reference that is normally closed by {@link #handleResponse(Response)}.
+         *
+         * This is to be used in error situations, where handleResponse() may not be invoked.
+         */
+        public void unrefer() {
+            if (closed.getAndSet(true)) {
+                // This simply means that handleResponse() has been run, in which case we are
+                // guaranteed that the reference is closed.
+                return;
+            }
+            requestReference.close();
+        }
+    }
+
+    private static class ReferenceCountingContentChannel implements ContentChannel {
+
+        final SharedResource request;
+        final ResourceReference requestReference;
+        final ContentChannel delegate;
+
+        ReferenceCountingContentChannel(SharedResource request, ContentChannel delegate) {
+            Objects.requireNonNull(request, "request");
+            Objects.requireNonNull(delegate, "delegate");
+            this.request = request;
+            this.delegate = delegate;
+            this.requestReference = request.refer();
+        }
+
+        @Override
+        public void write(ByteBuffer buf, CompletionHandler completionHandler) {
+            final CompletionHandler referenceCountingCompletionHandler
+                    = new ReferenceCountingCompletionHandler(request, completionHandler);
+            try {
+                delegate.write(buf, referenceCountingCompletionHandler);
+            } catch (Throwable t) {
+                try {
+                    referenceCountingCompletionHandler.failed(t);
+                } catch (AlreadyCompletedException ignored) {
+                } catch (Throwable failFailure) {
+                    log.log(Level.WARNING, "Failure during call to CompletionHandler.failed()", failFailure);
+                }
+                throw t;
+            }
+        }
+
+        @Override
+        public void close(CompletionHandler completionHandler) {
+            final CompletionHandler referenceCountingCompletionHandler
+                    = new ReferenceCountingCompletionHandler(request, completionHandler);
+            try (final ResourceReference ref = requestReference) {
+                delegate.close(referenceCountingCompletionHandler);
+            } catch (Throwable t) {
+                try {
+                    referenceCountingCompletionHandler.failed(t);
+                } catch (AlreadyCompletedException ignored) {
+                } catch (Throwable failFailure) {
+                    log.log(Level.WARNING, "Failure during call to CompletionHandler.failed()", failFailure);
+                }
+                throw t;
+            }
+        }
+
+        @Override
+        public String toString() {
+            return delegate.toString();
+        }
+    }
+
+    private static class AlreadyCompletedException extends IllegalStateException {
+        public AlreadyCompletedException(final CompletionHandler completionHandler) {
+            super(completionHandler + " is already called.");
+        }
+    }
+
+    private static class ReferenceCountingCompletionHandler implements CompletionHandler {
+
+        final ResourceReference requestReference;
+        final CompletionHandler delegate;
+        final AtomicBoolean closed = new AtomicBoolean(false);
+
+        public ReferenceCountingCompletionHandler(SharedResource request, CompletionHandler delegate) {
+            this.delegate = delegate;
+            this.requestReference = request.refer();
+        }
+
+        @Override
+        public void completed() {
+            if (closed.getAndSet(true)) {
+                throw new AlreadyCompletedException(delegate);
+            }
+            try {
+                if (delegate != null) {
+                    delegate.completed();
+                }
+            } finally {
+                requestReference.close();
+            }
+        }
+
+        @Override
+        public void failed(Throwable t) {
+            if (closed.getAndSet(true)) {
+                throw new AlreadyCompletedException(delegate);
+            }
+            try (final ResourceReference ref = requestReference) {
+                if (delegate != null) {
+                    delegate.failed(t);
+                } else {
+                    log.log(Level.WARNING, "Uncaught completion failure.", t);
+                }
+            }
+        }
+
+        @Override
+        public String toString() {
+            return String.valueOf(delegate);
+        }
+    }
+
+    private static class NullContentResponseHandler implements ResponseHandler {
+
+        final ResponseHandler delegate;
+
+        NullContentResponseHandler(ResponseHandler delegate) {
+            Objects.requireNonNull(delegate, "delegate");
+            this.delegate = delegate;
+        }
+
+        @Override
+        public ContentChannel handleResponse(Response response) {
+            ContentChannel contentChannel = delegate.handleResponse(response);
+            if (contentChannel == null) {
+                contentChannel = NullContent.INSTANCE;
+            }
+            return contentChannel;
+        }
+
+        @Override
+        public String toString() {
+            return delegate.toString();
+        }
+    }
+}