Rework the prepared statement cache and evict the statement from the … (#597)

* rework the prepared statement cache and evict the statement from the cache when it's closed - fixes #572

Signed-off-by: Billy Yuan <billy112487983@gmail.com>

* prepared statement cache will only cache successful responses now

Signed-off-by: Billy Yuan <billy112487983@gmail.com>

* evict the invalid prepared statement when the postgres table is changed - fixes #619

Signed-off-by: Billy Yuan <billy112487983@gmail.com>

* remove unnecessary lines

Signed-off-by: Billy Yuan <billy112487983@gmail.com>

* use sql string as the prepared statement cache entry removing identifier and some minors for naming

Signed-off-by: Billy Yuan <billy112487983@gmail.com>

Author: BillyYccc

vertx-pg-client/docker/postgres/resources/create-postgres.sql

@@ -347,6 +347,19 @@ CREATE TABLE mutable
 );
 -- mutable for insert,update,delete query testing
 
+-- unstable table for table schema changing testing
+DROP TABLE IF EXISTS unstable;
+CREATE TABLE unstable
+(
+  id      integer       NOT NULL,
+  message varchar(2048) NOT NULL,
+  PRIMARY KEY (id)
+);
+
+INSERT INTO unstable (id, message)
+VALUES (1, 'fortune: No such file or directory');
+-- unstable table for table schema changing testing
+
 -- table for test ANSI SQL data type codecs
 DROP TABLE IF EXISTS basicdatatype;
 CREATE TABLE basicdatatype

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/ExtendedBatchQueryCommandCodec.java

