com.julienviet.pgclient.impl.ConnectionPool Maven / Gradle / Ivy

Go to download
/*
 * 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 com.julienviet.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 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;
      available.add(this);
      check();
    }

    @Override
    public void handleClosed() {
      if (all.remove(this)) {
        size--;
        if (holder != null) {
          available.remove(this);
        }
        if (holder != null) {
          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 doAcq(Handler> handler) {
    if (available.size() > 0) {
      PooledConnection proxy = available.poll();
      handler.handle(Future.succeededFuture(proxy));
    } else {
      if (size < maxSize) {
        size++;
        connector.accept(ar -> {
          if (ar.succeeded()) {
            Connection conn = ar.result();
            PooledConnection proxy = new PooledConnection(conn);
            all.add(proxy);
            available.add(proxy);
            conn.init(proxy);
            doAcq(handler);
          } else {
            handler.handle(Future.failedFuture(ar.cause()));
          }
        });
      }
    }
  }

  private void check() {
    if (waiters.size() > 0) {
      doAcq(ar -> {
        if (ar.succeeded()) {
          PooledConnection proxy = ar.result();
          Future waiter = waiters.poll();
          waiter.complete(proxy);
        } else {
          Future waiter;
          while ((waiter = waiters.poll()) != null) {
            waiter.fail(ar.cause());
          }
        }
      });
    }
  }
}