    001/*
    002 * Copyright (c) 2007-2017 Xplenty, Inc. All Rights Reserved.
    003 *
    004 * Project and contact information: http://www.cascading.org/
    005 *
    006 * This file is part of the Cascading project.
    007 *
    008 * Licensed under the Apache License, Version 2.0 (the "License");
    009 * you may not use this file except in compliance with the License.
    010 * You may obtain a copy of the License at
    011 *
    012 *     http://www.apache.org/licenses/LICENSE-2.0
    013 *
    014 * Unless required by applicable law or agreed to in writing, software
    015 * distributed under the License is distributed on an "AS IS" BASIS,
    016 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    017 * See the License for the specific language governing permissions and
    018 * limitations under the License.
    019 */
    020
    021package cascading.util;
    022
    023import java.util.Comparator;
    024import java.util.Queue;
    025import java.util.concurrent.PriorityBlockingQueue;
    026
    027import org.slf4j.Logger;
    028import org.slf4j.LoggerFactory;
    029
    030/**
    031 * ShutdownUtil is a private helper class for registering dependent shutdown hooks to maintain internal state
    032 * information reliably when a jvm is shutting down.
    033 * <p/>
    034 * This api is not intended for public use.
    035 */
    036public class ShutdownUtil
    037  {
    038  /**
    039   * System property set to true when we have entered shutdown
    040   */
    041  public static final String SHUTDOWN_EXECUTING = "cascading.jvm.shutdown.executing";
    042  public static final String SHUTDOWN_FORCE_NON_DAEMON = "cascading.jvm.shutdown.non-daemon";
    043
    044  private static final Logger LOG = LoggerFactory.getLogger( ShutdownUtil.class );
    045
    046  public static abstract class Hook
    047    {
    048    public enum Priority
    049      {
    050        FIRST, WORK_PARENT, WORK_CHILD, SERVICE_CONSUMER, SERVICE_PROVIDER, LAST
    051      }
    052
    053    public abstract Priority priority();
    054
    055    public abstract void execute();
    056    }
    057
    058  private static Queue<Hook> queue = new PriorityBlockingQueue<>( 20, new Comparator<Hook>()
    059  {
    060  @Override
    061  public int compare( Hook lhs, Hook rhs )
    062    {
    063    if( lhs == rhs )
    064      return 0;
    065
    066    if( lhs == null )
    067      return -1;
    068    else if( rhs == null )
    069      return 1;
    070
    071    return lhs.priority().compareTo( rhs.priority() );
    072    }
    073  }
    074  );
    075
    076  private static Thread shutdownHook;
    077
    078  public static void addHook( Hook hook )
    079    {
    080    if( hook == null )
    081      throw new IllegalArgumentException( "hook may not be null" );
    082
    083    registerShutdownHook();
    084
    085    queue.add( hook );
    086    }
    087
    088  public static boolean removeHook( Hook hook )
    089    {
    090    return queue.remove( hook );
    091    }
    092
    093  public static synchronized void registerShutdownHook()
    094    {
    095    if( shutdownHook != null )
    096      return;
    097
    098    final boolean isForceNonDaemon = Boolean.getBoolean( SHUTDOWN_FORCE_NON_DAEMON );
    099
    100    shutdownHook = new Thread( "cascading shutdown hooks" )
    101    {
    102
    103    {
    104    if( isForceNonDaemon )
    105      this.setDaemon( false );
    106    }
    107
    108    @Override
    109    public void run()
    110      {
    111      System.setProperty( SHUTDOWN_EXECUTING, "true" );
    112
    113      try
    114        {
    115        // These are not threads, so each one will be run in priority order
    116        // blocking until the previous is complete
    117        while( !queue.isEmpty() )
    118          {
    119          Hook hook = null;
    120
    121          try
    122            {
    123            hook = queue.poll();
    124
    125            // may get removed while shutdown is executing
    126            if( hook == null )
    127              continue;
    128
    129            hook.execute();
    130            }
    131          catch( Exception exception )
    132            {
    133            LOG.error( "failed executing hook: {}, with exception: {}", hook, exception.getMessage() );
    134            LOG.debug( "with exception trace", exception );
    135            }
    136          }
    137        }
    138      finally
    139        {
    140        System.setProperty( SHUTDOWN_EXECUTING, "false" );
    141        }
    142      }
    143    };
    144
    145    Runtime.getRuntime().addShutdownHook( shutdownHook );
    146    }
    147
    148  public static void deregisterShutdownHook()
    149    {
    150    Runtime.getRuntime().removeShutdownHook( shutdownHook );
    151    }
    152  }