package org.apache.logging.log4j.core.util;
import java.lang.ref.Reference;
import java.lang.ref.SoftReference;
import java.lang.ref.WeakReference;
import java.util.Collection;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.logging.log4j.Logger;
import org.apache.logging.log4j.core.AbstractLifeCycle;
import org.apache.logging.log4j.core.LifeCycle2;
import org.apache.logging.log4j.status.StatusLogger;
public class DefaultShutdownCallbackRegistry implements ShutdownCallbackRegistry, LifeCycle2, Runnable {
protected static final Logger LOGGER = StatusLogger.getLogger();
private final AtomicReference<State> state = new AtomicReference<>(State.INITIALIZED);
private final ThreadFactory threadFactory;
private final Collection<Cancellable> hooks = new CopyOnWriteArrayList<>();
private Reference<Thread> shutdownHookRef;
public DefaultShutdownCallbackRegistry() {
this(Executors.defaultThreadFactory());
}
protected DefaultShutdownCallbackRegistry(final ThreadFactory threadFactory) {
this.threadFactory = threadFactory;
}
@Override
public void run() {
if (state.compareAndSet(State.STARTED, State.STOPPING)) {
for (final Runnable hook : hooks) {
try {
hook.run();
} catch (final Throwable t1) {
try {
LOGGER.error(SHUTDOWN_HOOK_MARKER, "Caught exception executing shutdown hook {}", hook, t1);
} catch (final Throwable t2) {
System.err.println("Caught exception " + t2.getClass() + " logging exception " + t1.getClass());
t1.printStackTrace();
}
}
}
state.set(State.STOPPED);
}
}
private static class RegisteredCancellable implements Cancellable {
private final Reference<Runnable> hook;
private Collection<Cancellable> registered;
RegisteredCancellable(final Runnable callback, final Collection<Cancellable> registered) {
this.registered = registered;
hook = new SoftReference<>(callback);
}
@Override
public void cancel() {
hook.clear();
registered.remove(this);
registered = null;
}
@Override
public void run() {
final Runnable runnableHook = this.hook.get();
if (runnableHook != null) {
runnableHook.run();
this.hook.clear();
}
}
@Override
public String toString() {
return String.valueOf(hook.get());
}
}
@Override
public Cancellable addShutdownCallback(final Runnable callback) {
if (isStarted()) {
final Cancellable receipt = new RegisteredCancellable(callback, hooks);
hooks.add(receipt);
return receipt;
}
throw new IllegalStateException("Cannot add new shutdown hook as this is not started. Current state: " +
state.get().name());
}
@Override
public void initialize() {
}
@Override
public void start() {
if (state.compareAndSet(State.INITIALIZED, State.STARTING)) {
try {
addShutdownHook(threadFactory.newThread(this));
state.set(State.STARTED);
} catch (final IllegalStateException ex) {
state.set(State.STOPPED);
throw ex;
} catch (final Exception e) {
LOGGER.catching(e);
state.set(State.STOPPED);
}
}
}
private void addShutdownHook(final Thread thread) {
shutdownHookRef = new WeakReference<>(thread);
Runtime.getRuntime().addShutdownHook(thread);
}
@Override
public void stop() {
stop(AbstractLifeCycle.DEFAULT_STOP_TIMEOUT, AbstractLifeCycle.DEFAULT_STOP_TIMEUNIT);
}
@Override
public boolean stop(final long timeout, final TimeUnit timeUnit) {
if (state.compareAndSet(State.STARTED, State.STOPPING)) {
try {
removeShutdownHook();
} finally {
state.set(State.STOPPED);
}
}
return true;
}
private void removeShutdownHook() {
final Thread shutdownThread = shutdownHookRef.get();
if (shutdownThread != null) {
Runtime.getRuntime().removeShutdownHook(shutdownThread);
shutdownHookRef.enqueue();
}
}
@Override
public State getState() {
return state.get();
}
@Override
public boolean isStarted() {
return state.get() == State.STARTED;
}
@Override
public boolean isStopped() {
return state.get() == State.STOPPED;
}
}