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    }