changing totalCount to AtomicInteger data type

qubole · Jul 23, 2020 · 568bfda · 568bfda
1 parent 941fb29
commit 568bfda
Showing 1 changed file with 13 additions and 12 deletions.
diff --git a/rubix-spi/src/main/java/com/qubole/rubix/spi/fop/ObjectPoolPartition.java b/rubix-spi/src/main/java/com/qubole/rubix/spi/fop/ObjectPoolPartition.java
@@ -23,6 +23,7 @@
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static com.google.common.base.Preconditions.checkState;
 
@@ -36,11 +37,11 @@ public class ObjectPoolPartition<T>
   private final PoolConfig config;
   private final BlockingQueue<Poolable<T>> objectQueue;
   private final ObjectFactory<T> objectFactory;
-  private int totalCount;
   private final String host;
   private final int socketTimeout;
   private final int connectTimeout;
   private final Semaphore takeSemaphore;
+  private final AtomicInteger totalCount;
 
   public ObjectPoolPartition(ObjectPool<T> pool, PoolConfig config,
           ObjectFactory<T> objectFactory, BlockingQueue<Poolable<T>> queue, String host, String name)
@@ -52,14 +53,14 @@ public ObjectPoolPartition(ObjectPool<T> pool, PoolConfig config,
     this.host = host;
     this.socketTimeout = config.getSocketTimeoutMilliseconds();
     this.connectTimeout = config.getConnectTimeoutMilliseconds();
-    this.totalCount = 0;
+    this.totalCount = new AtomicInteger();
     this.log = new CustomLogger(name, host);
     this.takeSemaphore = new Semaphore(config.getMaxSize(), true);
     try {
       for (int i = 0; i < config.getMinSize(); i++) {
           T object = objectFactory.create(host, socketTimeout, connectTimeout);
           objectQueue.add(new Poolable<>(object, pool, host));
-          totalCount++;
+          totalCount.incrementAndGet();
       }
     }
     catch (Exception e) {
@@ -116,11 +117,11 @@ private Poolable<T> tryGetObject() throws Exception
       try {
         T object = objectFactory.create(host, socketTimeout, connectTimeout);
         poolable = new Poolable<>(object, pool, host);
-        totalCount++;
-        log.debug(String.format("Added a connection, new totalCount: %d, queueSize: %d", totalCount, objectQueue.size()));
+        totalCount.incrementAndGet();
+        log.debug(String.format("Added a connection, Pool state: totalCount: %s, queueSize: %d", totalCount, objectQueue.size()));
       }
       catch (Exception e) {
-        log.warn(String.format("Unable create a connection. Pool state: totalCount=%d queueSize=%d", totalCount, objectQueue.size()), e);
+        log.warn(String.format("Unable create a connection. Pool state: totalCount=%s queueSize=%d", totalCount, objectQueue.size()), e);
         if (poolable != null) {
           objectFactory.destroy(poolable.getObject());
           poolable.destroy();
@@ -144,20 +145,20 @@ public boolean decreaseObject(Poolable<T> obj)
     return true;
   }
 
-  private synchronized void objectRemoved()
+  private void objectRemoved()
   {
-    totalCount--;
+    totalCount.decrementAndGet();
   }
 
-  public synchronized int getTotalCount()
+  public int getTotalCount()
   {
-    return totalCount;
+    return totalCount.get();
   }
 
   // set the scavenge interval carefully
   public void scavenge() throws InterruptedException
   {
-    int delta = this.totalCount - config.getMinSize();
+    int delta = this.totalCount.get() - config.getMinSize();
     if (delta <= 0) {
       log.debug("Scavenge for delta <= 0, Skipping !!!");
       return;
@@ -195,7 +196,7 @@ public void scavenge() throws InterruptedException
   public synchronized int shutdown()
   {
     int removed = 0;
-    while (this.totalCount > 0) {
+    while (this.totalCount.get() > 0) {
       Poolable<T> obj = objectQueue.poll();
       if (obj != null) {
         decreaseObject(obj);