io.vertx.sqlclient.impl.PoolImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vertx-sql-client Show documentation
The Reactive SQL Client
There is a newer version: 5.0.0.CR1
/*
 * Copyright (C) 2017 Julien Viet
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License 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 io.vertx.sqlclient.impl;

import io.vertx.codegen.annotations.Nullable;
import io.vertx.core.*;
import io.vertx.core.impl.CloseFuture;
import io.vertx.core.impl.ContextInternal;
import io.vertx.core.impl.VertxInternal;
import io.vertx.core.impl.future.PromiseInternal;
import io.vertx.sqlclient.*;
import io.vertx.sqlclient.impl.command.CommandBase;
import io.vertx.sqlclient.impl.pool.SqlConnectionPool;
import io.vertx.sqlclient.spi.Driver;

import java.util.function.Function;

import static java.util.concurrent.TimeUnit.MILLISECONDS;

/**
 * @author Julien Viet
 * @author Emad Alblueshi
 */
public class PoolImpl extends SqlClientBase implements Pool, Closeable {

  private final VertxInternal vertx;
  private final SqlConnectionPool pool;
  private final CloseFuture closeFuture;
  private final long idleTimeout;
  private final long connectionTimeout;
  private final long maxLifetime;
  private final long cleanerPeriod;
  private final boolean pipelined;
  private volatile Handler connectionInitializer;
  private long timerID;
  private volatile Function> connectionProvider;

  public static final String PROPAGATABLE_CONNECTION = "propagatable_connection";

  public PoolImpl(VertxInternal vertx,
                  Driver driver,
                  boolean pipelined,
                  PoolOptions poolOptions,
                  Function> afterAcquire,
                  Function> beforeRecycle,
                  CloseFuture closeFuture) {
    super(driver);

    this.idleTimeout = MILLISECONDS.convert(poolOptions.getIdleTimeout(), poolOptions.getIdleTimeoutUnit());
    this.connectionTimeout = MILLISECONDS.convert(poolOptions.getConnectionTimeout(), poolOptions.getConnectionTimeoutUnit());
    this.maxLifetime = MILLISECONDS.convert(poolOptions.getMaxLifetime(), poolOptions.getMaxLifetimeUnit());
    this.cleanerPeriod = poolOptions.getPoolCleanerPeriod();
    this.timerID = -1L;
    this.pipelined = pipelined;
    this.vertx = vertx;
    this.pool = new SqlConnectionPool(ctx -> connectionProvider.apply(ctx), () -> connectionInitializer,
      afterAcquire, beforeRecycle, vertx, idleTimeout, maxLifetime, poolOptions.getMaxSize(), pipelined,
      poolOptions.getMaxWaitQueueSize(), poolOptions.getEventLoopSize());
    this.closeFuture = closeFuture;
  }

  public Pool init() {
    closeFuture.add(this);
    if ((idleTimeout > 0 || maxLifetime > 0) && cleanerPeriod > 0) {
      synchronized (this) {
        timerID = vertx.setTimer(cleanerPeriod, id -> {
          runEviction();
        });
      }
    }
    return this;
  }

  public Pool connectionProvider(Function> connectionProvider) {
    if (connectionProvider == null) {
      throw new NullPointerException();
    }
    this.connectionProvider = connectionProvider;
    return this;
  }

  private void runEviction() {
    synchronized (this) {
      if (timerID == -1) {
        // Cancelled
        return;
      }
      timerID = vertx.setTimer(cleanerPeriod, id -> {
        runEviction();
      });
    }
    pool.evict();
  }

  @Override
  protected  PromiseInternal promise() {
    return vertx.promise();
  }

  protected ContextInternal context() {
    return vertx.getOrCreateContext();
  }

  @Override
  protected  PromiseInternal promise(Handler> handler) {
    return vertx.promise(handler);
  }

  @Override
  public void getConnection(Handler> handler) {
    Future fut = getConnection();
    if (handler != null) {
      fut.onComplete(handler);
    }
  }

  @Override
  public Future getConnection() {
    ContextInternal current = vertx.getOrCreateContext();
    if (pipelined) {
      return current.failedFuture("Cannot acquire a connection on a pipelined pool");
    }
    Promise promise = current.promise();
    acquire(current, connectionTimeout, promise);
    return promise.future().map(conn -> {
      SqlConnectionInternal wrapper = driver.wrapConnection(current, conn.factory(), conn);
      conn.init(wrapper);
      return wrapper;
    });
  }

  public static  Future<@Nullable T> startPropagatableConnection(Pool pool, Function> function) {
    ContextInternal context = (ContextInternal) Vertx.currentContext();
    return pool.getConnection().onComplete(handler -> context.putLocal(PROPAGATABLE_CONNECTION, handler.result()))
      .flatMap(conn -> conn
        .begin()
        .flatMap(tx -> function
          .apply(conn)
          .compose(
            res -> tx
              .commit()
              .flatMap(v -> context.succeededFuture(res)),
            err -> {
              if (err instanceof TransactionRollbackException) {
                return context.failedFuture(err);
              } else {
                return tx
                  .rollback()
                  .compose(v -> context.failedFuture(err), failure -> context.failedFuture(err));
              }
            }))
        .onComplete(ar -> conn.close(v -> context.removeLocal(PROPAGATABLE_CONNECTION))));
  }

  @Override
  public  Future schedule(ContextInternal context, CommandBase cmd) {
    return pool.execute(context, cmd);
  }

  private void acquire(ContextInternal context, long timeout, Handler> completionHandler) {
    pool.acquire(context, timeout, completionHandler);
  }

  @Override
  public void close(Promise completion) {
    doClose().onComplete(completion);
  }

  @Override
  public Future close() {
    Promise promise = vertx.promise();
    closeFuture.close(promise);
    return promise.future();
  }

  @Override
  public void close(Handler> handler) {
    closeFuture.close(vertx.promise(handler));
  }

  @Override
  public Pool connectHandler(Handler handler) {
    if (handler != null) {
      connectionInitializer = conn -> {
        ContextInternal current = vertx.getContext();
        SqlConnectionInternal wrapper = driver.wrapConnection(current, conn.factory(), conn);
        conn.init(wrapper);
        current.dispatch(wrapper, handler);
      };
    } else {
      connectionInitializer = null;
    }
    return this;
  }

  private Future doClose() {
    synchronized (this) {
      if (timerID >= 0) {
        vertx.cancelTimer(timerID);
        timerID = -1;
      }
    }
    return pool.close();
  }

  public int size() {
    return pool.size();
  }
}