corrected implementation

Author: ramitg254

common/src/java/org/apache/hadoop/hive/conf/HiveConf.java

@@ -5679,9 +5679,9 @@ public static enum ConfVars {
     HIVE_QUERY_RESULTS_CACHE_ENABLED("hive.query.results.cache.enabled", true,
         "If the query results cache is enabled. This will keep results of previously executed queries " +
         "to be reused if the same query is executed again."),
-    HIVE_QUERY_RESULTS_SAFE_CACHE_WRITE_ENABLED("hive.query.results.safe.cache.write.enabled", false,
-        "If the query results safe cache is enabled. This will safely write to cache directory by first evaluating " +
-            "the cache entry is not overspilling the the cache directory before writing it to cache directory "),
+    HIVE_QUERY_RESULTS_SAFE_CACHE_WRITE_ENABLED("hive.query.results.cache.safe.write.enabled", false,
+        "Write results to the query output path, then copy them into the results cache if size limits allow. " +
+            "Adds overhead from that copy versus writing result files once under the cache directory."),
     HIVE_QUERY_RESULTS_CACHE_NONTRANSACTIONAL_TABLES_ENABLED("hive.query.results.cache.nontransactional.tables.enabled", false,
         "If the query results cache is enabled for queries involving non-transactional tables." +
         "Users who enable this setting should be willing to tolerate some amount of stale results in the cache."),

ql/src/java/org/apache/hadoop/hive/ql/Context.java

@@ -107,6 +107,12 @@ public class Context {
   // keeps track of result cache dir for the query, later cleaned up by context cleanup
   private Path fsResultCacheDirs = null;
 
+  /**
+   * When HIVE_QUERY_RESULTS_SAFE_CACHE_WRITE_ENABLED is true, the absolute path prefix under which
+   * execution wrote fetchable result files.
+   */
+  private String cacheSafeWriteSourceDir;
+
   private Configuration conf;
   protected int pathid = 10000;
   private int moveTaskId = 0;
@@ -452,6 +458,7 @@ protected Context(Context ctx) {
     this.scratchDirPermission = ctx.scratchDirPermission;
     this.fsScratchDirs.putAll(ctx.fsScratchDirs);
     this.fsResultCacheDirs = ctx.fsResultCacheDirs;
+    this.cacheSafeWriteSourceDir = ctx.cacheSafeWriteSourceDir;
     this.conf = ctx.conf;
     this.pathid = ctx.pathid;
     this.explainConfig = ctx.explainConfig;
@@ -492,6 +499,14 @@ public Path getFsResultCacheDirs() {
     return this.fsResultCacheDirs;
   }
 
+  public void setCacheSafeWriteSourceDir(String cacheSafeWriteSourceDir) {
+    this.cacheSafeWriteSourceDir = cacheSafeWriteSourceDir;
+  }
+
+  public String getCacheSafeWriteSourceDir() {
+    return cacheSafeWriteSourceDir;
+  }
+
   public Map<LoadTableDesc, WriteEntity> getLoadTableOutputMap() {
     return loadTableOutputMap;
   }

ql/src/java/org/apache/hadoop/hive/ql/Executor.java

@@ -227,6 +227,7 @@ private void preExecutionCacheActions() throws Exception {
       CacheEntry pendingCacheEntry =
           QueryResultsCache.getInstance().addToCache(driverContext.getCacheUsage().getQueryInfo(), txnWriteIdList);
       if (pendingCacheEntry != null) {
+        pendingCacheEntry.setSafeSourceDir(context.getCacheSafeWriteSourceDir());
         // Update cacheUsage to reference the pending entry.
         this.driverContext.getCacheUsage().setCacheEntry(pendingCacheEntry);
       }
@@ -436,7 +437,7 @@ private void postExecutionCacheActions() throws Exception {
 
       CacheEntry cacheEntry = driverContext.getCacheUsage().getCacheEntry();
       boolean savedToCache = QueryResultsCache.getInstance().setEntryValid(cacheEntry,
-          driverContext.getPlan().getFetchTask().getWork());
+          driverContext.getPlan().getFetchTask().getWork(), driverContext.getConf());
       LOG.info("savedToCache: {} ({})", savedToCache, cacheEntry);
       if (savedToCache) {
         useFetchFromCache(driverContext.getCacheUsage().getCacheEntry());

ql/src/java/org/apache/hadoop/hive/ql/cache/results/QueryResultsCache.java

@@ -195,6 +195,12 @@ public static class CacheEntry {
     private Path cachedResultsPath;
     private Set<FileStatus> cachedResultPaths;
 
+    /**
+     * Absolute path prefix for result files when using safe cache write; see
+     * {@link org.apache.hadoop.hive.ql.Context#getCacheSafeWriteSourceDir()}.
+     */
+    private String safeSourceDir;
+
     // Cache administration
     private long size;
     private AtomicInteger readers = new AtomicInteger(0);
@@ -288,6 +294,14 @@ public QueryInfo getQueryInfo() {
       return queryInfo;
     }
 
+    public void setSafeSourceDir(String safeSourceDir) {
+      this.safeSourceDir = safeSourceDir;
+    }
+
+    public String getSafeSourceDir() {
+      return safeSourceDir;
+    }
+
     public Path getCachedResultsPath() {
       return cachedResultsPath;
     }
@@ -354,13 +368,11 @@ public Stream<String> getTableNames() {
   private Path cacheDirPath;
   private Path zeroRowsPath;
   private long cacheSize = 0;
-  private boolean isSafeCacheWriteEnabled;
   private long maxCacheSize;
   private long maxEntrySize;
   private long maxEntryLifetime;
   private ReadWriteLock rwLock = new ReentrantReadWriteLock();
   private ScheduledFuture<?> invalidationPollFuture;
-  private String safeDir;
 
   private QueryResultsCache(HiveConf configuration) throws IOException {
     this.conf = configuration;
@@ -381,8 +393,6 @@ private QueryResultsCache(HiveConf configuration) throws IOException {
     // Results cache directory should be cleaned up at process termination.
     fs.deleteOnExit(cacheDirPath);
 
-    isSafeCacheWriteEnabled = conf.getBoolVar(HiveConf.ConfVars.HIVE_QUERY_RESULTS_SAFE_CACHE_WRITE_ENABLED);
-
     maxCacheSize = conf.getLongVar(HiveConf.ConfVars.HIVE_QUERY_RESULTS_CACHE_MAX_SIZE);
     maxEntrySize = conf.getLongVar(HiveConf.ConfVars.HIVE_QUERY_RESULTS_CACHE_MAX_ENTRY_SIZE);
     maxEntryLifetime = conf.getTimeVar(
@@ -420,12 +430,17 @@ public Path getCacheDirPath() {
     return cacheDirPath;
   }
 
-  public void setSafeDir(String dirName) {
-    safeDir = dirName;
-  }
-
-  public String getSafeDir() {
-    return safeDir;
+  /**
+   * Runs {@code action} while holding {@link #rwLock} in exclusive (write) mode.
+   */
+  private void withWriteLock(Runnable action) {
+    Lock writeLock = rwLock.writeLock();
+    try {
+      writeLock.lock();
+      action.run();
+    } finally {
+      writeLock.unlock();
+    }
   }
 
   /**
@@ -507,10 +522,7 @@ public CacheEntry addToCache(QueryInfo queryInfo, ValidTxnWriteIdList txnWriteId
     addedEntry.queryInfo = queryInfo;
     addedEntry.txnWriteIdList = txnWriteIdList;
 
-    Lock writeLock = rwLock.writeLock();
-    try {
-      writeLock.lock();
-
+    withWriteLock(() -> {
       LOG.info("Adding placeholder cache entry for query '{}'", queryText);
 
       // Add the entry to the cache structures while under write lock.
@@ -519,16 +531,13 @@ public CacheEntry addToCache(QueryInfo queryInfo, ValidTxnWriteIdList txnWriteId
       // Index of entries by table usage.
       addedEntry.getTableNames()
           .forEach(tableName -> addToEntryMap(tableToEntryMap, tableName, addedEntry));
-    } finally {
-      writeLock.unlock();
-    }
+    });
 
     return addedEntry;
   }
 
   public void removeInvalidStaleFiles(FileSystem fs, Set<FileStatus> files) {
-    rwLock.writeLock().lock();
-    try {
+    withWriteLock(() -> {
       for (FileStatus f : files) {
         try {
           fs.delete(f.getPath(), true);
@@ -537,9 +546,7 @@ public void removeInvalidStaleFiles(FileSystem fs, Set<FileStatus> files) {
               f.getPath(), e);
         }
       }
-    } finally {
-      rwLock.writeLock().unlock();
-    }
+    });
   }
 
   /**
@@ -549,9 +556,11 @@ public void removeInvalidStaleFiles(FileSystem fs, Set<FileStatus> files) {
    * CacheEntry.releaseReader() should be called when the caller is done with the cache entry.
    * @param cacheEntry
    * @param fetchWork
+   * @param queryConf session (or query) Hive configuration; used for safe-cache-write and filesystem access
+   *                  so per-session {@code SET hive.query.results.cache.safe.write.enabled} is honored
    * @return
    */
-  public boolean setEntryValid(CacheEntry cacheEntry, FetchWork fetchWork) {
+  public boolean setEntryValid(CacheEntry cacheEntry, FetchWork fetchWork, HiveConf queryConf) {
     Path queryResultsPath = null;
     Path cachedResultsPath = null;
 
@@ -561,7 +570,7 @@ public boolean setEntryValid(CacheEntry cacheEntry, FetchWork fetchWork) {
 
       boolean requiresCaching = true;
       queryResultsPath = fetchWork.getTblDir();
-      FileSystem resultsFs = queryResultsPath.getFileSystem(conf);
+      FileSystem resultsFs = queryResultsPath.getFileSystem(queryConf);
 
       long resultSize = 0;
       for(FileStatus fs:fetchWork.getFilesToFetch()) {
@@ -578,37 +587,8 @@ public boolean setEntryValid(CacheEntry cacheEntry, FetchWork fetchWork) {
         return false;
       }
 
-      if (isSafeCacheWriteEnabled) {
-        Path resultDir = new Path(cacheDirPath, UUID.randomUUID().toString());
-        FileSystem cacheFs = resultDir.getFileSystem(conf);
-        cacheFs.mkdirs(resultDir);
-
-        Set<FileStatus> cacheFilesToFetch = new HashSet<>();
-        rwLock.writeLock().lock();
-        boolean succeeded = true;
-        try {
-          for (FileStatus fs : fetchWork.getFilesToFetch()) {
-            FileSystem srcFs = fs.getPath().getFileSystem(conf);
-            Path srcFile = fs.getPath();
-            Path destFile = new Path(resultDir,
-                new Path(fs.getPath().toString().substring(safeDir.length() + 1)));
-            succeeded = FileUtil.copy(srcFs, srcFile, cacheFs, destFile, false, conf);
-            if (!succeeded) {
-              throw new IOException("File copy failed for " + srcFile + " -> " + destFile);
-            }
-            cacheFilesToFetch.add(cacheFs.getFileStatus(destFile));
-          }
-        } catch (IOException e) {
-          LOG.warn("Failed to write cache entry to {}", resultDir, e);
-        } finally {
-          rwLock.writeLock().unlock();
-        }
-        if (!succeeded) {
-          removeInvalidStaleFiles(cacheFs, cacheFilesToFetch);
-          return false;
-        }
-        fetchWork.setFilesToFetch(cacheFilesToFetch);
-        fetchWork.setTblDir(new Path(resultDir, fetchWork.getTblDir().toString().substring(safeDir.length() + 1)));
+      if (!rewriteFetchWorkForSafeCacheWrite(cacheEntry, fetchWork, queryConf)) {
+        return false;
       }
 
       // Synchronize on the cache entry so that no one else can invalidate this entry
@@ -663,10 +643,68 @@ public boolean setEntryValid(CacheEntry cacheEntry, FetchWork fetchWork) {
     return true;
   }
 
+  private boolean rewriteFetchWorkForSafeCacheWrite(CacheEntry cacheEntry, FetchWork fetchWork, HiveConf queryConf)
+      throws IOException {
+    if (!queryConf.getBoolVar(HiveConf.ConfVars.HIVE_QUERY_RESULTS_SAFE_CACHE_WRITE_ENABLED)) {
+      return true;
+    }
+    String safeDir = cacheEntry.getSafeSourceDir();
+    if (safeDir == null) {
+      LOG.error("Safe cache write enabled but cache entry has no safe source dir; query: {}",
+          cacheEntry.getQueryInfo().getLookupInfo().getQueryText());
+      return false;
+    }
+    final int safeDirAndSepLen = safeDir.length() + Path.SEPARATOR.length();
+    Path resultDir = new Path(cacheDirPath, UUID.randomUUID().toString());
+    FileSystem cacheFs = resultDir.getFileSystem(queryConf);
+    cacheFs.mkdirs(resultDir);
+
+    Set<FileStatus> cacheFilesToFetch = new HashSet<>();
+    boolean succeeded =
+        copyFetchWorkFilesIntoCacheDirUnderWriteLock(fetchWork, resultDir, cacheFs, safeDirAndSepLen,
+            queryConf, cacheFilesToFetch);
+    if (!succeeded) {
+      removeInvalidStaleFiles(cacheFs, cacheFilesToFetch);
+      return false;
+    }
+    fetchWork.setFilesToFetch(cacheFilesToFetch);
+    fetchWork.setTblDir(new Path(resultDir, fetchWork.getTblDir().toString().substring(safeDirAndSepLen)));
+    return true;
+  }
+
+  /**
+   * Copies each file from the fetch work into {@code resultDir}, preserving relative paths below
+   * the safe-write prefix. Runs under {@link #rwLock} write lock so cache metadata and filesystem
+   * layout stay consistent with other cache mutations.
+   *
+   * @param destFileStatuses output: destination {@link FileStatus} entries for successfully copied files
+   * @return {@code true} if every file was copied and listed; {@code false} otherwise
+   */
+  private boolean copyFetchWorkFilesIntoCacheDirUnderWriteLock(FetchWork fetchWork, Path resultDir,
+      FileSystem cacheFs, int safeDirAndSepLen, HiveConf queryConf, Set<FileStatus> destFileStatuses) {
+    final boolean[] succeeded = { true };
+    withWriteLock(() -> {
+      try {
+        for (FileStatus fs : fetchWork.getFilesToFetch()) {
+          FileSystem srcFs = fs.getPath().getFileSystem(queryConf);
+          Path srcFile = fs.getPath();
+          Path destFile = new Path(resultDir,
+              new Path(fs.getPath().toString().substring(safeDirAndSepLen)));
+          succeeded[0] = FileUtil.copy(srcFs, srcFile, cacheFs, destFile, false, queryConf);
+          if (!succeeded[0]) {
+            throw new IOException("File copy failed for " + srcFile + " -> " + destFile);
+          }
+          destFileStatuses.add(cacheFs.getFileStatus(destFile));
+        }
+      } catch (IOException e) {
+        LOG.warn("Failed to write cache entry to {}", resultDir, e);
+      }
+    });
+    return succeeded[0];
+  }
+
   public void clear() {
-    Lock writeLock = rwLock.writeLock();
-    try {
-      writeLock.lock();
+    withWriteLock(() -> {
       LOG.info("Clearing the results cache");
       CacheEntry[] allEntries = null;
       synchronized (lru) {
@@ -679,9 +717,7 @@ public void clear() {
           LOG.error("Error removing cache entry " + entry, err);
         }
       }
-    } finally {
-      writeLock.unlock();
-    }
+    });
   }
 
   public long getSize() {
@@ -697,27 +733,21 @@ public long getSize() {
   public void notifyTableChanged(String dbName, String tableName, long updateTime) {
     LOG.debug("Table changed: {}.{}, at {}", dbName, tableName, updateTime);
     // Invalidate all cache entries using this table.
-    List<CacheEntry> entriesToInvalidate = null;
-    rwLock.writeLock().lock();
-    try {
+    withWriteLock(() -> {
       String key = (dbName.toLowerCase() + "." + tableName.toLowerCase());
       Set<CacheEntry> entriesForTable = tableToEntryMap.get(key);
       if (entriesForTable != null) {
         // Possible concurrent modification issues if we try to remove cache entries while
         // traversing the cache structures. Save the entries to remove in a separate list.
-        entriesToInvalidate = new ArrayList<>(entriesForTable);
-      }
-      if (entriesToInvalidate != null) {
+        List<CacheEntry> entriesToInvalidate = new ArrayList<>(entriesForTable);
         for (CacheEntry entry : entriesToInvalidate) {
           // Ignore updates that occured before this cached query was created.
           if (entry.getQueryInfo().getQueryTime() <= updateTime) {
             removeEntry(entry);
           }
         }
       }
-    } finally {
-      rwLock.writeLock().unlock();
-    }
+    });
   }
 
   private static final int INITIAL_LRU_SIZE = 16;
@@ -799,15 +829,12 @@ private boolean entryMatches(LookupInfo lookupInfo, CacheEntry entry, Set<CacheE
 
   public void removeEntry(CacheEntry entry) {
     entry.invalidate();
-    rwLock.writeLock().lock();
-    try {
+    withWriteLock(() -> {
       removeFromLookup(entry);
       lru.remove(entry);
       // Should the cache size be updated here, or after the result data has actually been deleted?
       cacheSize -= entry.size;
-    } finally {
-      rwLock.writeLock().unlock();
-    }
+    });
   }
 
   private void removeFromLookup(CacheEntry entry) {