NSG Timeout basic implementation (lacks UI)

src/community/nsg-profile/pom.xml

@@ -74,7 +74,7 @@
     <dependency>
       <groupId>javax.servlet</groupId>
       <artifactId>javax.servlet-api</artifactId>
-      <scope>test</scope>
+      <scope>provided</scope>
     </dependency>
     <dependency>
       <groupId>org.geoserver.web</groupId>

src/community/nsg-profile/src/main/java/org/geoserver/nsg/timeout/TimeoutCallback.java

@@ -0,0 +1,143 @@
+/* (c) 2017 Open Source Geospatial Foundation - all rights reserved
+ * This code is licensed under the GPL 2.0 license, available at the root
+ * application directory.
+ */
+package org.geoserver.nsg.timeout;
+
+import net.opengis.wfs20.BaseRequestType;
+import org.geoserver.config.GeoServer;
+import org.geoserver.ows.AbstractDispatcherCallback;
+import org.geoserver.ows.Request;
+import org.geoserver.platform.Operation;
+import org.geoserver.wfs.WFSException;
+import org.geoserver.wfs.WFSInfo;
+import org.geoserver.wfs.request.FeatureCollectionResponse;
+import org.geotools.feature.FeatureCollection;
+import org.geotools.util.Converters;
+import org.geotools.util.Version;
+import org.geotools.util.logging.Logging;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import java.util.stream.Collectors;
+
+/**
+ * Callback implementing NSG timeout extension
+ */
+public class TimeoutCallback extends AbstractDispatcherCallback {
+
+    static final Logger LOGGER = Logging.getLogger(TimeoutCallback.class);
+
+    /**
+     * The key storing the NSG request timeout in the {@link WFSInfo#getMetadata()} map
+     */
+    static final String TIMEOUT_CONFIG_KEY = "org.geoserver.nsg.timeout";
+
+    /**
+     * The default timeout according to specification (5 minutes)
+     */
+    static final int TIMEOUT_CONFIG_DEFAULT = 300;
+
+    static final String TIMEOUT_REQUEST_ATTRIBUTE = "timeout";
+
+    static final Version V_20 = new Version("2.0");
+
+    GeoServer gs;
+
+    ThreadLocal<TimeoutVerifier> TIMEOUT_VERIFIER = new ThreadLocal<>();
+
+    public TimeoutCallback(GeoServer gs) {
+        this.gs = gs;
+    }
+
+    @Override
+    public Request init(Request request) {
+        return super.init(request);
+    }
+
+    @Override
+    public Operation operationDispatched(Request request, Operation operation) {
+        String version = request.getVersion();
+        String method = request.getRequest();
+        long timeout = getTimeoutMilliseconds(operation);
+        if ("WFS".equalsIgnoreCase(request.getService()) &&
+                (version == null || V_20.compareTo(new Version(version)) <= 0)
+                && method != null && (method.equalsIgnoreCase("GetFeature") ||
+                method.equalsIgnoreCase("GetFeatureWithLock") ||
+                method.equalsIgnoreCase("GetPropertyValue")) &&
+                timeout > 0 &&
+                operation.getParameters().length > 0 &&
+                operation.getParameters()[0] instanceof BaseRequestType
+                ) {
+
+            if (LOGGER.isLoggable(Level.FINE)) {
+                LOGGER.fine("Starting to track NSG timeout on this request");
+            }
+
+            // start tracking time
+            TimeoutVerifier timeoutVerifier = new TimeoutVerifier((BaseRequestType) operation.getParameters()[0],
+                    timeout);
+            // need to wrap the http response and its output stream
+            request.setHttpResponse(new TimeoutCancellingResponse(request.getHttpResponse(), timeoutVerifier));
+            // set in the thread local for later use
+            TIMEOUT_VERIFIER.set(timeoutVerifier);
+        }
+
+        return operation;
+    }
+
+
+    @Override
+    public void finished(Request request) {
+        TIMEOUT_VERIFIER.remove();
+    }
+
+    @Override
+    public Object operationExecuted(Request request, Operation operation, Object result) {
+        TimeoutVerifier timeoutVerifier = TIMEOUT_VERIFIER.get();
+        if (timeoutVerifier != null) {
+            // check before encode
+            timeoutVerifier.checkTimeout();
+
+            // wrap if needed
+            if (result instanceof FeatureCollectionResponse) {
+                FeatureCollectionResponse featureCollectionResponse = (FeatureCollectionResponse) result;
+                List<FeatureCollection> collections = featureCollectionResponse.getFeatures();
+                List<FeatureCollection> wrappers = collections.stream().map(fc -> TimeoutFeatureCollection.wrap
+                        (timeoutVerifier,
+                                fc)).collect(Collectors
+                        .toList());
+
+                featureCollectionResponse.setFeatures(wrappers);
+            }
+        }
+
+        return result;
+    }
+
+
+    private long getTimeoutMilliseconds(Operation operation) {
+        // check if there is a timeout parameter
+        Object[] parameters = operation.getParameters();
+        if (parameters != null && parameters.length > 0 && parameters[0] instanceof BaseRequestType) {
+            BaseRequestType request = (BaseRequestType) parameters[0];
+            Object timeout = request.getExtendedProperties().get(TIMEOUT_REQUEST_ATTRIBUTE);
+            if (timeout != null) {
+                Long converted = Converters.convert(timeout, Long.class);
+                if (converted != null && converted > 0) {
+                    return converted * 1000l;
+                } else {
+                    throw new WFSException(request, "Invalid timeout value: " + timeout);
+                }
+            }
+        }
+
+        // use the configured default
+        WFSInfo wfs = gs.getService(WFSInfo.class);
+        Integer timeoutSeconds = wfs.getMetadata().get(TIMEOUT_CONFIG_KEY, Integer.class);
+        return Optional.ofNullable(timeoutSeconds).orElse(TIMEOUT_CONFIG_DEFAULT) * 1000L;
+    }
+
+}

