From a2446080bc06219673b44453127277d14568dd08 Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <awilkinson@pivotal.io>
Date: Fri, 17 Jun 2016 13:10:02 +0100
Subject: [PATCH] Prevent GC pressure from causing an NPE in
 SimpleInMemoryRepository

Previously, SimpleInMemoryRepository used a ConcurrentReferenceHashMap
to store its locks. The type of map will discard its entries when the
JVM comes under GC pressure. With the code in its previous form, this
could lead to a NullPointerException when the following occurred:

1. putIfAbsent returned null indicating that a new entry has been added
   to the map
2. GC pressure caused the map to discard the new entry
3. get returned null as the entry has been discard

There are two problems with the existing code:

1. Its usage of a ConcurrentMap is incorrect. The correct usage is:
   a. Call get to see if the map already contains a lock
   b. If the lock is null, create a new one
   c. Call putIfAbsent to add the new lock
   d. If the return value is non-null, another thread has created the
      lock and it should be used. If the return value is null, use the
      new lock created in b.
2. Once the use of ConcurrentMap has been corrected, the fact that it is
   a ConcurrentReferenceHashMap means that different threads could
   access the same value using different locks. This would occur if one
   thread has retrieved a lock from the map and is using it, while GC
   causes the lock to be removed from the map. Another thread then
   attempts to get the lock and, as GC pressure has remove it, a new
   lock is created allowing concurrent access to the same value.

This commit updates the code to use the ConcurrentMap correctly and also
replaces the ConcurrentReferenceHashMap with a ConcurrentHashMap. This
means that the repository will now use slightly more memory but this is
outweighed by the benefits of thread-safe updates and no risk of an NPE.

Closes gh-6115
---
 .../util/SimpleInMemoryRepository.java        | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/util/SimpleInMemoryRepository.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/util/SimpleInMemoryRepository.java
index 81eb9477590..19105c53e40 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/util/SimpleInMemoryRepository.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/util/SimpleInMemoryRepository.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-2016 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.
@@ -18,29 +18,26 @@ package org.springframework.boot.actuate.metrics.util;
 
 import java.util.ArrayList;
 import java.util.NavigableMap;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentNavigableMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 
-import org.springframework.util.ConcurrentReferenceHashMap;
-
 /**
  * Repository utility that stores stuff in memory with period-separated String keys.
  *
  * @param <T> the type to store
  * @author Dave Syer
+ * @author Andy Wilkinson
  */
 public class SimpleInMemoryRepository<T> {
 
 	private ConcurrentNavigableMap<String, T> values = new ConcurrentSkipListMap<String, T>();
 
-	private final ConcurrentMap<String, Object> locks = new ConcurrentReferenceHashMap<String, Object>();
+	private final ConcurrentMap<String, Object> locks = new ConcurrentHashMap<String, Object>();
 
 	public T update(String name, Callback<T> callback) {
-		Object lock = this.locks.putIfAbsent(name, new Object());
-		if (lock == null) {
-			lock = this.locks.get(name);
-		}
+		Object lock = getLock(name);
 		synchronized (lock) {
 			T current = this.values.get(name);
 			T value = callback.modify(current);
@@ -49,6 +46,18 @@ public class SimpleInMemoryRepository<T> {
 		}
 	}
 
+	private Object getLock(String name) {
+		Object lock = this.locks.get(name);
+		if (lock == null) {
+			Object newLock = new Object();
+			lock = this.locks.putIfAbsent(name, newLock);
+			if (lock == null) {
+				lock = newLock;
+			}
+		}
+		return lock;
+	}
+
 	public void set(String name, T value) {
 		T current = this.values.get(name);
 		if (current != null) {