com.microsoft.sqlserver.jdbc.ScopeTokenCache Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of mssql-jdbc Show documentation

Microsoft JDBC Driver for SQL Server.

There is a newer version: 12.7.0.jre11-preview

/*
 * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
 * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
 */
package com.microsoft.sqlserver.jdbc;

import com.azure.core.credential.AccessToken;
import com.azure.core.credential.TokenRequestContext;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Mono;
import reactor.core.publisher.ReplayProcessor;


/**
 * A token cache that supports caching a token and refreshing it.
 */
class ScopeTokenCache {

    private final AtomicBoolean wip;
    private AccessToken cache;
    private final ReplayProcessor emitterProcessor = ReplayProcessor.create(1);
    private final FluxSink sink = emitterProcessor.sink(FluxSink.OverflowStrategy.BUFFER);
    private final Function> getNew;
    private TokenRequestContext request;

    /**
     * Creates an instance of RefreshableTokenCredential with default scheme "Bearer".
     *
     * @param getNew
     *        a method to get a new token
     */
    ScopeTokenCache(Function> getNew) {
        this.wip = new AtomicBoolean(false);
        this.getNew = getNew;
    }

    void setRequest(TokenRequestContext request) {
        this.request = request;
    }

    /**
     * Asynchronously get a token from either the cache or replenish the cache with a new token.
     * 
     * @return a Publisher that emits an AccessToken
     */
    Mono getToken() {
        if (cache != null && !cache.isExpired()) {
            return Mono.just(cache);
        }
        return Mono.defer(() -> {
            if (!wip.getAndSet(true)) {
                return getNew.apply(request).doOnNext(ac -> cache = ac).doOnNext(sink::next).doOnError(sink::error)
                        .doOnTerminate(() -> wip.set(false));
            } else {
                return emitterProcessor.next();
            }
        });
    }
}