src/community/nsg-profile/src/main/java/org/geoserver/nsg/timeout/TimeoutCancellingResponse.java

@@ -0,0 +1,31 @@
+/* (c) 2017 Open Source Geospatial Foundation - all rights reserved
+ * This code is licensed under the GPL 2.0 license, available at the root
+ * application directory.
+ */
+package org.geoserver.nsg.timeout;
+
+import javax.servlet.ServletOutputStream;
+import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpServletResponseWrapper;
+import java.io.IOException;
+import java.io.PrintWriter;
+
+class TimeoutCancellingResponse extends HttpServletResponseWrapper {
+    TimeoutVerifier timeoutVerifier;
+
+    public TimeoutCancellingResponse(HttpServletResponse response, TimeoutVerifier timeoutVerifier) {
+        super(response);
+        this.timeoutVerifier = timeoutVerifier;
+    }
+
+    @Override
+    public ServletOutputStream getOutputStream() throws IOException {
+        ServletOutputStream delegate = super.getOutputStream();
+        return new TimeoutCancellingStream(timeoutVerifier, delegate);
+    }
+
+    @Override
+    public PrintWriter getWriter() throws IOException {
+        return new PrintWriter(getOutputStream());
+    }
+}

src/community/nsg-profile/src/main/java/org/geoserver/nsg/timeout/TimeoutCancellingStream.java

@@ -0,0 +1,44 @@
+/* (c) 2017 Open Source Geospatial Foundation - all rights reserved
+ * This code is licensed under the GPL 2.0 license, available at the root
+ * application directory.
+ */
+package org.geoserver.nsg.timeout;
+
+import org.geotools.util.logging.Logging;
+
+import javax.servlet.ServletOutputStream;
+import java.io.IOException;
+import java.util.logging.Logger;
+
+/**
+ * A simple stream that will clear the timeout verifier when the first byte
+ * is written to the output
+ */
+public class TimeoutCancellingStream extends ServletOutputStream {
+    TimeoutVerifier timeoutVerifier;
+    ServletOutputStream delegate;
+
+    public TimeoutCancellingStream(TimeoutVerifier timeoutVerifier, ServletOutputStream delegate) {
+        this.timeoutVerifier = timeoutVerifier;
+        this.delegate = delegate;
+    }
+
+
+    @Override
+    public void write(int b) throws IOException {
+        this.timeoutVerifier.cancel();
+        delegate.write(b);
+    }
+
+    @Override
+    public void write(byte[] b) throws IOException {
+        this.timeoutVerifier.cancel();
+        super.write(b);
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) throws IOException {
+        this.timeoutVerifier.cancel();
+        super.write(b, off, len);
+    }
+}

src/community/nsg-profile/src/main/java/org/geoserver/nsg/timeout/TimeoutFeatureCollection.java

