org.apache.flink.util.ComponentClosingUtils Maven / Gradle / Ivy

Go to download
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.util;

import org.apache.flink.util.function.ThrowingRunnable;

import java.util.Optional;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;

/**
 * A util class to help with a clean component shutdown.
 */
public class ComponentClosingUtils {

	private ComponentClosingUtils() {}

	/**
	 * Close a component with a timeout.
	 *
	 * @param componentName the name of the component.
	 * @param closingSequence the closing logic which is a callable that can throw exceptions.
	 * @param closeTimeoutMs the timeout in milliseconds to waif for the component to close.
	 * @return An optional throwable which is non-empty if an error occurred when closing the component.
	 */
	public static Optional closeWithTimeout(
			String componentName,
			ThrowingRunnable closingSequence,
			long closeTimeoutMs) {
		return closeWithTimeout(
				componentName,
				(Runnable) () -> {
					try {
						closingSequence.run();
					} catch (Exception e) {
						throw new ClosingException(componentName, e);
					}
				}, closeTimeoutMs);
	}

	/**
	 * Close a component with a timeout.
	 *
	 * @param componentName the name of the component.
	 * @param closingSequence the closing logic.
	 * @param closeTimeoutMs the timeout in milliseconds to waif for the component to close.
	 * @return An optional throwable which is non-empty if an error occurred when closing the component.
	 */
	public static Optional closeWithTimeout(
			String componentName,
			Runnable closingSequence,
			long closeTimeoutMs) {
		AtomicReference closingError = new AtomicReference<>();
		Thread t = new Thread(closingSequence);
		t.setUncaughtExceptionHandler((thread, error) -> closingError.set(error));
		t.start();
		try {
			t.join(closeTimeoutMs);
		} catch (InterruptedException e) {
			// Only add the interrupted exception to suppressed if it is the first error.
			// Otherwise, either a user error or a TimeoutException would have been added.
			reportError(closingError, e);
		}
		if (t.isAlive()) {
			reportError(closingError, new TimeoutException(
					String.format("Failed to close the %s before timeout of %d milliseconds",
							componentName, closeTimeoutMs)));
			t.interrupt();
		}
		return Optional.ofNullable(closingError.get());
	}

	// ---------------------------

	private static void reportError(AtomicReference errorRef, Throwable toReport) {
		if (!errorRef.compareAndSet(null, toReport)) {
			errorRef.get().addSuppressed(toReport);
		}
	}

	private static class ClosingException extends RuntimeException {
		private static final long serialVersionUID = 2527474477287706295L;

		private ClosingException(String componentName, Exception e) {
			super(String.format("Caught exception when closing %s", componentName), e);
		}
	}
}