Make BufferingApplicationStartup thread safe

Update `BufferingApplicationStartup` to use thread safe data structures.

Prior to this commit, it was possible for calls from different threads
(for example due to request scope beans) to cause a
NoSuchElementException to be thrown.

Closes gh-25792

Author: philwebb

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/endpoint/web/documentation/StartupEndpointDocumentationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2021 the 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.
@@ -46,11 +46,10 @@ class StartupEndpointDocumentationTests extends MockMvcEndpointDocumentationTest
 	void appendSampleStartupSteps(@Autowired BufferingApplicationStartup applicationStartup) {
 		StartupStep starting = applicationStartup.start("spring.boot.application.starting");
 		starting.tag("mainApplicationClass", "com.example.startup.StartupApplication");
-		starting.end();
-
 		StartupStep instantiate = applicationStartup.start("spring.beans.instantiate");
 		instantiate.tag("beanName", "homeController");
 		instantiate.end();
+		starting.end();
 	}
 
 	@Test
@@ -67,14 +66,13 @@ void startup() throws Exception {
 				fieldWithPath("timeline.events.[].startupStep.name").description("The name of the StartupStep."),
 				fieldWithPath("timeline.events.[].startupStep.id").description("The id of this StartupStep."),
 				fieldWithPath("timeline.events.[].startupStep.parentId")
-						.description("The parent id for this StartupStep."),
+						.description("The parent id for this StartupStep.").optional(),
 				fieldWithPath("timeline.events.[].startupStep.tags")
 						.description("An array of key/value pairs with additional step info."),
 				fieldWithPath("timeline.events.[].startupStep.tags[].key")
 						.description("The key of the StartupStep Tag."),
 				fieldWithPath("timeline.events.[].startupStep.tags[].value")
 						.description("The value of the StartupStep Tag."));
-
 		this.mockMvc.perform(post("/actuator/startup")).andExpect(status().isOk())
 				.andDo(document("startup", responseFields));
 	}

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/context/metrics/buffering/BufferedStartupStep.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the 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.
@@ -16,43 +16,54 @@
 
 package org.springframework.boot.context.metrics.buffering;
 
