package com.tibbo.aggregate.common.context;

import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

public class AwaitableRequestController extends DefaultRequestController
{

    public enum Result { SUCCESS, FAILURE }

    private final CountDownLatch latch = new CountDownLatch(1);
    private final AtomicReference<Result> result = new AtomicReference<>(null);
    private final AtomicReference<Throwable> error = new AtomicReference<>(null);

    public void markSuccess()
    {
        if (result.compareAndSet(null, Result.SUCCESS))
        {
            latch.countDown();
        }
    }

    public void markFailure(Throwable t)
    {
        error.set(t);
        if (result.compareAndSet(null, Result.FAILURE))
        {
            latch.countDown();
        }
    }

    public Result await() throws InterruptedException
    {
        latch.await();
        return result.get();
    }

    public Optional<Result> await(long timeout, TimeUnit unit) throws InterruptedException
    {
        return latch.await(timeout, unit) ? Optional.ofNullable(result.get()) : Optional.empty();
    }

    public Optional<Throwable> getError()
    {
        return Optional.ofNullable(error.get());
    }
}
