Simplify zone db locking to avoid a race

src/main/java/com/nvidia/spark/rapids/jni/GpuTimeZoneDB.java

@@ -52,67 +52,34 @@ public class GpuTimeZoneDB {
   // For the timezone database, we store the transitions in a ColumnVector that is a list of 
   // structs. The type of this column vector is:
   //   LIST<STRUCT<utcInstant: int64, localInstant: int64, offset: int32>>
-  private Map<String, Integer> zoneIdToTable;
+  private static Map<String, Integer> zoneIdToTable;
 
   // use this reference to indicate if time zone cache is initialized.
-  private HostColumnVector fixedTransitions;
+  private static HostColumnVector fixedTransitions;
 
-  // Guarantee singleton instance
-  private GpuTimeZoneDB() {
-  }
-
-  // singleton instance
-  private static final GpuTimeZoneDB instance = new GpuTimeZoneDB();
-
-  // This method is default visibility for testing purposes only.
-  // The instance will be never be exposed publicly for this class.
-  static GpuTimeZoneDB getInstance() {
-    return instance;
-  }
-
-  static class LoadingLock {
-    Boolean isLoading = false;
-
-    // record whether a shutdown is called ever.
-    // if `isCloseCalledEver` is true, then the following loading should be skipped.
-    Boolean isShutdownCalledEver = false;
-  }
-
-  private static final LoadingLock lock = new LoadingLock();
+  private static boolean isShutdownCalledEver = false;
 
   /**
    * This should be called on startup of an executor.
    * Runs in a thread asynchronously.
    * If `shutdown` was called ever, then will not load the cache
    */
   public static void cacheDatabaseAsync() {
-    synchronized (lock) {
-      if (lock.isShutdownCalledEver) {
-        // shutdown was called ever, will never load cache again.
+    // This has a race in that we could still launch a thread after
+    // shutting down. This is just to prevent the thread from launching
+    // in some cases.
+    synchronized (GpuTimeZoneDB.class) {
+      if (isShutdownCalledEver) {
+        log.error("cache async called after DB already loaded");
         return;
       }
-
-      if (lock.isLoading) {
-        // another thread is loading(), return
-        return;
-      } else {
-        lock.isLoading = true;
-      }
     }
-
     // start a new thread to load
     Runnable runnable = () -> {
       try {
-        instance.cacheDatabaseImpl();
+        cacheDatabaseImpl();
       } catch (Exception e) {
         log.error("cache time zone transitions cache failed", e);
-      } finally {
-        synchronized (lock) {
-          // now loading is done
-          lock.isLoading = false;
-          // `cacheDatabase` and `shutdown` may wait loading is done.
-          lock.notify();
-        }
       }
     };
     Thread thread = Executors.defaultThreadFactory().newThread(runnable);
@@ -127,55 +94,21 @@ public static void cacheDatabaseAsync() {
    * If cache is exits, do not load cache again.
    */
   public static void cacheDatabase() {
-    synchronized (lock) {
-      if (lock.isLoading) {
-        // another thread is loading(), wait loading is done
-        while (lock.isLoading) {
-          try {
-            lock.wait();
-          } catch (InterruptedException e) {
-            throw new IllegalStateException("cache time zone transitions cache failed", e);
-          }
-        }
-        return;
-      } else {
-        lock.isLoading = true;
-      }
-    }
-
-    try {
-      instance.cacheDatabaseImpl();
-    } finally {
-      // loading is done.
-      synchronized (lock) {
-        lock.isLoading = false;
-        // `cacheDatabase` and/or `shutdown` may wait loading is done.
-        lock.notify();
-      }
-    }
+    cacheDatabaseImpl();
   }
 
   /**
    * close the cache, used when Plugin is closing
    */
-  public static void shutdown() {
-    synchronized (lock) {
-      lock.isShutdownCalledEver = true;
-      while (lock.isLoading) {
-        // wait until loading is done
-        try {
-          lock.wait();
-        } catch (InterruptedException e) {
-          throw new IllegalStateException("shutdown time zone transitions cache failed", e);
-        }
-      }
-      instance.shutdownImpl();
-      // `cacheDatabase` and/or `shutdown` may wait loading is done.
-      lock.notify();
-    }
+  public static synchronized void shutdown() {
+    isShutdownCalledEver = true;
+    closeResources();
   }
 
-  private void cacheDatabaseImpl() {
+  private static synchronized void cacheDatabaseImpl() {
+    if (isShutdownCalledEver) {
+      throw new IllegalStateException("GpuTimeZoneDB has already been shut down");
+    }
     if (fixedTransitions == null) {
       try {
         loadData();
@@ -186,11 +119,7 @@ private void cacheDatabaseImpl() {
     }
   }
 
-  private void shutdownImpl() {
-    closeResources();
-  }
-
-  private void closeResources()  {
+  private static synchronized void closeResources()  {
     if (zoneIdToTable != null) {
       zoneIdToTable.clear();
       zoneIdToTable = null;
@@ -208,9 +137,12 @@ public static ColumnVector fromTimestampToUtcTimestamp(ColumnVector input, ZoneI
       throw new IllegalArgumentException(String.format("Unsupported timezone: %s",
           currentTimeZone.toString()));
     }
+    // there is technically a race condition on shutdown. Shutdown could be called after
+    // the database is cached. This would result in a null pointer exception at some point
+    // in the processing. This should be rare enough that it is not a big deal.
     cacheDatabase();
-    Integer tzIndex = instance.getZoneIDMap().get(currentTimeZone.normalized().toString());
-    try (Table transitions = instance.getTransitions()) {
+    Integer tzIndex = zoneIdToTable.get(currentTimeZone.normalized().toString());
+    try (Table transitions = getTransitions()) {
       return new ColumnVector(convertTimestampColumnToUTC(input.getNativeView(),
           transitions.getNativeView(), tzIndex));
     }
@@ -223,9 +155,12 @@ public static ColumnVector fromUtcTimestampToTimestamp(ColumnVector input, ZoneI
       throw new IllegalArgumentException(String.format("Unsupported timezone: %s",
           desiredTimeZone.toString()));
     }
+    // there is technically a race condition on shutdown. Shutdown could be called after
+    // the database is cached. This would result in a null pointer exception at some point
+    // in the processing. This should be rare enough that it is not a big deal.
     cacheDatabase();
-    Integer tzIndex = instance.getZoneIDMap().get(desiredTimeZone.normalized().toString());
-    try (Table transitions = instance.getTransitions()) {
+    Integer tzIndex = zoneIdToTable.get(desiredTimeZone.normalized().toString());
+    try (Table transitions = getTransitions()) {
       return new ColumnVector(convertUTCTimestampColumnToTimeZone(input.getNativeView(),
           transitions.getNativeView(), tzIndex));
     }
@@ -258,7 +193,7 @@ public static ZoneId getZoneId(String timeZoneId) {
   }
 
   @SuppressWarnings("unchecked")
-  private void loadData() {
+  private static synchronized void loadData() {
     try {
       List<List<HostColumnVector.StructData>> masterTransitions = new ArrayList<>();
       zoneIdToTable = new HashMap<>();
@@ -334,17 +269,13 @@ private void loadData() {
     }
   }
 
-  private Map<String, Integer> getZoneIDMap() {
-    return zoneIdToTable;
-  }
-
-  private Table getTransitions() {
+  private static synchronized Table getTransitions() {
     try (ColumnVector fixedTransitions = getFixedTransitions()) {
       return new Table(fixedTransitions);
     }
   }
 
-  private ColumnVector getFixedTransitions() {
+  private static synchronized ColumnVector getFixedTransitions() {
     return fixedTransitions.copyToDevice();
   }
 
@@ -358,16 +289,15 @@ private ColumnVector getFixedTransitions() {
    * @param zoneId
    * @return list of fixed transitions
    */
-  List getHostFixedTransitions(String zoneId) {
+  static synchronized List getHostFixedTransitions(String zoneId) {
     zoneId = ZoneId.of(zoneId).normalized().toString(); // we use the normalized form to dedupe
-    Integer idx = getZoneIDMap().get(zoneId);
+    Integer idx = zoneIdToTable.get(zoneId);
     if (idx == null) {
       return null;
     }
     return fixedTransitions.getList(idx);
   }
 
-
   private static native long convertTimestampColumnToUTC(long input, long transitions, int tzIndex);
 
   private static native long convertUTCTimestampColumnToTimeZone(long input, long transitions, int tzIndex);

src/test/java/com/nvidia/spark/rapids/jni/TimeZoneTest.java

@@ -1,5 +1,5 @@
 /*
-* Copyright (c)  2023, NVIDIA CORPORATION.
+* Copyright (c)  2023-2024, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -44,11 +44,10 @@ static void cleanup() {
   @Test
   void databaseLoadedTest() {
     // Check for a few timezones
-    GpuTimeZoneDB instance = GpuTimeZoneDB.getInstance();
-    List transitions = instance.getHostFixedTransitions("UTC+8");
+    List transitions = GpuTimeZoneDB.getHostFixedTransitions("UTC+8");
     assertNotNull(transitions);
     assertEquals(1, transitions.size());
-    transitions = instance.getHostFixedTransitions("Asia/Shanghai");
+    transitions = GpuTimeZoneDB.getHostFixedTransitions("Asia/Shanghai");
     assertNotNull(transitions);
     ZoneId shanghai = ZoneId.of("Asia/Shanghai").normalized();
     assertEquals(shanghai.getRules().getTransitions().size() + 1, transitions.size());