Simple in memory cache with plain Java

The following example is simple in memory cache using plain java without using any third party cache libraries. The below implementation shows expiry time based cache eviction and setting maximum cache size.

If used with Spring Framework, it is ideal to make this class as DisposableBean to safely execute the shutdown method during application shutdown event. As below

@Component
public class SimpleCache implements DisposableBean {
  @Override
  public void destroy() throws Exception {
   shutdown();
  }
}

Full implementation:

import java.time.Duration;
import java.time.Instant;
import java.util.Comparator;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class SimpleCache {
  public record HeavyObject(
   UUID id,
   String firstName,
   String lastName
  ) {
  }

  private record CacheValue(
   Instant timestamp,
   HeavyObject value
  ) {
  }

  private final ConcurrentHashMap<UUID, CacheValue> cache;
  private final ScheduledExecutorService executorService;
  private final Duration cacheExpiry;
  private final int maxSize;

  public SimpleCache(Duration cacheExpiry, int maxSize) {
   this.cacheExpiry = cacheExpiry;
   this.maxSize = maxSize;
   cache = new ConcurrentHashMap<>();
   executorService = Executors.newSingleThreadScheduledExecutor();
   executorService.scheduleWithFixedDelay(
    this::cleanUp,
    1,
    1,
    TimeUnit.SECONDS);
  }

  public HeavyObject get(UUID id) {
   CacheValue cacheEntry = cache.compute(
    id,
    (UUID idKey, CacheValue existingValue) -> {
     if (Objects.isNull(existingValue)) {
      // id is not cached, hence we need to lookup
      return new CacheValue(Instant.now(), lookup(idKey));
     } else {
      // id is already found,
      // we can return the value as is or update the timestamp if needed
      return existingValue;
     }
    });
   return cacheEntry.value;
  }

  public void shutdown() {
   executorService.shutdown();
  }

  private HeavyObject lookup(UUID id) {
   // this method lookup by id from database or api or heavy computation
   System.out.println("Looking up id = " + id);
   return new HeavyObject(id, "first-name" + id, "last-name" + id);
  }

  private void remove(UUID id) {
   System.out.println("Removing id = " + id);
   cache.remove(id);
  }

  private void cleanUp() {
   System.out.println("Initiating clean up");
   removeExpiredEntries();
   removeOldEntries();
  }

  // method to remove cached values by expiry time
  private void removeExpiredEntries() {
   Instant expiryTime = Instant.now().minus(cacheExpiry);
   cache.entrySet().stream()
    .filter(entry -> entry.getValue().timestamp.isBefore(expiryTime))
    .forEach(entry -> remove(entry.getKey()));
  }

  // method to remove old cached values once the cache reaches the maximum capacity
  private void removeOldEntries() {
   int cacheSize = cache.size();
   if (cacheSize > maxSize) {
    int removalCount = cacheSize - maxSize;
    cache.entrySet().stream()
     .sorted(Comparator.comparing(entry -> entry.getValue().timestamp))
     .limit(removalCount)
     .forEach(entry -> remove(entry.getKey()));
   }
  }

}