@@ -0,0 +1,168 @@
+/* (c) 2017 Open Source Geospatial Foundation - all rights reserved
+ * This code is licensed under the GPL 2.0 license, available at the root
+ * application directory.
+ */
+package org.geoserver.nsg.timeout;
+
+import org.geotools.data.simple.SimpleFeatureCollection;
+import org.geotools.data.simple.SimpleFeatureIterator;
+import org.geotools.feature.FeatureCollection;
+import org.geotools.feature.FeatureIterator;
+import org.geotools.geometry.jts.ReferencedEnvelope;
+import org.opengis.feature.Feature;
+import org.opengis.feature.FeatureVisitor;
+import org.opengis.feature.simple.SimpleFeature;
+import org.opengis.feature.simple.SimpleFeatureType;
+import org.opengis.feature.type.FeatureType;
+import org.opengis.filter.Filter;
+import org.opengis.filter.sort.SortBy;
+import org.opengis.util.ProgressListener;
+
+import java.io.IOException;
+import java.util.Collection;
+
+/**
+ * A {@link FeatureCollection} decorator that checks if the timeout expired, and throws an exception
+ * in case it is
+ * @param <T>
+ * @param <F>
+ */
+class TimeoutFeatureCollection<T extends FeatureType, F extends Feature> implements FeatureCollection<T, F> {
+
+
+    /**
+     * Wraps a feature collection into a timing out decorator, keeping its {@link SimpleFeature} nature
+     * if possible
+     *
+     * @param <R>
+     * @param timeoutVerifier
+     * @param fc
+     * @return
+     */
+    public static FeatureCollection wrap(TimeoutVerifier timeoutVerifier, FeatureCollection fc) {
+        if (fc instanceof SimpleFeatureCollection) {
+            return new SimpleTimeoutCollection(timeoutVerifier, fc);
+        } else {
+            return new TimeoutFeatureCollection<>(timeoutVerifier, fc);
+        }
+    }
+
+    /**
+     * Simple feature version of {@link TimeoutFeatureCollection}
+     */
+    static class SimpleTimeoutCollection extends TimeoutFeatureCollection<SimpleFeatureType, SimpleFeature> implements SimpleFeatureCollection {
+
+
+        public SimpleTimeoutCollection(TimeoutVerifier timeoutVerifier, FeatureCollection<SimpleFeatureType,
+                SimpleFeature> delegate) {
+            super(timeoutVerifier, delegate);
+        }
+
+        @Override
+        public SimpleFeatureIterator features() {
+            return new TimeoutFeatureIterator.SimpleTimeoutFeatureIterator(timeoutVerifier, super.features());
+        }
+
+        @Override
+        public SimpleFeatureCollection subCollection(Filter filter) {
+            timeoutVerifier.checkTimeout();
+            return new SimpleTimeoutCollection(timeoutVerifier, delegate.subCollection(filter));
+        }
+
+        @Override
+        public SimpleFeatureCollection sort(SortBy order) {
+            return new SimpleTimeoutCollection(timeoutVerifier, delegate.sort(order));
+        }
+    }
+
+    TimeoutVerifier timeoutVerifier;
+    FeatureCollection<T, F> delegate;
+
+    public TimeoutFeatureCollection(TimeoutVerifier timeoutVerifier, FeatureCollection<T, F> delegate) {
+        this.timeoutVerifier = timeoutVerifier;
+        this.delegate = delegate;
+    }
+
+    // Timeout delegate creating methods
+
+    @Override
+    public FeatureIterator<F> features() {
+        timeoutVerifier.checkTimeout();
+        return new TimeoutFeatureIterator<>(timeoutVerifier, delegate.features());
+    }
+
+    @Override
+    public void accepts(FeatureVisitor visitor, ProgressListener progress) throws IOException {
+        timeoutVerifier.checkTimeout();
+        TimeoutFeatureVisitor timeoutVisitor = new TimeoutFeatureVisitor(timeoutVerifier, visitor);
+        delegate.accepts(timeoutVisitor, progress);
+    }
+
+    @Override
+    public FeatureCollection<T, F> subCollection(Filter filter) {
+        timeoutVerifier.checkTimeout();
+        return new TimeoutFeatureCollection<>(timeoutVerifier, delegate.subCollection(filter));
+    }
+
+    @Override
+    public FeatureCollection<T, F> sort(SortBy order) {
+        timeoutVerifier.checkTimeout();
+        return new TimeoutFeatureCollection<>(timeoutVerifier, delegate.sort(order));
+    }
+
+    // Simple check and delegate methods
+
+    @Override
+    public T getSchema() {
+        timeoutVerifier.checkTimeout();
+        return delegate.getSchema();
+    }
+
+    @Override
+    public String getID() {
+        timeoutVerifier.checkTimeout();
+        return delegate.getID();
+    }
+
+    @Override
+    public ReferencedEnvelope getBounds() {
+        timeoutVerifier.checkTimeout();
+        return delegate.getBounds();
+    }
+
+    @Override
+    public boolean contains(Object o) {
+        timeoutVerifier.checkTimeout();
+        return delegate.contains(o);
+    }
+
+    @Override
+    public boolean containsAll(Collection<?> o) {
+        timeoutVerifier.checkTimeout();
+        return delegate.containsAll(o);
+    }
+
+    @Override
+    public boolean isEmpty() {
+        timeoutVerifier.checkTimeout();
+        return delegate.isEmpty();
+    }
+
+    @Override
+    public int size() {
+        timeoutVerifier.checkTimeout();
+        return delegate.size();
+    }
+
+    @Override
+    public Object[] toArray() {
+        timeoutVerifier.checkTimeout();
+        return delegate.toArray();
+    }
+
+    @Override
+    public <O> O[] toArray(O[] a) {
+        timeoutVerifier.checkTimeout();
+        return delegate.toArray(a);
+    }
+}