HIVE-29465:Prevent excessive query results cache usage at runtime

Author: ramitg254

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

@@ -5670,7 +5670,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_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/cache/results/QueryResultsCache.java

@@ -50,6 +50,7 @@
 import org.apache.hadoop.fs.ContentSummary;
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.FileUtil;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.permission.FsPermission;
 import org.apache.hadoop.hive.common.metrics.common.Metrics;
@@ -353,11 +354,13 @@ 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;
@@ -378,6 +381,8 @@ 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(
@@ -415,6 +420,14 @@ public Path getCacheDirPath() {
     return cacheDirPath;
   }
 
+  public void setSafeDir(String dirName) {
+    safeDir = dirName;
+  }
+
+  public String getSafeDir() {
+    return safeDir;
+  }
+
   /**
    * Check if the cache contains an entry for the requested LookupInfo.
    * @param request
@@ -513,6 +526,22 @@ public CacheEntry addToCache(QueryInfo queryInfo, ValidTxnWriteIdList txnWriteId
     return addedEntry;
   }
 
+  public void removeInvalidStaleFiles(FileSystem fs, Set<FileStatus> files) {
+    rwLock.writeLock().lock();
+    try {
+      for (FileStatus f : files) {
+        try {
+          fs.delete(f.getPath(), true);
+        } catch (IOException e) {
+          LOG.warn("Failed to clean up stale invalid file: {}",
+              f.getPath(), e);
+        }
+      }
+    } finally {
+      rwLock.writeLock().unlock();
+    }
+  }
+
   /**
    * Updates a pending cache entry with a FetchWork result from a finished query.
    * If successful the cache entry will be set to valid status and be usable for cached queries.
@@ -549,6 +578,41 @@ 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)));
+            boolean isMoved =
+                FileUtil.copy(srcFs, srcFile, cacheFs, destFile, false, conf);
+            if (!isMoved) {
+              succeeded = false;
+              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)));
+      }
+
       // Synchronize on the cache entry so that no one else can invalidate this entry
       // while we are in the process of setting it to valid.
       synchronized (cacheEntry) {

ql/src/java/org/apache/hadoop/hive/ql/parse/SemanticAnalyzer.java

@@ -7500,19 +7500,24 @@ protected Table getTargetTable(QB qb, String dest) throws SemanticException {
   }
 
   private Path getDestinationFilePath(QB qb, final String destinationFile, boolean isMmTable) {
+    Path defaultPath = new Path(destinationFile);
     if (this.isResultsCacheEnabled() && this.queryTypeCanUseCache(qb)) {
       assert (!isMmTable);
       QueryResultsCache instance = QueryResultsCache.getInstance();
       // QueryResultsCache should have been initialized by now
       if (instance != null) {
-        Path resultCacheTopDir = instance.getCacheDirPath();
-        String dirName = UUID.randomUUID().toString();
-        Path resultDir = new Path(resultCacheTopDir, dirName);
-        this.ctx.setFsResultCacheDirs(resultDir);
-        return resultDir;
+        if (!conf.getBoolVar(ConfVars.HIVE_QUERY_RESULTS_SAFE_CACHE_WRITE_ENABLED)) {
+          Path resultCacheTopDir = instance.getCacheDirPath();
+          String dirName = UUID.randomUUID().toString();
+          Path resultDir = new Path(resultCacheTopDir, dirName);
+          this.ctx.setFsResultCacheDirs(resultDir);
+          return resultDir;
+        } else {
+          instance.setSafeDir(defaultPath.toString());
+        }
       }
     }
-    return new Path(destinationFile);
+    return defaultPath;
   }
 
   @SuppressWarnings("nls")