se.fortnox.reactivewizard.client.UnsubscribeAwareConnectionProviderFactory Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of reactivewizard-client Show documentation

There is a newer version: 24.6.0

Show newest version

package se.fortnox.reactivewizard.client;

import io.netty.buffer.ByteBuf;
import io.reactivex.netty.channel.Connection;
import io.reactivex.netty.client.ConnectionProvider;
import io.reactivex.netty.client.ConnectionProviderFactory;
import io.reactivex.netty.client.HostConnector;
import io.reactivex.netty.client.pool.PoolConfig;
import io.reactivex.netty.client.pool.PoolLimitDeterminationStrategy;
import rx.Observable;

import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Bugfix for RxNetty, which will not release a connection pool permit if the connection observable is unsubscribed
 * (from a timeout) before the connection is returned or before an error occurs.
 */
public class UnsubscribeAwareConnectionProviderFactory implements ConnectionProviderFactory {
    private final ConnectionProviderFactory wrapped;
    private final PoolConfig                poolConfig;

    public UnsubscribeAwareConnectionProviderFactory(ConnectionProviderFactory wrapped, PoolConfig poolConfig) {
        this.wrapped = wrapped;
        this.poolConfig = poolConfig;
    }

    @Override
    public ConnectionProvider newProvider(Observable> hosts) {
        return new UnsubscribeAwareConnectionProvider(wrapped.newProvider(hosts), poolConfig.getPoolLimitDeterminationStrategy());
    }

    private static class UnsubscribeAwareConnectionProvider implements ConnectionProvider {
        private final ConnectionProvider wrapped;
        private final PoolLimitDeterminationStrategy       poolLimitDeterminationStrategy;

        public UnsubscribeAwareConnectionProvider(ConnectionProvider wrapped, PoolLimitDeterminationStrategy poolLimitDeterminationStrategy) {
            this.wrapped = wrapped;
            this.poolLimitDeterminationStrategy = poolLimitDeterminationStrategy;
        }

        @Override
        public Observable> newConnectionRequest() {
            AtomicBoolean connectionReturned = new AtomicBoolean();
            return wrapped.newConnectionRequest()
                .doOnNext(c -> connectionReturned.set(true))
                .doOnError(e -> connectionReturned.set(true))
                .doOnUnsubscribe(() -> {
                    if (!connectionReturned.get()) {
                        // Unsubscribed before a connection was returned. Must give back a permit.
                        poolLimitDeterminationStrategy.releasePermit();
                    }
                });
        }
    }
}