fi.evolver.basics.spring.job.TaskStatusService Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of evolver-basics-spring
There is a newer version: 6.5.1
package fi.evolver.basics.spring.job;

import java.lang.management.ManagementFactory;
import java.lang.management.ThreadMXBean;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;

import fi.evolver.basics.spring.job.JobStatusService.Job;
import fi.evolver.basics.spring.job.entity.TaskStatus;
import fi.evolver.basics.spring.job.entity.TaskStatus.TaskState;
import fi.evolver.basics.spring.job.entity.TaskStatusMetadata;
import fi.evolver.basics.spring.util.MessageChainUtils;


@Service
public class TaskStatusService {
	private static final Logger LOG = LoggerFactory.getLogger(TaskStatusService.class);

	private static final ResultState DEFAULT_STATE = new ResultState(TaskState.FAILED, "Result state not set");

	private static final ThreadLocal taskStatusHolder = new ThreadLocal<>();


	private final JobStatusService jobStatusService;
	private final TaskStatusRepository taskStatusRepository;
	private final TaskStatusMetadataRepository taskStatusMetadataRepository;


	@Value("${git.commit.id.describe:?}")
	private String gitDescription;


	@Autowired
	public TaskStatusService(
			JobStatusService jobStatusService,
			TaskStatusRepository taskStatusRepository,
			TaskStatusMetadataRepository taskStatusMetadataRepository) {

		this.jobStatusService = jobStatusService;
		this.taskStatusRepository = taskStatusRepository;
		this.taskStatusMetadataRepository = taskStatusMetadataRepository;
	}


	/**
	 * Associate the current thread with the given task.
	 *
	 * @param taskStatus The task status to push.
	 * @return An auto-closer for finishing the task.
	 */
	public Task couple(TaskStatus taskStatus) {
		TaskStatus oldTask = taskStatusHolder.get();
		if (oldTask != null) {
			Task closer = new Task(oldTask, null);
			closer.setResultState(new ResultState(TaskState.FAILED, "Not closed correctly"));
			finish(closer);
		}
		taskStatusHolder.set(taskStatus);
		if (!taskStatus.getMetadata("MessageChainId").isPresent())
			addMetadata("MessageChainId", MessageChainUtils.getMessageChainId());
		return new Task(taskStatus, jobStatusService.start(taskStatus.getGroup().replaceAll(".*\\.", "")));
	}


	public Optional getTaskStatus() {
		return Optional.ofNullable(taskStatusHolder.get());
	}


	/**
	 * Decouple the current task from the current thread.
	 *
	 * @return The task if one is associated with the thread.
	 */
	public static Optional decouple() {
		Optional result = Optional.ofNullable(taskStatusHolder.get());
		if (result.isPresent())
			taskStatusHolder.remove();
		return result;
	}



	/**
	 * Creates a new task status but doesn't persist it nor associate it with the current thread.
	 * Should only be used in special situations.
	 *
	 * @param clazz The task class.
	 * @param description Description of the task.
	 * @param metadata Metadata for the task;
	 * @return The created task;
	 */
	public TaskStatus create(Class clazz, String description, Map metadata) {
		return new TaskStatus(
				clazz.getCanonicalName(),
				description,
				"...",
				gitDescription,
				metadata);
	}


	/**
	 * Creates and persists a new task status but doesn't associate it with the current thread.
	 * Should only be used in special situations.
	 *
	 * @param clazz The task class.
	 * @param description Description of the task.
	 * @param metadata Metadata for the task;
	 * @return The created task;
	 */
	public TaskStatus initialize(Class clazz, String description, Map metadata) {
		TaskStatus taskStatus = create(clazz, description, metadata);
		save(taskStatus);
		return taskStatus;
	}


	/**
	 * Creates and persists a new task status and associates it with the current thread.
	 * This should be used in most situations for creating a task status.
	 *
	 * @param clazz The task class.
	 * @param description Description of the task.
	 * @param metadata Metadata for the task;
	 * @return An auto-closer for finishing the task.
	 */
	@Transactional(propagation = Propagation.REQUIRES_NEW)
	public Task start(Class clazz, String description) {
		return start(clazz, description, Collections.emptyMap());
	}


	/**
	 * Creates and persists a new task status and associates it with the current thread.
	 * This should be used in most situations for creating a task status.
	 *
	 * @param clazz The task class.
	 * @param description Description of the task.
	 * @param metadata Metadata for the task;
	 * @return An auto-closer for finishing the task.
	 */
	@Transactional(propagation = Propagation.REQUIRES_NEW)
	public Task start(Class clazz, String description, Map metadata) {
		return couple(initialize(clazz, description, metadata));
	}


