    001    /*
    002     * Licensed to the Apache Software Foundation (ASF) under one or more
    003     * contributor license agreements. See the NOTICE file distributed with
    004     * this work for additional information regarding copyright ownership.
    005     * The ASF licenses this file to You under the Apache license, Version 2.0
    006     * (the "License"); you may not use this file except in compliance with
    007     * the License. You may obtain a copy of the License at
    008     *
    009     *      http://www.apache.org/licenses/LICENSE-2.0
    010     *
    011     * Unless required by applicable law or agreed to in writing, software
    012     * distributed under the License is distributed on an "AS IS" BASIS,
    013     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    014     * See the license for the specific language governing permissions and
    015     * limitations under the license.
    016     */
    017
    018    package org.apache.logging.log4j.core.util;
    019
    020    import java.io.Serializable;
    021    import java.lang.ref.Reference;
    022    import java.lang.ref.SoftReference;
    023    import java.lang.ref.WeakReference;
    024    import java.util.Collection;
    025    import java.util.concurrent.CopyOnWriteArrayList;
    026    import java.util.concurrent.Executors;
    027    import java.util.concurrent.ThreadFactory;
    028    import java.util.concurrent.atomic.AtomicReference;
    029
    030    import org.apache.logging.log4j.Logger;
    031    import org.apache.logging.log4j.core.LifeCycle;
    032    import org.apache.logging.log4j.core.LifeCycle.State;
    033    import org.apache.logging.log4j.status.StatusLogger;
    034
    035    /**
    036     * ShutdownRegistrationStrategy that simply uses {@link Runtime#addShutdownHook(Thread)}. If no strategy is specified,
    037     * this one is used for shutdown hook registration.
    038     *
    039     * @since 2.1
    040     */
    041    public class DefaultShutdownCallbackRegistry implements ShutdownCallbackRegistry, LifeCycle, Runnable, Serializable {
    042
    043        private static final long serialVersionUID = 1L;
    044        protected static final Logger LOGGER = StatusLogger.getLogger();
    045
    046        private final AtomicReference<State> state = new AtomicReference<State>(State.INITIALIZED);
    047        private final ThreadFactory threadFactory;
    048        private final Collection<Cancellable> hooks = new CopyOnWriteArrayList<Cancellable>();
    049        private Reference<Thread> shutdownHookRef;
    050
    051        /**
    052         * Constructs a DefaultShutdownRegistrationStrategy.
    053         */
    054        public DefaultShutdownCallbackRegistry() {
    055            this(Executors.defaultThreadFactory());
    056        }
    057
    058        /**
    059         * Constructs a DefaultShutdownRegistrationStrategy using the given {@link ThreadFactory}.
    060         *
    061         * @param threadFactory the ThreadFactory to use to create a {@link Runtime} shutdown hook thread
    062         */
    063        protected DefaultShutdownCallbackRegistry(final ThreadFactory threadFactory) {
    064            this.threadFactory = threadFactory;
    065        }
    066
    067        /**
    068         * Executes the registered shutdown callbacks.
    069         */
    070        @Override
    071        public void run() {
    072            if (state.compareAndSet(State.STARTED, State.STOPPING)) {
    073                for (final Runnable hook : hooks) {
    074                    try {
    075                        hook.run();
    076                    } catch (final Throwable t) {
    077                        LOGGER.error(SHUTDOWN_HOOK_MARKER, "Caught exception executing shutdown hook {}", hook, t);
    078                    }
    079                }
    080                state.set(State.STOPPED);
    081            }
    082        }
    083
    084        @Override
    085        public Cancellable addShutdownCallback(final Runnable callback) {
    086            if (isStarted()) {
    087                final Cancellable receipt = new Cancellable() {
    088                    // use a reference to prevent memory leaks
    089                    private final Reference<Runnable> hook = new SoftReference<Runnable>(callback);
    090
    091                    @Override
    092                    public void cancel() {
    093                        hook.clear();
    094                        hooks.remove(this);
    095                    }
    096
    097                    @Override
    098                    public void run() {
    099                        final Runnable hook = this.hook.get();
    100                        if (hook != null) {
    101                            hook.run();
    102                            this.hook.clear();
    103                        }
    104                    }
    105
    106                    @Override
    107                    public String toString() {
    108                        return String.valueOf(hook.get());
    109                    }
    110                };
    111                hooks.add(receipt);
    112                return receipt;
    113            }
    114            throw new IllegalStateException("Cannot add new shutdown hook as this is not started. Current state: " +
    115                state.get().name());
    116        }
    117
    118        /**
    119         * Registers the shutdown thread only if this is initialized.
    120         */
    121        @Override
    122        public void start() {
    123            if (state.compareAndSet(State.INITIALIZED, State.STARTING)) {
    124                try {
    125                    addShutdownHook(threadFactory.newThread(this));
    126                    state.set(State.STARTED);
    127                } catch (final Exception e) {
    128                    LOGGER.catching(e);
    129                    state.set(State.STOPPED);
    130                }
    131            }
    132        }
    133
    134        private void addShutdownHook(final Thread thread) {
    135            shutdownHookRef = new WeakReference<Thread>(thread);
    136            Runtime.getRuntime().addShutdownHook(thread);
    137        }
    138
    139        /**
    140         * Cancels the shutdown thread only if this is started.
    141         */
    142        @Override
    143        public void stop() {
    144            if (state.compareAndSet(State.STARTED, State.STOPPING)) {
    145                try {
    146                    removeShutdownHook();
    147                } finally {
    148                    state.set(State.STOPPED);
    149                }
    150            }
    151        }
    152
    153        private void removeShutdownHook() {
    154            final Thread shutdownThread = shutdownHookRef.get();
    155            if (shutdownThread != null) {
    156                Runtime.getRuntime().removeShutdownHook(shutdownThread);
    157                shutdownHookRef.enqueue();
    158            }
    159        }
    160
    161        @Override
    162        public State getState() {
    163            return state.get();
    164        }
    165
    166        /**
    167         * Indicates if this can accept shutdown hooks.
    168         *
    169         * @return true if this can accept shutdown hooks
    170         */
    171        @Override
    172        public boolean isStarted() {
    173            return state.get() == State.STARTED;
    174        }
    175
    176        @Override
    177        public boolean isStopped() {
    178            return state.get() == State.STOPPED;
    179        }
    180
    181    }