-import java.util.Iterator;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import org.springframework.core.metrics.StartupStep;
+import org.springframework.util.Assert;
 
 /**
  * {@link StartupStep} implementation to be buffered by
  * {@link BufferingApplicationStartup}. Its processing time is recorded using
  * {@link System#nanoTime()}.
  *
  * @author Brian Clozel
+ * @author Phillip Webb
  */
 class BufferedStartupStep implements StartupStep {
 
 	private final String name;
 
 	private final long id;
 
-	private final Long parentId;
+	private final BufferedStartupStep parent;
 
-	private long startTime;
+	private final List<Tag> tags = new ArrayList<>();
 
-	private long endTime;
+	private final Consumer<BufferedStartupStep> recorder;
 
-	private final DefaultTags tags;
+	private final Instant startTime;
 
-	private final Consumer<BufferedStartupStep> recorder;
+	private final AtomicBoolean ended = new AtomicBoolean();
 
-	BufferedStartupStep(long id, String name, Long parentId, Consumer<BufferedStartupStep> recorder) {
-		this.id = id;
-		this.parentId = parentId;
-		this.tags = new DefaultTags();
+	BufferedStartupStep(BufferedStartupStep parent, String name, long id, Instant startTime,
+			Consumer<BufferedStartupStep> recorder) {
+		this.parent = parent;
 		this.name = name;
+		this.id = id;
+		this.startTime = startTime;
 		this.recorder = recorder;
 	}
 
+	BufferedStartupStep getParent() {
+		return this.parent;
+	}
+
 	@Override
 	public String getName() {
 		return this.name;
@@ -63,88 +74,40 @@ public long getId() {
 		return this.id;
 	}
 
+	Instant getStartTime() {
+		return this.startTime;
+	}
+
 	@Override
 	public Long getParentId() {
-		return this.parentId;
+		return (this.parent != null) ? this.parent.getId() : null;
 	}
 
 	@Override
 	public Tags getTags() {
-		return this.tags;
+		return Collections.unmodifiableList(this.tags)::iterator;
 	}
 
 	@Override
-	public StartupStep tag(String key, String value) {
-		if (this.endTime != 0L) {
-			throw new IllegalStateException("StartupStep has already ended.");
-		}
-		this.tags.add(key, value);
-		return this;
+	public StartupStep tag(String key, Supplier<String> value) {
+		return this.tag(key, value.get());
 	}
 
 	@Override
-	public StartupStep tag(String key, Supplier<String> value) {
-		return this.tag(key, value.get());
+	public StartupStep tag(String key, String value) {
+		Assert.state(!this.ended.get(), "StartupStep has already ended.");
+		this.tags.add(new DefaultTag(key, value));
+		return this;
 	}
 
 	@Override
 	public void end() {
+		this.ended.set(true);
 		this.recorder.accept(this);
 	}
 
-	long getStartTime() {
-		return this.startTime;
-	}
-
-	void recordStartTime(long startTime) {
-		this.startTime = startTime;
-	}
-
-	long getEndTime() {
-		return this.endTime;
-	}
-
-	void recordEndTime(long endTime) {
-		this.endTime = endTime;
-	}
-
-	static class DefaultTags implements Tags {
-
-		private Tag[] tags = new Tag[0];
-
-		void add(String key, String value) {
-			Tag[] newTags = new Tag[this.tags.length + 1];
-			System.arraycopy(this.tags, 0, newTags, 0, this.tags.length);
-			newTags[newTags.length - 1] = new DefaultTag(key, value);
-			this.tags = newTags;
-		}
-
-		@Override
-		public Iterator<Tag> iterator() {
-			return new TagsIterator();
-		}
-
-		private class TagsIterator implements Iterator<Tag> {
-
-			private int index = 0;
-
-			@Override
-			public boolean hasNext() {
-				return this.index < DefaultTags.this.tags.length;
-			}
-
-			@Override
-			public Tag next() {
-				return DefaultTags.this.tags[this.index++];
-			}
-
-			@Override
-			public void remove() {
-				throw new UnsupportedOperationException("tags are append only");
-			}
-
-		}
-
+	boolean isEnded() {
+		return this.ended.get();
 	}
 
 	static class DefaultTag implements Tag {

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/context/metrics/buffering/BufferingApplicationStartup.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 the 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.
@@ -16,15 +16,17 @@
 
 package org.springframework.boot.context.metrics.buffering;
 
+import java.time.Clock;
 import java.time.Instant;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Deque;
+import java.util.Iterator;
 import java.util.List;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 
+import org.springframework.boot.context.metrics.buffering.StartupTimeline.TimelineEvent;
 import org.springframework.core.metrics.ApplicationStartup;
 import org.springframework.core.metrics.StartupStep;
 import org.springframework.util.Assert;
@@ -45,32 +47,40 @@
  * </ul>
  *
  * @author Brian Clozel
+ * @author Phillip Webb
  * @since 2.4.0
  */
 public class BufferingApplicationStartup implements ApplicationStartup {
 
-	private Instant recordingStartTime;
+	private final int capacity;
 
-	private long recordingStartNanos;
+	private final Clock clock;
 
-	private long currentSequenceId = 0;
+	private Instant startTime;
 
-	private final Deque<Long> currentSteps;
+	private final AtomicInteger idSeq = new AtomicInteger();
 
-	private final BlockingQueue<BufferedStartupStep> recordedSteps;
+	private Predicate<StartupStep> filter = (step) -> true;
 
-	private Predicate<StartupStep> stepFilters = (step) -> true;
+	private final AtomicReference<BufferedStartupStep> current = new AtomicReference<>();
+
+	private final AtomicInteger estimatedSize = new AtomicInteger();
+
+	private final ConcurrentLinkedQueue<TimelineEvent> events = new ConcurrentLinkedQueue<>();
 
 	/**
 	 * Create a new buffered {@link ApplicationStartup} with a limited capacity and starts
 	 * the recording of steps.
 	 * @param capacity the configured capacity; once reached, new steps are not recorded.
 	 */
 	public BufferingApplicationStartup(int capacity) {
-		this.currentSteps = new ArrayDeque<>();
-		this.currentSteps.offerFirst(this.currentSequenceId);
-		this.recordedSteps = new LinkedBlockingQueue<>(capacity);
-		startRecording();
+		this(capacity, Clock.systemDefaultZone());
+	}
+
+	BufferingApplicationStartup(int capacity, Clock clock) {
+		this.capacity = capacity;
+		this.clock = clock;
+		this.startTime = clock.instant();
 	}
 
 	/**
@@ -81,9 +91,8 @@ public BufferingApplicationStartup(int capacity) {
 	 * already.
 	 */
 	public void startRecording() {
-		Assert.state(this.recordedSteps.isEmpty(), "Cannot restart recording once steps have been buffered.");
-		this.recordingStartTime = Instant.now();
-		this.recordingStartNanos = getCurrentTime();
+		Assert.state(this.events.isEmpty(), "Cannot restart recording once steps have been buffered.");
+		this.startTime = this.clock.instant();
 	}
 
 	/**
@@ -93,7 +102,42 @@ public void startRecording() {
 	 * @param filter the predicate filter to add.
 	 */
 	public void addFilter(Predicate<StartupStep> filter) {
-		this.stepFilters = this.stepFilters.and(filter);
+		this.filter = this.filter.and(filter);
+	}
+
+	@Override
+	public StartupStep start(String name) {
+		int id = this.idSeq.getAndIncrement();
+		Instant start = this.clock.instant();
+		while (true) {
+			BufferedStartupStep current = this.current.get();
+			BufferedStartupStep parent = getLatestActive(current);
+			BufferedStartupStep next = new BufferedStartupStep(parent, name, id, start, this::record);
+			if (this.current.compareAndSet(current, next)) {
+				return next;
+			}
+		}
+	}
+
+	private void record(BufferedStartupStep step) {
+		if (this.filter.test(step) && this.estimatedSize.get() < this.capacity) {
+			this.estimatedSize.incrementAndGet();
+			this.events.add(new TimelineEvent(step, this.clock.instant()));
+		}
+		while (true) {
+			BufferedStartupStep current = this.current.get();
+			BufferedStartupStep next = getLatestActive(current);
+			if (this.current.compareAndSet(current, next)) {
+				return;
+			}
+		}
+	}
+
+	private BufferedStartupStep getLatestActive(BufferedStartupStep step) {
+		while (step != null && step.isEnded()) {
+			step = step.getParent();
+		}
+		return step;
 	}
 
 	/**
@@ -105,7 +149,7 @@ public void addFilter(Predicate<StartupStep> filter) {
 	 * @return a snapshot of currently buffered steps.
 	 */
 	public StartupTimeline getBufferedTimeline() {
-		return new StartupTimeline(this.recordingStartTime, this.recordingStartNanos, this.recordedSteps);
+		return new StartupTimeline(this.startTime, new ArrayList<>(this.events));
 	}
 
 	/**
@@ -116,30 +160,14 @@ public StartupTimeline getBufferedTimeline() {
 	 * @return buffered steps drained from the buffer.
 	 */
 	public StartupTimeline drainBufferedTimeline() {
-		List<BufferedStartupStep> steps = new ArrayList<>(this.recordedSteps.size());
-		this.recordedSteps.drainTo(steps);
-		return new StartupTimeline(this.recordingStartTime, this.recordingStartNanos, steps);
-	}
-
-	@Override
-	public StartupStep start(String name) {
-		BufferedStartupStep step = new BufferedStartupStep(++this.currentSequenceId, name,
-				this.currentSteps.peekFirst(), this::record);
-		step.recordStartTime(getCurrentTime());
-		this.currentSteps.offerFirst(this.currentSequenceId);
-		return step;
-	}
-
-	private void record(BufferedStartupStep step) {
-		step.recordEndTime(getCurrentTime());
-		if (this.stepFilters.test(step)) {
-			this.recordedSteps.offer(step);
+		List<TimelineEvent> events = new ArrayList<>();
+		Iterator<TimelineEvent> iterator = this.events.iterator();
+		while (iterator.hasNext()) {
+			events.add(iterator.next());
+			iterator.remove();
 		}
-		this.currentSteps.removeFirst();
-	}
-
-	private long getCurrentTime() {
-		return System.nanoTime();
+		this.estimatedSize.set(0);
+		return new StartupTimeline(this.startTime, events);
 	}
 
 }