	/**
	 * Finishes any task status currently associated with the thread.
	 *
	 * @param taskStatusCloser The auto-closer for the task.
	 */
	@Transactional(propagation = Propagation.REQUIRES_NEW)
	public void finish(Task taskStatusCloser) {
		try {
			Optional current = Optional.ofNullable(taskStatusHolder.get())
					.filter(taskStatusCloser.taskStatus::equals)
					.flatMap(t -> decouple());

			if (!current.isPresent()) {
				LOG.warn("The task {} of type {} is not associated with the current thread: finishing FAILED",
						taskStatusCloser.taskStatus.getId(),
						taskStatusCloser.taskStatus.getGroup());
				return;
			}

			TaskStatus task = current.get();
			ResultState resultState = taskStatusCloser.getResultState();
			task.setState(resultState.getState());
			task.setMessage(resultState.getMessage());
			getThreadCpuTimeMs().map(t -> t - taskStatusCloser.startCpuTimeMs.orElse(0L)).ifPresent(t -> task.addMetadata("CpuTimeMs", t));
			getThreadAllocatedBytes().map(b -> b - taskStatusCloser.startAllocatedBytes.orElse(0L)).ifPresent(b -> task.addMetadata("AllocatedBytes", b));
			save(task);
		}
		catch (RuntimeException e) {
			LOG.warn("Failed finishing task status", e);
		}
	}


	/**
	 * Adds a single metadata for the message. The metadata is not persisted before
	 * the next updateMessage or finish method call.
	 *
	 * @param key The key for the metadata.
	 * @param value The value for the metadata.
	 */
	public void addMetadata(String key, Object value) {
		if (key != null && value != null) {
			try {
				getTaskStatus().ifPresent(t -> t.addMetadata(key, value.toString()));
			}
			catch (RuntimeException e) {
				LOG.warn("Failed adding metadata", e);
			}
		}
	}


	/**
	 * Updates the message for the task status currently associated with the thread.
	 *
	 * @param message The updated message.
	 * @param args Arguments for the updated message.
	 */
	@Transactional(propagation = Propagation.REQUIRES_NEW)
	public void updateMessage(String message, Object... args) {
		assertRunning();
		try {
			getTaskStatus().ifPresent(t -> {
				t.setMessage(format(message, args));
				save(t);
			});
		}
		catch (RuntimeException e) {
			LOG.warn("Failed updating message", e);
		}
	}


	private void save(TaskStatus taskStatus) {
		List metadata = new ArrayList<>(taskStatus.getMetadata());
		taskStatusRepository.saveAndFlush(taskStatus);
		metadata.stream().filter(m -> m.getId() == 0L)
				.forEach(taskStatusMetadataRepository::save);
		taskStatusMetadataRepository.flush();
		taskStatus.setMetadata(metadata);
	}


	private static String format(String message, Object... args) {
		if (message == null) {
			LOG.warn("Got null message, defaulting to empty string");
			message = "";
		}

		if (args == null || args.length == 0)
			return message;

		try {
			return String.format(message, args);
		}
		catch (Exception e) {
			LOG.warn("Failed formatting log message", e);
			return message + ": " + args;
		}
	}


	private static Optional getThreadCpuTimeMs() {
		return Optional.ofNullable(ManagementFactory.getThreadMXBean())
				.filter(ThreadMXBean::isThreadCpuTimeSupported)
				.filter(ThreadMXBean::isThreadCpuTimeEnabled)
				.map(b -> b.getThreadCpuTime(Thread.currentThread().getId()))
				.map(v -> v / 1000000);
	}


	private static Optional getThreadAllocatedBytes() {
		ThreadMXBean bean = ManagementFactory.getThreadMXBean();
		if (bean == null)
			return Optional.empty();

		try {
			Method method = bean.getClass().getMethod("getThreadAllocatedBytes", long.class);
			method.setAccessible(true);
			return Optional.of((Long)method.invoke(bean, Thread.currentThread().getId()));
		}
		catch (Exception e) {
			LOG.warn("Could not get thread allocated bytes count", e);
		}

		return Optional.empty();
	}


	public void assertRunning() {
		long taskStatusId = getTaskStatus().map(TaskStatus::getId).orElse(0L);
		if (taskStatusId == 0L)
			return;

		TaskState state = taskStatusRepository.findStateById(taskStatusId);
		if (TaskState.CANCELLING == state)
			throw new TriggerableCancelledException();
	}


	public class Task implements AutoCloseable {
		private final TaskStatus taskStatus;
		private final Job job;
		private ResultState resultState = DEFAULT_STATE;

		private final Optional startCpuTimeMs;
		private final Optional startAllocatedBytes;


		public Task(TaskStatus taskStatus, Job job) {
			this.taskStatus = taskStatus;
			this.job = job;
			this.startCpuTimeMs = getThreadCpuTimeMs();
			this.startAllocatedBytes = getThreadAllocatedBytes();
		}


		public void setResultState(TaskState state, String message, Object... args) {
			setResultState(new ResultState(state, message, args));
		}

		public void setResultState(ResultState resultState) {
			if (resultState != null && resultState.getState().isFinished()) {
				this.resultState = resultState;
				if (this.job != null)
					this.job.setResultState(resultState);
			}
		}

		public void addMetadata(String key, Object value) {
			taskStatus.addMetadata(key, value);
		}


		public ResultState getResultState() {
			return resultState;
		}


		@Override
		public void close() {
			finish(this);
			if (job != null)
				job.close();
		}
	}

}