package com.hubspot.singularity.s3uploader;

import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.AmazonS3Client;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

  private final Map<String, AmazonS3> clientByKey;
  private final Map<String, AtomicInteger> clientHolds;

  @Inject
  public S3ClientProvider() {
    this.clientByKey = new HashMap<>();
    this.clientHolds = new HashMap<>();
  }

  public synchronized AmazonS3 getClient(BasicAWSCredentials credentials) {
    String key = credsToKey(credentials);
    AmazonS3 s3 = clientByKey.computeIfAbsent(key, k -> new AmazonS3Client(credentials));
    clientHolds.computeIfAbsent(key, k -> new AtomicInteger(0)).incrementAndGet();
    return s3;
  }

  public synchronized void returnClient(BasicAWSCredentials credentials) {
    String key = credsToKey(credentials);
    AtomicInteger holds = clientHolds.get(key);
    if (holds == null) {
      LOG.error("Client returned before being checked out!");
    } else {
      int remaining = holds.decrementAndGet();
      if (remaining == 0) {
        clientHolds.remove(key);
        clientByKey.remove(key);
      }
    }
  }

  public static String credsToKey(BasicAWSCredentials credentials) {
    // BasicAWSCredentials doesn't implement any sort of hash code, use this
    return String.format(
      "%s%s",
      credentials.getAWSAccessKeyId(),
      credentials.getAWSSecretKey()
    );
  }
}

@@ -262,42 +262,148 @@
 ava.
-util.HashMap;%0Aimport java.util.Map
+io.Closeable;%0Aimport java.io.IOException;%0Aimport java.util.HashMap;%0Aimport java.util.Map;%0Aimport java.util.Timer;%0Aimport java.util.TimerTask
 ;%0Aim
@@ -549,16 +549,37 @@
 rovider 
+implements Closeable 
 %7B%0A  priv
@@ -658,16 +658,45 @@
 lass);%0A%0A
+  private final Timer timer;%0A
   privat
@@ -831,24 +831,54 @@
 rovider() %7B%0A
+    this.timer = new Timer();%0A
     this.cli
@@ -1613,26 +1613,242 @@
     
-clientHolds.remove
+// Wait a bit to clean up in case anything else is about to use this client%0A        timer.schedule(%0A          new TimerTask() %7B%0A%0A            @Override%0A            public void run() %7B%0A              try %7B%0A                tryRemoveClient
 (key
@@ -1854,24 +1854,375 @@
 y);%0A        
+   
+   %7D catch (Exception e) %7B%0A                LOG.error(%22Could not clean up s3 client%22, e);%0A              %7D%0A            %7D%0A          %7D,%0A          15000%0A        );%0A      %7D%0A    %7D%0A  %7D%0A%0A  public synchronized void tryRemoveClient(String key) %7B%0A    AtomicInteger holds = clientHolds.get(key);%0A    if (holds.get() == 0) %7B%0A      clientHolds.remove(key);%0A      
 clientByKey.
@@ -2233,24 +2233,16 @@
 ve(key);
-%0A      %7D
 %0A    %7D%0A 
@@ -2516,11 +2516,91 @@
   );%0A  %7D
+%0A%0A  @Override%0A  public void close() throws IOException %7B%0A    timer.cancel();%0A  %7D
 %0A%7D%0A

Add delay for client cleanup