Issue 475: PgServer: add support for CancelRequest, patch from Andrew Franklin

h2/src/docsrc/html/changelog.html

@@ -35,6 +35,7 @@ Change Log
 </li><li>Improve error message when check constraint is broken, test case from Gili (cowwoc)
 </li><li>Improve error message when we have a unique constraint violation, displays the offending key in the error message
 </li><li>Issue 478: Support for "SHOW TRANSACTION ISOLATION LEVEL", patch from Andrew Franklin
+</li><li>Issue 475: PgServer: add support for CancelRequest, patch from Andrew Franklin
 </li></ul>
 
 <h2>Version 1.3.172 (2013-05-25)</h2>

h2/src/main/org/h2/jdbc/JdbcStatement.java

@@ -39,6 +39,7 @@ public class JdbcStatement extends TraceObject implements Statement {
     private int lastExecutedCommandType;
     private ArrayList<String> batchCommands;
     private boolean escapeProcessing = true;
+    private boolean cancelled = false;
 
     JdbcStatement(JdbcConnection conn, int id, int resultSetType, int resultSetConcurrency, boolean closeWithResultSet) {
         this.conn = conn;
@@ -530,6 +531,7 @@ public class JdbcStatement extends TraceObject implements Statement {
             try {
                 if (c != null) {
                     c.cancel();
+                    cancelled = true;
                 }
             } finally {
                 setExecutingStatement(null);
@@ -539,6 +541,13 @@ public class JdbcStatement extends TraceObject implements Statement {
         }
     }
 
+    /**
+     * @return true if this statement has been cancelled.
+     */
+    public boolean isCancelled() {
+        return cancelled;
+    }
+
     /**
      * Gets the current query timeout in seconds.
      * This method will return 0 if no query timeout is set.

h2/src/main/org/h2/server/pg/PgServer.java

@@ -19,6 +19,7 @@ import java.sql.Types;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.h2.engine.Constants;
 import org.h2.message.DbException;
 import org.h2.server.Service;
@@ -78,6 +79,7 @@ public class PgServer implements Service {
     private boolean trace;
     private ServerSocket serverSocket;
     private final Set<PgServerThread> running = Collections.synchronizedSet(new HashSet<PgServerThread>());
+    private final AtomicInteger pid= new AtomicInteger(0);
     private String baseDir;
     private boolean allowOthers;
     private boolean isDaemon;
@@ -191,7 +193,7 @@ public class PgServer implements Service {
                 } else {
                     PgServerThread c = new PgServerThread(s, this);
                     running.add(c);
-                    c.setProcessId(running.size());
+                    c.setProcessId(pid.incrementAndGet());
                     Thread thread = new Thread(c, threadName+" thread");
                     thread.setDaemon(isDaemon);
                     c.setThread(thread);
@@ -252,6 +254,18 @@ public class PgServer implements Service {
         }
     }
 
+    /**
+     * @return the thread with the given process-id
+     */
+    PgServerThread getThread(int processId) {
+        for (PgServerThread c : New.arrayList(running)) {
+            if (c.getProcessId()==processId) {
+                return c;
+            }
+        }
+        return null;
+    }
+
     String getBaseDir() {
         return baseDir;
     }

h2/src/main/org/h2/server/pg/PgServerThread.java

@@ -28,6 +28,7 @@ import java.sql.Statement;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Properties;
+import java.util.Random;
 import org.h2.command.CommandInterface;
 import org.h2.constant.SysProperties;
 import org.h2.engine.ConnectionInfo;
@@ -62,6 +63,8 @@ public class PgServerThread implements Runnable {
     private String userName;
     private String databaseName;
     private int processId;
+    private int secret;
+    private JdbcStatement activeRequest;
     private String clientEncoding = SysProperties.PG_DEFAULT_CLIENT_ENCODING;
     private String dateStyle = "ISO";
     private final HashMap<String, Prepared> prepared = new CaseInsensitiveMap<Prepared>();
@@ -70,6 +73,7 @@ public class PgServerThread implements Runnable {
     PgServerThread(Socket socket, PgServer server) {
         this.server = server;
         this.socket = socket;
+        this.secret = new Random().nextInt();
     }
 
     @Override
@@ -142,9 +146,16 @@ public class PgServerThread implements Runnable {
             server.trace("Init");
             int version = readInt();
             if (version == 80877102) {
-                server.trace("CancelRequest (not supported)");
-                server.trace(" pid: " + readInt());
-                server.trace(" key: " + readInt());
+                server.trace("CancelRequest");
+                int pid = readInt();
+                int key = readInt();
+                PgServerThread c = server.getThread(pid);
+                if (c!=null) {
+                    c.cancelRequest(key);
+                } else {
+                    sendErrorResponse("unknown process: "+pid);
+                }
+                close();
             } else if (version == 80877103) {
                 server.trace("SSLRequest");
                 out.write('N');
@@ -328,6 +339,7 @@ public class PgServerThread implements Runnable {
             server.trace(prepared.sql);
             try {
                 prep.setMaxRows(maxRows);
+                setActiveRequest(prep);
                 boolean result = prep.execute();
                 if (result) {
                     try {
@@ -345,8 +357,14 @@ public class PgServerThread implements Runnable {
                     sendCommandComplete(prep, prep.getUpdateCount());
                 }
             } catch (Exception e) {
+                if (prep.isCancelled()) {
+                    sendCancelQueryResponse();
+                } else {
                     sendErrorResponse(e);
                 }
+            } finally {
+                setActiveRequest(null);
+            }
             break;
         }
         case 'S': {
@@ -367,6 +385,7 @@ public class PgServerThread implements Runnable {
                     }
                     s = getSQL(s);
                     stat = (JdbcStatement) conn.createStatement();
+                    setActiveRequest(stat);
                     boolean result = stat.execute(s);
                     if (result) {
                         ResultSet rs = stat.getResultSet();
@@ -385,10 +404,15 @@ public class PgServerThread implements Runnable {
                         sendCommandComplete(stat, stat.getUpdateCount());
                     }
                 } catch (SQLException e) {
+                    if (stat!=null && stat.isCancelled()) {
+                        sendCancelQueryResponse();
+                    } else {
                         sendErrorResponse(e);
+                    }
                     break;
                 } finally {
                     JdbcUtils.closeSilently(stat);
+                    setActiveRequest(null);
                 }
             }
             sendReadyForQuery();
@@ -513,6 +537,19 @@ public class PgServerThread implements Runnable {
         sendMessage();
     }
 
+    private void sendCancelQueryResponse() throws IOException {
+        server.trace("CancelSuccessResponse");
+        startMessage('E');
+        write('S');
+        writeString("ERROR");
+        write('C');
+        writeString("57014");
+        write('M');
+        writeString("canceling statement due to user request");
+        write(0);
+        sendMessage();
+    }
+
     private void sendParameterDescription(Prepared p) throws IOException {
         try {
             PreparedStatement prep = p.prep;
@@ -746,7 +783,7 @@ public class PgServerThread implements Runnable {
     private void sendBackendKeyData() throws IOException {
         startMessage('K');
         writeInt(processId);
-        writeInt(processId);
+        writeInt(secret);
         sendMessage();
     }
 
@@ -811,6 +848,34 @@ public class PgServerThread implements Runnable {
         this.processId = id;
     }
 
+    int getProcessId() {
+        return this.processId;
+    }
+
+    synchronized void setActiveRequest(JdbcStatement statement) {
+        activeRequest = statement;
+    }
+
+    /**
+     * Kill a currently running query on this thread.
+     * @param secret the private key of the command
+     * @return true if the command was successfully killed
+     */
+    synchronized boolean cancelRequest(int secret) throws IOException {
+        if (activeRequest != null)
+        {
+            if (secret != this.secret) throw new IOException("invalid cancel secret");
+            try {
+                activeRequest.cancel();
+                activeRequest = null;
+            } catch (SQLException e) {
+                throw DbException.convert(e);
+            }
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Represents a PostgreSQL Prepared object.
      */
@@ -857,5 +922,4 @@ public class PgServerThread implements Runnable {
          */
         Prepared prep;
     }
-
 }

h2/src/test/org/h2/test/unit/TestPgServer.java

@@ -16,6 +16,7 @@ import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.sql.Types;
+import java.util.concurrent.*;
 
 import org.h2.test.TestBase;
 import org.h2.tools.Server;
@@ -50,6 +51,50 @@ public class TestPgServer extends TestBase {
         } finally {
             server.stop();
         }
+    }
+
+    private void testCancelQuery() throws Exception {
+        Server server = Server.createPgServer("-pgPort", "5535", "-pgDaemon", "-key", "test", "mem:test");
+        server.start();
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            Class.forName("org.postgresql.Driver");
+
+            Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5535/test", "sa", "sa");
+            final Statement stat = conn.createStatement();
+            stat.execute("create alias sleep for \"java.lang.Thread.sleep\"");
+
+            // create a table with 200 rows (cancel interval is 127)
+            stat.execute("create table test(id int)");
+            for (int i=0; i<200; i++) {
+                stat.execute("insert into test (id) values (rand())");
+            }
+
+            Future<Boolean> future = executor.submit(new Callable<Boolean>() {
+                public Boolean call() throws Exception {
+                    return stat.execute("select id, sleep(5) from test");
+                }
+            });
+
+            // give it a little time to start and then cancel it
+            Thread.sleep(100);
+            stat.cancel();
+
+            try {
+                future.get();
+                throw new IllegalStateException();
+            } catch (ExecutionException e) {
+                assertStartsWith(e.getCause().getMessage(), "ERROR: Statement was canceled");
+            } finally {
+                conn.close();
+            }
+        } catch (ClassNotFoundException e) {
+            println("PostgreSQL JDBC driver not found - PgServer not tested");
+        } finally {
+            server.stop();
+            executor.shutdown();
+        }
         deleteDb("test");
     }