runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/logging/DataflowWorkerLoggingInitializer.java - beam - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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.beam.runners.dataflow.worker.logging;

 import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.DEBUG;
 import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.ERROR;
 import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.INFO;
 import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.OFF;
 import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.TRACE;
 import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.WARN;

 import java.io.File;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.List;
 import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.LogManager;
 import java.util.logging.Logger;
 import org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;

 /**
  * Sets up {@link java.util.logging} configuration on the Dataflow worker with a rotating file
  * logger. The file logger uses the {@link DataflowWorkerLoggingHandler} format. A user can override
  * the logging level by customizing the options found within {@link DataflowWorkerLoggingOptions}. A
  * user can override the location by specifying the Java system property
  * "dataflow.worker.logging.basepath" and the file size in MB before rolling over to a new file by
  * specifying the Java system property "dataflow.worker. loggging.filesize_mb". The default log
  * level is INFO, the default location is a file named dataflow-json.log within the system temporary
  * directory and the default file size is 1 GB.
  */
 public class DataflowWorkerLoggingInitializer {
   private static final String ROOT_LOGGER_NAME = "";

   @VisibleForTesting
   static final String DEFAULT_RUNNER_LOGGING_LOCATION =
       new File(System.getProperty("java.io.tmpdir"), "dataflow-json.log").getPath();

   @VisibleForTesting
   static final String DEFAULT_SDK_LOGGING_LOCATION =
       new File(System.getProperty("java.io.tmpdir"), "sdk-json.log").getPath();

   @VisibleForTesting
   public static final String RUNNER_FILEPATH_PROPERTY = "dataflow.worker.logging.filepath";

   @VisibleForTesting
   static final String SDK_FILEPATH_PROPERTY = "dataflow.worker.logging.sdkfilepath";

   private static final String FILESIZE_MB_PROPERTY = "dataflow.worker.logging.filesize_mb";

   private static final String SYSTEM_OUT_LOG_NAME = "System.out";
   private static final String SYSTEM_ERR_LOG_NAME = "System.err";

   static final ImmutableBiMap<Level, DataflowWorkerLoggingOptions.Level> LEVELS =
       ImmutableBiMap.<Level, DataflowWorkerLoggingOptions.Level>builder()
           .put(Level.OFF, OFF)
           .put(Level.SEVERE, ERROR)
           .put(Level.WARNING, WARN)
           .put(Level.INFO, INFO)
           .put(Level.FINE, DEBUG)
           .put(Level.FINEST, TRACE)
           .build();

   /**
    * This default log level is overridden by the log level found at {@code
    * DataflowWorkerLoggingOptions#getDefaultWorkerLogLevel()}.
    */
   private static final DataflowWorkerLoggingOptions.Level DEFAULT_LOG_LEVEL =
       LEVELS.get(Level.INFO);

   /* We need to store a reference to the configured loggers so that they are not
    * garbage collected. java.util.logging only has weak references to the loggers
    * so if they are garbage collection, our hierarchical configuration will be lost. */
   @SuppressWarnings("MismatchedQueryAndUpdateOfCollection")
   private static List<Logger> configuredLoggers = Lists.newArrayList();

   private static DataflowWorkerLoggingHandler loggingHandler;
   private static DataflowWorkerLoggingHandler sdkLoggingHandler;
   private static PrintStream originalStdOut;
   private static PrintStream originalStdErr = System.err;
   private static boolean initialized = false;

   private static DataflowWorkerLoggingHandler makeLoggingHandler(
       String filepathProperty, String defaultFilePath) throws IOException {
     String filepath = System.getProperty(filepathProperty, defaultFilePath);
     int filesizeMb = Integer.parseInt(System.getProperty(FILESIZE_MB_PROPERTY, "1024"));
     DataflowWorkerLoggingHandler handler =
         new DataflowWorkerLoggingHandler(filepath, filesizeMb * 1024L * 1024L);
     handler.setLevel(Level.ALL);
     return handler;
   }

   public static PrintStream getOriginalStdErr() {
     return originalStdErr;
   }

   /** Sets up the initial logging configuration. */
   public static synchronized void initialize() {
     if (initialized) {
       return;
     }

     try {
       loggingHandler =
           makeLoggingHandler(RUNNER_FILEPATH_PROPERTY, DEFAULT_RUNNER_LOGGING_LOCATION);
       Logger rootLogger = LogManager.getLogManager().getLogger(ROOT_LOGGER_NAME);

       Level logLevel = getJulLevel(DEFAULT_LOG_LEVEL);
       rootLogger.setLevel(logLevel);
       rootLogger.addHandler(loggingHandler);

       originalStdOut = System.out;
       originalStdErr = System.err;
       System.setOut(
           JulHandlerPrintStreamAdapterFactory.create(
               loggingHandler, SYSTEM_OUT_LOG_NAME, Level.INFO));
       System.setErr(
           JulHandlerPrintStreamAdapterFactory.create(
               loggingHandler, SYSTEM_ERR_LOG_NAME, Level.SEVERE));

       // Initialize the SDK Logging Handler, which will only be used for the LoggingService
       sdkLoggingHandler = makeLoggingHandler(SDK_FILEPATH_PROPERTY, DEFAULT_SDK_LOGGING_LOCATION);

       initialized = true;
     } catch (SecurityException | IOException | NumberFormatException e) {
       throw new ExceptionInInitializerError(e);
     }
   }

   /** Reconfigures logging with the passed in options. */
   public static synchronized void configure(DataflowWorkerLoggingOptions options) {
     if (!initialized) {
       throw new RuntimeException("configure() called before initialize()");
     }
     if (options.getDefaultWorkerLogLevel() != null) {
       Level defaultLevel = getJulLevel(options.getDefaultWorkerLogLevel());
       LogManager.getLogManager().getLogger(ROOT_LOGGER_NAME).setLevel(defaultLevel);
     }

     if (options.getWorkerLogLevelOverrides() != null) {
       for (Map.Entry<String, DataflowWorkerLoggingOptions.Level> loggerOverride :
           options.getWorkerLogLevelOverrides().entrySet()) {
         Logger logger = Logger.getLogger(loggerOverride.getKey());
         logger.setLevel(getJulLevel(loggerOverride.getValue()));
         configuredLoggers.add(logger);
       }
     }

     // If the options specify a level for messages logged to System.out/err, we need to reconfigure
     // the corresponding stream adapter.
     if (options.getWorkerSystemOutMessageLevel() != null) {
       System.out.close();
       System.setOut(
           JulHandlerPrintStreamAdapterFactory.create(
               loggingHandler,
               SYSTEM_OUT_LOG_NAME,
               getJulLevel(options.getWorkerSystemOutMessageLevel())));
     }

     if (options.getWorkerSystemErrMessageLevel() != null) {
       System.err.close();
       System.setErr(
           JulHandlerPrintStreamAdapterFactory.create(
               loggingHandler,
               SYSTEM_ERR_LOG_NAME,
               getJulLevel(options.getWorkerSystemErrMessageLevel())));
     }
   }

   /**
    * Returns the underlying {@link DataflowWorkerLoggingHandler}.
    *
    * <p>Generally, code should just use logging interface.
    */
   public static DataflowWorkerLoggingHandler getLoggingHandler() {
     if (!initialized) {
       throw new RuntimeException("getLoggingHandler() called before initialize()");
     }
     return loggingHandler;
   }

   /**
    * Returns the underlying {@link DataflowWorkerLoggingHandler} for logs from the SDK.
    *
    * <p>Initializes Dataflow worker logging if not initialized already.
    */
   public static DataflowWorkerLoggingHandler getSdkLoggingHandler() {
     if (!initialized) {
       throw new RuntimeException("getSdkLoggingHandler() called before initialize()");
     }
     return sdkLoggingHandler;
   }

   private static Level getJulLevel(DataflowWorkerLoggingOptions.Level level) {
     return LEVELS.inverse().get(level);
   }

   @VisibleForTesting
   public static synchronized void reset() {
     if (!initialized) {
       return;
     }
     configuredLoggers = Lists.newArrayList();
     System.setOut(originalStdOut);
     System.setErr(originalStdErr);
     JulHandlerPrintStreamAdapterFactory.reset();
     initialized = false;
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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.beam.runners.dataflow.worker.logging;

	import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.DEBUG;
	import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.ERROR;
	import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.INFO;
	import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.OFF;
	import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.TRACE;
	import static org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions.Level.WARN;

	import java.io.File;
	import java.io.IOException;
	import java.io.PrintStream;
	import java.util.List;
	import java.util.Map;
	import java.util.logging.Level;
	import java.util.logging.LogManager;
	import java.util.logging.Logger;
	import org.apache.beam.runners.dataflow.options.DataflowWorkerLoggingOptions;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;

	/**
	* Sets up {@link java.util.logging} configuration on the Dataflow worker with a rotating file
	* logger. The file logger uses the {@link DataflowWorkerLoggingHandler} format. A user can override
	* the logging level by customizing the options found within {@link DataflowWorkerLoggingOptions}. A
	* user can override the location by specifying the Java system property
	* "dataflow.worker.logging.basepath" and the file size in MB before rolling over to a new file by
	* specifying the Java system property "dataflow.worker. loggging.filesize_mb". The default log
	* level is INFO, the default location is a file named dataflow-json.log within the system temporary
	* directory and the default file size is 1 GB.
	*/
	public class DataflowWorkerLoggingInitializer {
	private static final String ROOT_LOGGER_NAME = "";

	@VisibleForTesting
	static final String DEFAULT_RUNNER_LOGGING_LOCATION =
	new File(System.getProperty("java.io.tmpdir"), "dataflow-json.log").getPath();

	@VisibleForTesting
	static final String DEFAULT_SDK_LOGGING_LOCATION =
	new File(System.getProperty("java.io.tmpdir"), "sdk-json.log").getPath();

	@VisibleForTesting
	public static final String RUNNER_FILEPATH_PROPERTY = "dataflow.worker.logging.filepath";

	@VisibleForTesting
	static final String SDK_FILEPATH_PROPERTY = "dataflow.worker.logging.sdkfilepath";

	private static final String FILESIZE_MB_PROPERTY = "dataflow.worker.logging.filesize_mb";

	private static final String SYSTEM_OUT_LOG_NAME = "System.out";
	private static final String SYSTEM_ERR_LOG_NAME = "System.err";

	static final ImmutableBiMap<Level, DataflowWorkerLoggingOptions.Level> LEVELS =
	ImmutableBiMap.<Level, DataflowWorkerLoggingOptions.Level>builder()
	.put(Level.OFF, OFF)
	.put(Level.SEVERE, ERROR)
	.put(Level.WARNING, WARN)
	.put(Level.INFO, INFO)
	.put(Level.FINE, DEBUG)
	.put(Level.FINEST, TRACE)
	.build();

	/**
	* This default log level is overridden by the log level found at {@code
	* DataflowWorkerLoggingOptions#getDefaultWorkerLogLevel()}.
	*/
	private static final DataflowWorkerLoggingOptions.Level DEFAULT_LOG_LEVEL =
	LEVELS.get(Level.INFO);

	/* We need to store a reference to the configured loggers so that they are not
	* garbage collected. java.util.logging only has weak references to the loggers
	* so if they are garbage collection, our hierarchical configuration will be lost. */
	@SuppressWarnings("MismatchedQueryAndUpdateOfCollection")
	private static List<Logger> configuredLoggers = Lists.newArrayList();

	private static DataflowWorkerLoggingHandler loggingHandler;
	private static DataflowWorkerLoggingHandler sdkLoggingHandler;
	private static PrintStream originalStdOut;
	private static PrintStream originalStdErr = System.err;
	private static boolean initialized = false;

	private static DataflowWorkerLoggingHandler makeLoggingHandler(
	String filepathProperty, String defaultFilePath) throws IOException {
	String filepath = System.getProperty(filepathProperty, defaultFilePath);
	int filesizeMb = Integer.parseInt(System.getProperty(FILESIZE_MB_PROPERTY, "1024"));
	DataflowWorkerLoggingHandler handler =
	new DataflowWorkerLoggingHandler(filepath, filesizeMb * 1024L * 1024L);
	handler.setLevel(Level.ALL);
	return handler;
	}

	public static PrintStream getOriginalStdErr() {
	return originalStdErr;
	}

	/** Sets up the initial logging configuration. */
	public static synchronized void initialize() {
	if (initialized) {
	return;
	}

	try {
	loggingHandler =
	makeLoggingHandler(RUNNER_FILEPATH_PROPERTY, DEFAULT_RUNNER_LOGGING_LOCATION);
	Logger rootLogger = LogManager.getLogManager().getLogger(ROOT_LOGGER_NAME);

	Level logLevel = getJulLevel(DEFAULT_LOG_LEVEL);
	rootLogger.setLevel(logLevel);
	rootLogger.addHandler(loggingHandler);

	originalStdOut = System.out;
	originalStdErr = System.err;
	System.setOut(
	JulHandlerPrintStreamAdapterFactory.create(
	loggingHandler, SYSTEM_OUT_LOG_NAME, Level.INFO));
	System.setErr(
	JulHandlerPrintStreamAdapterFactory.create(
	loggingHandler, SYSTEM_ERR_LOG_NAME, Level.SEVERE));

	// Initialize the SDK Logging Handler, which will only be used for the LoggingService
	sdkLoggingHandler = makeLoggingHandler(SDK_FILEPATH_PROPERTY, DEFAULT_SDK_LOGGING_LOCATION);

	initialized = true;
	} catch (SecurityException \| IOException \| NumberFormatException e) {
	throw new ExceptionInInitializerError(e);
	}
	}

	/** Reconfigures logging with the passed in options. */
	public static synchronized void configure(DataflowWorkerLoggingOptions options) {
	if (!initialized) {
	throw new RuntimeException("configure() called before initialize()");
	}
	if (options.getDefaultWorkerLogLevel() != null) {
	Level defaultLevel = getJulLevel(options.getDefaultWorkerLogLevel());
	LogManager.getLogManager().getLogger(ROOT_LOGGER_NAME).setLevel(defaultLevel);
	}

	if (options.getWorkerLogLevelOverrides() != null) {
	for (Map.Entry<String, DataflowWorkerLoggingOptions.Level> loggerOverride :
	options.getWorkerLogLevelOverrides().entrySet()) {
	Logger logger = Logger.getLogger(loggerOverride.getKey());
	logger.setLevel(getJulLevel(loggerOverride.getValue()));
	configuredLoggers.add(logger);
	}
	}

	// If the options specify a level for messages logged to System.out/err, we need to reconfigure
	// the corresponding stream adapter.
	if (options.getWorkerSystemOutMessageLevel() != null) {
	System.out.close();
	System.setOut(
	JulHandlerPrintStreamAdapterFactory.create(
	loggingHandler,
	SYSTEM_OUT_LOG_NAME,
	getJulLevel(options.getWorkerSystemOutMessageLevel())));
	}

	if (options.getWorkerSystemErrMessageLevel() != null) {
	System.err.close();
	System.setErr(
	JulHandlerPrintStreamAdapterFactory.create(
	loggingHandler,
	SYSTEM_ERR_LOG_NAME,
	getJulLevel(options.getWorkerSystemErrMessageLevel())));
	}
	}

	/**
	* Returns the underlying {@link DataflowWorkerLoggingHandler}.
	*
	* <p>Generally, code should just use logging interface.
	*/
	public static DataflowWorkerLoggingHandler getLoggingHandler() {
	if (!initialized) {
	throw new RuntimeException("getLoggingHandler() called before initialize()");
	}
	return loggingHandler;
	}

	/**
	* Returns the underlying {@link DataflowWorkerLoggingHandler} for logs from the SDK.
	*
	* <p>Initializes Dataflow worker logging if not initialized already.
	*/
	public static DataflowWorkerLoggingHandler getSdkLoggingHandler() {
	if (!initialized) {
	throw new RuntimeException("getSdkLoggingHandler() called before initialize()");
	}
	return sdkLoggingHandler;
	}

	private static Level getJulLevel(DataflowWorkerLoggingOptions.Level level) {
	return LEVELS.inverse().get(level);
	}

	@VisibleForTesting
	public static synchronized void reset() {
	if (!initialized) {
	return;
	}
	configuredLoggers = Lists.newArrayList();
	System.setOut(originalStdOut);
	System.setErr(originalStdErr);
	JulHandlerPrintStreamAdapterFactory.reset();
	initialized = false;
	}
	}