org.yx.common.ThreadContext Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of sumk Show documentation
A quick developing framewort for internet company
There is a newer version: 4.0.2
/**
 * Copyright (C) 2016 - 2030 youtongluan.
 *
 * 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.yx.common;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import org.yx.conf.AppInfo;
import org.yx.rpc.Attachable;
import org.yx.rpc.client.Req;
import org.yx.util.StringUtil;

public final class ThreadContext implements Attachable {

	private static final String TEST = "sumk.test";

	public static enum ActionType {
		HTTP, RPC, OTHER
	}

	private final ActionType type;

	private volatile LogContext logContext;

	/**
	 * 用来做自增长的
	 */
	private int spanSeed;

	public boolean isTest() {
		return logContext.test;
	}

	private boolean parseTest(boolean isTest) {
		if (!isTest) {
			return false;
		}
		return AppInfo.getBoolean(TEST, false);
	}

	private ThreadContext(ActionType type, String act, String traceId, String spanId, String userId, boolean isTest,
			Map attachments) {
		this.type = type;
		this.logContext = new LogContext(act, traceId, spanId, userId, this.parseTest(isTest), attachments);
	}

	public ActionType type() {
		return type;
	}

	public String act() {
		return logContext.act;
	}

	public String traceId() {
		return logContext.traceId;
	}

	public String spanId() {
		return logContext.spanId;
	}

	private static final ThreadLocal holder = new ThreadLocal() {

		@Override
		protected ThreadContext initialValue() {

			return new ThreadContext(ActionType.OTHER, null, null, null, null, false, null);
		}

	};

	public static ThreadContext httpContext(String act, String thisIsTest) {
		boolean test = false;
		if (thisIsTest != null && thisIsTest.equals(AppInfo.get(TEST))) {
			test = true;
		}
		ThreadContext c = new ThreadContext(ActionType.HTTP, act, null, null, null, test, null);
		holder.set(c);
		return c;
	}

	public static ThreadContext rpcContext(Req req, boolean isTest) {
		String traceId = StringUtil.isEmpty(req.getTraceId()) ? null : req.getTraceId();
		ThreadContext c = new ThreadContext(ActionType.RPC, req.getApi(), traceId, req.getSpanId(), req.getUserId(),
				isTest, req.getAttachments());
		holder.set(c);
		return c;
	}

	public static ThreadContext get() {
		return holder.get();
	}

	public static void remove() {
		holder.remove();
	}

	public void setTraceIdIfAbsent(String traceId) {
		LogContext lc = this.logContext;
		if (lc.traceId != null) {
			return;
		}
		this.logContext = new LogContext(lc.act, traceId, lc.spanId, lc.userId, lc.test, lc.attachments);
	}

	public String userId() {
		return logContext.userId;
	}

	public void userId(String userId) {
		LogContext lc = this.logContext;
		this.logContext = new LogContext(lc.act, lc.traceId, lc.spanId, userId, lc.test, lc.attachments);
	}

	@Override
	public Map attachmentView() {
		return logContext.attachments;
	}

	@Override
	public void setAttachments(Map attachments) {
		this.logContext = new LogContext(this.logContext, attachments);
	}

	@Override
	public void setAttachment(String key, String value) {
		Map attachments = this.logContext.attachments;
		attachments = attachments == null ? new HashMap<>() : new HashMap<>(attachments);
		attachments.put(key, value);
		this.logContext = new LogContext(this.logContext, attachments);
	}

	@Override
	public String getAttachment(String key) {
		Map attachments = this.logContext.attachments;
		if (attachments == null) {
			return null;
		}
		return attachments.get(key);
	}

	public String nextSpanId() {
		LogContext lc = this.logContext;
		if (lc.traceId == null) {
			return "1";
		}
		int seed;
		synchronized (this) {
			seed = ++this.spanSeed;
		}
		String sp = lc.spanId;
		if (sp == null) {
			return String.valueOf(seed);
		}
		return new StringBuilder().append(lc.spanId).append('.').append(seed).toString();
	}

	public static void recover(ThreadContext context) {
		holder.set(context);
	}

	public static final class LogContext {

		public final String act;

		public final String traceId;

		public final String spanId;

		public final String userId;

		public final boolean test;

		private final Map attachments;

		public LogContext(String act, String traceId, String spanId, String userId, boolean test,
				Map attachments) {
			this.act = act;
			this.traceId = StringUtil.isEmpty(traceId) ? null : traceId;
			this.spanId = spanId;
			this.userId = userId;
			this.test = test;
			this.attachments = attachments;
		}

		public LogContext(LogContext lc, Map attachments) {
			this.act = lc.act;
			this.traceId = lc.traceId;
			this.spanId = lc.spanId;
			this.userId = lc.userId;
			this.test = lc.test;
			this.attachments = attachments == null || attachments.isEmpty() ? null
					: Collections.unmodifiableMap(attachments);
		}

		/**
		 * @return
		 */
		public Map unmodifyAttachs() {
			return this.attachments;
		}
	}

	public LogContext logContext() {
		return this.logContext;
	}
}