/*
 *  Copyright 2018 original author or authors.
 *
 *  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.springframework.cloud.gcp.autoconfigure.logging;

import com.google.cloud.logging.LogEntry;
import com.google.cloud.logging.LoggingEnhancer;

import org.springframework.cloud.gcp.core.DefaultGcpProjectIdProvider;
import org.springframework.cloud.gcp.core.GcpProjectIdProvider;

/**
 * Adds the trace ID to the logging entry, in its correct format to be displayed in the Logs viewer.
 *
 * @author João André Martins
 */
public class TraceIdLoggingEnhancer implements LoggingEnhancer {

	private static final ThreadLocal<String> traceId = new ThreadLocal<>();

	private GcpProjectIdProvider projectIdProvider = new DefaultGcpProjectIdProvider();

	public static void setCurrentTraceId(String id) {
		if (id == null) {
			traceId.remove();
		}
		else {
			traceId.set(id);
		}
	}

	/**
	 * @return the trace ID stored through {@link #setCurrentTraceId(String)}
	 */
	public static String getCurrentTraceId() {
		return traceId.get();
	}

	/**
	 * Set the trace field of the log entry to the current trace ID.
	 * <p>The current trace ID is either the trace ID stored in the Mapped Diagnostic Context (MDC)
	 * under the "X-B3-TraceId" key or, if none set, the current trace ID set by
	 * {@link #setCurrentTraceId(String)}.
	 * <p>The trace ID is set in the log entry in the "projects/[GCP_PROJECT_ID]/traces/[TRACE_ID]"
	 * format, in order to be associated to traces by the Google Cloud Console.
	 * @param builder log entry builder
	 */
	@Override
	public void enhanceLogEntry(LogEntry.Builder builder) {
		// In order not to duplicate the whole google-cloud-logging-logback LoggingAppender to add
		// the trace ID from the MDC there, we're doing it here.
		// This requires a call to the org.slf4j package.
		String traceId = org.slf4j.MDC.get(StackdriverTraceConstants.MDC_FIELD_TRACE_ID);
		if (traceId == null) {
			traceId = getCurrentTraceId();
		}

		if (traceId != null) {
			builder.setTrace(StackdriverTraceConstants.composeFullTraceName(
					this.projectIdProvider.getProjectId(), traceId));
		}
	}
}

@@ -1201,82 +1201,251 @@
 ate 
-GcpProjectIdProvider projectIdProvider = new DefaultGcpProjectIdProvider()
+static final String APP_ENGINE_LABEL_NAME = %22appengine.googleapis.com/trace_id%22;%0A%0A%09private GcpProjectIdProvider projectIdProvider = new DefaultGcpProjectIdProvider();%0A%0A%09private boolean runningOnAppEngine = System.getenv(%22GAE_INSTANCE%22) != null
 ;%0A%0A%09
@@ -2193,16 +2193,215 @@
 onsole.%0A
+%09 * %3Cp%3EIf an application is running on Google App Engine, the trace ID is also stored in the%0A%09 * %22appengine.googleapis.com/trace_id%22 field, in order to enable log correlation on the logs%0A%09 * viewer.%0A
 %09 * @par
@@ -2947,16 +2947,16 @@
 ceName(%0A
-
 %09%09%09%09%09thi
@@ -3002,16 +3002,109 @@
 ceId));%0A
+%09%09%09if (this.runningOnAppEngine) %7B%0A%09%09%09%09builder.addLabel(APP_ENGINE_LABEL_NAME, traceId);%0A%09%09%09%7D%0A
 %09%09%7D%0A%09%7D%0A%7D

Set trace ID on GAE label (#740)