TimeoutScope.java

package com.mongodb.client;

import java.time.Duration;
import java.util.Optional;

public class TimeoutScope implements AutoCloseable {
    private static final ThreadLocal<TimeoutScope> THREAD_LOCAL_TIMEOUT_SCOPE = new ThreadLocal<>();

    public static TimeoutScope withDuration(final Duration duration) {
        if (THREAD_LOCAL_TIMEOUT_SCOPE.get() != null) {
            throw new IllegalStateException("Not supported yet"); // TODO: support nested timeout scopes
        }
        TimeoutScope timeoutScope = new TimeoutScope(duration);
        THREAD_LOCAL_TIMEOUT_SCOPE.set(timeoutScope);
        return timeoutScope;
    }

    public static Optional<TimeoutScope> current() {
        return Optional.ofNullable(THREAD_LOCAL_TIMEOUT_SCOPE.get());
    }

    private final long startTime = System.nanoTime();
    private final Duration duration;
    private TimeoutScope(final Duration duration) {
        this.duration = duration;
    }

    public long getStartTime() {
        return startTime;
    }

    private Duration getDuration() {
        return duration;
    }

    /**
     * May be negative
     *
     * @return the remaining duration on the timeout, which may be negative
     */
    public Duration remaining() {
        return duration.minus(Duration.ofNanos(System.nanoTime() - startTime));
    }

    @Override
    public void close() throws Exception {
        THREAD_LOCAL_TIMEOUT_SCOPE.remove();
    }
}