io.reactiverse.pgclient.impl.ConnectionPool Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of reactive-pg-client Show documentation
The reactive Postgres client
There is a newer version: 0.11.4
/*
 * 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.reactiverse.pgclient.impl;

import io.vertx.core.*;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Consumer;

/**
 * @author Julien Viet
 */
public class ConnectionPool {

  private final Consumer>> connector;
  private final int maxSize;
  private final ArrayDeque> waiters = new ArrayDeque<>();
  private final Set all = new HashSet<>();
  private final ArrayDeque available = new ArrayDeque<>();
  private int size;

  public ConnectionPool(Consumer>> connector, int maxSize) {
    this.maxSize = maxSize;
    this.connector = connector;
  }

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

  public void acquire(Handler> holder) {
    waiters.add(Future.future().setHandler(holder));
    check();
  }

  public void close() {
    for (PooledConnection pooled : new ArrayList<>(all)) {
      pooled.close();
    }
  }

  class PooledConnection implements Connection, Connection.Holder  {

    private final Connection conn;
    private Holder holder;

    PooledConnection(Connection conn) {
      this.conn = conn;
    }

    @Override
    public Connection connection() {
      return this;
    }

    @Override
    public boolean isSsl() {
      return conn.isSsl();
    }

    @Override
    public void schedule(CommandBase cmd) {
      conn.schedule(cmd);
    }

    /**
     * Close the underlying connection
     */
    void close() {
      conn.close(this);
    }

    @Override
    public void init(Holder holder) {
      if (this.holder != null) {
        throw new IllegalStateException();
      }
      this.holder = holder;
    }

    @Override
    public void close(Holder holder) {
      if (holder != this.holder) {
        throw new IllegalStateException();
      }
      this.holder = null;
      release(this);
    }

    @Override
    public void handleClosed() {
      if (all.remove(this)) {
        size--;
        if (holder == null) {
          available.remove(this);
        } else {
          holder.handleClosed();
        }
        check();
      } else {
        throw new IllegalStateException();
      }
    }

    @Override
    public void handleNotification(int processId, String channel, String payload) {
      if (holder != null) {
        holder.handleNotification(processId, channel, payload);
      }
    }

    @Override
    public void handleException(Throwable err) {
      if (holder != null) {
        holder.handleException(err);
      }
    }
  }

  private void release(PooledConnection proxy) {
    available.add(proxy);
    check();
  }

  private void check() {
    if (waiters.size() > 0) {
      if (available.size() > 0) {
        PooledConnection proxy = available.poll();
        Future waiter = waiters.poll();
        waiter.complete(proxy);
      } else {
        if (size < maxSize) {
          size++;
          connector.accept(ar -> {
            if (ar.succeeded()) {
              Connection conn = ar.result();
              PooledConnection proxy = new PooledConnection(conn);
              all.add(proxy);
              conn.init(proxy);
              release(proxy);
            } else {
              Future waiter;
              while ((waiter = waiters.poll()) != null) {
                waiter.fail(ar.cause());
              }
            }
          });
        }
      }
    }
  }
}