@@ -30,6 +30,7 @@ class ExtendedBatchQueryCommandCodec<R> extends ExtendedQueryCommandBaseCodec<R,
 
   @Override
   void encode(PgEncoder encoder) {
+    super.encode(encoder);
     if (cmd.isSuspended()) {
       encoder.writeExecute(cmd.cursorId(), cmd.fetch());
       encoder.writeSync();

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/ExtendedQueryCommandBaseCodec.java

@@ -16,16 +16,26 @@
  */
 package io.vertx.pgclient.impl.codec;
 
+import io.vertx.sqlclient.impl.codec.InvalidCachedStatementEvent;
 import io.vertx.sqlclient.impl.RowDesc;
 import io.vertx.sqlclient.impl.command.ExtendedQueryCommandBase;
 
 abstract class ExtendedQueryCommandBaseCodec<R, C extends ExtendedQueryCommandBase<R>> extends QueryCommandBaseCodec<R, C> {
 
+  private PgEncoder encoder;
+
+  private static final String TABLE_SCHEMA_CHANGE_ERROR_MESSAGE_PATTERN = "bind message has \\d result formats but query has \\d columns";
+
   ExtendedQueryCommandBaseCodec(C cmd) {
     super(cmd);
     decoder = new RowResultDecoder<>(cmd.collector(), ((PgPreparedStatement)cmd.preparedStatement()).rowDesc());
   }
 
+  @Override
+  void encode(PgEncoder encoder) {
+    this.encoder = encoder;
+  }
+
   @Override
   void handleRowDescription(PgRowDesc rowDescription) {
     decoder = new RowResultDecoder<>(cmd.collector(), rowDescription);
@@ -51,4 +61,12 @@ void handlePortalSuspended() {
   void handleBindComplete() {
     // Response to Bind
   }
+
+  @Override
+  public void handleErrorResponse(ErrorResponse errorResponse) {
+    if (cmd.preparedStatement().cacheable() && errorResponse.getMessage().matches(TABLE_SCHEMA_CHANGE_ERROR_MESSAGE_PATTERN)) {
+      encoder.channelHandlerContext().fireChannelRead(new InvalidCachedStatementEvent(cmd.preparedStatement().sql()));
+    }
+    super.handleErrorResponse(errorResponse);
+  }
 }

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/ExtendedQueryCommandCodec.java

@@ -28,6 +28,7 @@ class ExtendedQueryCommandCodec<R> extends ExtendedQueryCommandBaseCodec<R, Exte
 
   @Override
   void encode(PgEncoder encoder) {
+    super.encode(encoder);
     if (cmd.isSuspended()) {
       encoder.writeExecute(cmd.cursorId(), cmd.fetch());
       encoder.writeSync();

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgEncoder.java

@@ -440,4 +440,8 @@ private void ensureBuffer() {
   long nextStatementName() {
     return psSeq.next();
   }
+
+  public ChannelHandlerContext channelHandlerContext() {
+    return ctx;
+  }
 }

vertx-pg-client/src/test/java/io/vertx/pgclient/PreparedStatementCachedTest.java

@@ -17,11 +17,45 @@
 
 package io.vertx.pgclient;
 
+import io.vertx.ext.unit.Async;
+import io.vertx.ext.unit.TestContext;
+import io.vertx.sqlclient.Tuple;
+import org.junit.Test;
+
 public class PreparedStatementCachedTest extends PreparedStatementTestBase {
 
   @Override
   protected PgConnectOptions options() {
     return new PgConnectOptions(options).setCachePreparedStatements(true);
   }
 
+  @Test
+  public void testOneShotPreparedQueryCacheRefreshOnTableSchemaChange(TestContext ctx) {
+    Async async = ctx.async();
+    PgConnection.connect(vertx, options(), ctx.asyncAssertSuccess(conn -> {
+      conn.preparedQuery("SELECT * FROM unstable WHERE id=$1").execute(Tuple.of(1), ctx.asyncAssertSuccess(res1 -> {
+        ctx.assertEquals(1, res1.size());
+        Tuple row1 = res1.iterator().next();
+        ctx.assertEquals(1, row1.getInteger(0));
+        ctx.assertEquals("fortune: No such file or directory", row1.getString(1));
+
+        // change table schema
+        conn.query("ALTER TABLE unstable DROP COLUMN message").execute(ctx.asyncAssertSuccess(dropColumn -> {
+          // failure due to schema change
+          conn.preparedQuery("SELECT * FROM unstable WHERE id=$1").execute(Tuple.of(1), ctx.asyncAssertFailure(failure -> {
+            // recover because the cache is refreshed
+            conn.preparedQuery("SELECT * FROM unstable WHERE id=$1").execute(Tuple.of(1), ctx.asyncAssertSuccess(res2 -> {
+              ctx.assertEquals(1, res2.size());
+              Tuple row2 = res2.iterator().next();
+              ctx.assertEquals(1, row2.getInteger(0));
+              ctx.assertEquals(null, row2.getString(1)); // the message column is removed
+              conn.close();
+              async.complete();
+            }));
+          }));
+        }));
+      }));
+    }));
+  }
+
 }

vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/PreparedStatementCache.java

@@ -30,10 +30,11 @@
 import io.vertx.core.net.impl.NetSocketInternal;
 import io.vertx.core.logging.Logger;
 import io.vertx.core.logging.LoggerFactory;
+import io.vertx.sqlclient.impl.cache.PreparedStatementCache;
+import io.vertx.sqlclient.impl.codec.InvalidCachedStatementEvent;
 import io.vertx.sqlclient.impl.command.*;
 
 import java.util.ArrayDeque;
-import java.util.Deque;
 
 /**
  * @author <a href="mailto:julien@julienviet.com">Julien Viet</a>
@@ -68,7 +69,7 @@ public SocketConnectionBase(NetSocketInternal socket,
     this.socket = socket;
     this.context = context;
     this.pipeliningLimit = pipeliningLimit;
-    this.psCache = cachePreparedStatements ? new PreparedStatementCache(preparedStatementCacheSize, this) : null;
+    this.psCache = cachePreparedStatements ? new PreparedStatementCache(this, preparedStatementCacheSize) : null;
     this.preparedStatementCacheSqlLimit = preparedStatementCacheSqlLimit;
   }
 
@@ -162,30 +163,33 @@ protected <R> void doSchedule(CommandBase<R> cmd, Handler<AsyncResult<R>> handle
 
     // Special handling for cache
     PreparedStatementCache psCache = this.psCache;
-    if (psCache != null && cmd instanceof PrepareStatementCommand) {
-      PrepareStatementCommand psCmd = (PrepareStatementCommand) cmd;
-      if (!psCmd.cacheable()) {
-        // we don't cache non one-shot preparedQuery
-      } else if (psCmd.sql().length() > preparedStatementCacheSqlLimit) {
-        // do not cache the statements
-      } else {
-        // TODO fix the auto-closing logic
-        CachedPreparedStatement cached = psCache.get(psCmd.sql());
-        Handler<AsyncResult<PreparedStatement>> orig = (Handler) handler;
-        if (cached != null) {
-          psCmd.handler = orig;
-          cached.get(psCmd::complete);
-          return;
+    if (psCache != null) {
+      // cache is enabled
+      if (cmd instanceof PrepareStatementCommand) {
+        PrepareStatementCommand psCmd = (PrepareStatementCommand) cmd;
+        if (!psCmd.cacheable()) {
+          // we don't cache non one-shot preparedQuery
+        } else if (psCmd.sql().length() > preparedStatementCacheSqlLimit) {
+          // do not cache the statements if it exceeds the sql length limit
         } else {
-          if (psCache.size() >= psCache.getCapacity() && !psCache.isReady()) {
-            // only if the prepared statement is ready then it can be evicted
+          Handler<AsyncResult<PreparedStatement>> originalHandler = (Handler) handler;
+          Handler<AsyncResult<PreparedStatement>> newHandler = psCache.appendStmtReq(psCmd.sql(), originalHandler);
+          if (newHandler == null) {
+            // we don't need to schedule it if the result is cached or the request has been sent
+            return;
           } else {
-            cached = new CachedPreparedStatement();
-            cached.get(orig);
-            psCache.put(psCmd.sql(), cached);
-            handler = (Handler) cached;
+            handler = (Handler) newHandler;
           }
         }
+      } else if (cmd instanceof CloseStatementCommand) {
+        CloseStatementCommand closeStmtCommand = (CloseStatementCommand) cmd;
+        /*
+         * We need to know how we handle the close statement command, this cmd might origin from PreparedStatement#close
+         * or it's automatically sent by the cache once the stmt is evicted, we should clean up the cache for those closing cached prepared statements.
+         */
+        if (closeStmtCommand.statement().cacheable()) {
+          psCache.remove(closeStmtCommand.statement().sql());
+        }
       }
     }
 
@@ -201,29 +205,6 @@ protected <R> void doSchedule(CommandBase<R> cmd, Handler<AsyncResult<R>> handle
     }
   }
 
-  static class CachedPreparedStatement implements Handler<AsyncResult<PreparedStatement>> {
-
-    private final Deque<Handler<AsyncResult<PreparedStatement>>> waiters = new ArrayDeque<>();
-    AsyncResult<PreparedStatement> resp;
-
-    void get(Handler<AsyncResult<PreparedStatement>> handler) {
-      if (resp != null) {
-        handler.handle(resp);
-      } else {
-        waiters.add(handler);
-      }
-    }
-
-    @Override
-    public void handle(AsyncResult<PreparedStatement> event) {
-      resp = event;
-      Handler<AsyncResult<PreparedStatement>> waiter;
-      while ((waiter = waiters.poll()) != null) {
-        waiter.handle(resp);
-      }
-    }
-  }
-
   private void checkPending() {
     ChannelHandlerContext ctx = socket.channelHandlerContext();
     if (inflight < pipeliningLimit) {
@@ -242,6 +223,9 @@ protected void handleMessage(Object msg) {
       checkPending();
       CommandResponse resp =(CommandResponse) msg;
       resp.fire();
+    } else if (msg instanceof InvalidCachedStatementEvent) {
+      InvalidCachedStatementEvent event = (InvalidCachedStatementEvent) msg;
+      removeCachedStatement(event.sql());
     }
   }
 
@@ -251,6 +235,12 @@ protected void handleEvent(Object event) {
     }
   }
 
+  private void removeCachedStatement(String sql) {
+    if (this.psCache != null) {
+      this.psCache.remove(sql);
+    }
+  }
+
   private void handleClosed(Void v) {
     handleClose(null);
   }

vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/SocketConnectionBase.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2011-2020 Contributors to the Eclipse Foundation
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
+ * which is available at https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
+ */
+
+package io.vertx.sqlclient.impl.cache;
+
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.sqlclient.impl.PreparedStatement;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+
+class InflightCachingStmtEntry implements Handler<AsyncResult<PreparedStatement>> {
+  private final Deque<Handler<AsyncResult<PreparedStatement>>> waiters = new ArrayDeque<>();
+  private final String sql;
+  private final PreparedStatementCache psCache;
+
+  InflightCachingStmtEntry(String sql, PreparedStatementCache psCache) {
+    this.sql = sql;
+    this.psCache = psCache;
+  }
+
+  void addWaiter(Handler<AsyncResult<PreparedStatement>> handler) {
+    waiters.add(handler);
+  }
+
+  @Override
+  public void handle(AsyncResult<PreparedStatement> preparedStatementResult) {
+    if (preparedStatementResult.succeeded()) {
+      // put it in the cache since the response is ready
+      // we only need to cache successful responses here
+      psCache.cache().put(sql, preparedStatementResult);
+    }
+    psCache.inflight().remove(sql);
+    Handler<AsyncResult<PreparedStatement>> waiter;
+    while ((waiter = waiters.poll()) != null) {
+      waiter.handle(preparedStatementResult);
+    }
+  }
+}