[taskserver] Handle concurrent uploads using a counter and locking mechanism and add new configs (TASK_SERVER_V2) (#3733)

## What changes were proposed in this pull request?

- Handle concurrent uploads using a counter and locking mechnaism. 
  The "upload_available_space" key is set in Redis when Hue boots up. It calculates the actual available /tmp space and stores it. Whenever a new file upload is started, a new key is created in the format "upload_<unique_uuid_hash>" and stores the upload file size in Redis. Also, a second key in the format "upload_<same_uuid>_timestamp" (stores the timestamp) is created, and we reduce the "upload_available_space" value by the new upload file size. When the upload completes, the "upload_<unique_uuid_hash>" and "upload_<same_uuid>_timestamp" keys are deleted, and the reserved space is given back to "upload_available_space". 
  
  Assume we have 1GB free space:
  
  - `upload_available_space = 1GB`
  
  When file1 (100MB) upload is triggered:
  
  - `upload_<uuid1> = 100MB`
  - `upload_<uuid1>_timestamp = <timestamp>`
  - `upload_available_space = 1GB - 100MB = 900MB`
  
  When file2 (100MB) upload is triggered:
  
  - `upload_<uuid2> = 100MB`
  - `upload_<uuid2>_timestamp = <timestamp>`
  - `upload_available_space = 900MB - 100MB = 800MB`
  
  When the uploads are complete, the keys (`upload_<uuid1>`, `upload_<uuid1>_timestamp`) and (`upload_<uuid2>`, `upload_<uuid2>_timestamp`) are deleted, and `upload_available_space` is updated back:
  
  - `upload_available_space = 800MB + 100MB + 100MB = 1GB`
  
  In case of failed uploads, the same process is repeated for each retry made by the user. However, we would have leftover `upload_` keys for each retry since the keys will only be deleted on a successful upload. So we run a periodic job `cleanup_stale_uploads` to clean up these keys. This job runs every `CLEANUP_STALE_UPLOADS_IN_REDIS_PERIODIC_INTERVAL` minutes and deletes `upload_*` keys if the timestamp difference is greater than 60 minutes.
  
- Moved task server configs to TASK_SERVER_V2.
- Moved reserve, release upload space methods to filebrowser utils. New configs for periodic scheduling under task server
- Display max_file_upload_size_limit on upload modal. Parse redis broker url from configs.
- Changed task server configuration parsing from awk to grep. Set autorestart to false in redis and celery template.
- Show task_server tab in Admin Server, based on task_server_v2 configs
- Moving Uploaded chunk log message to Debug from Info
- Setting celery default log level to info
- Pull timezone from hue.ini. /tmp_cleaner job checks the timestamp of each file and deletes it based on timedelta=60mins

## How was this patch tested?
Tested on local machine and using docker builds. 

Change-Id: Id167701873d10426c7f6e5064b3feb731966b2a7

Co-authored-by: Athithyaa Selvam <[email protected]>

Author: athithyaaselvam

apps/filebrowser/src/filebrowser/conf.py

@@ -18,9 +18,8 @@
 import os
 import sys
 
-from desktop.conf import ENABLE_DOWNLOAD
+from desktop.conf import ENABLE_DOWNLOAD, is_oozie_enabled
 from desktop.lib.conf import Config, coerce_bool
-from desktop.conf import is_oozie_enabled
 
 if sys.version_info[0] > 2:
   from django.utils.translation import gettext_lazy as _
@@ -31,7 +30,7 @@
   key="max_snappy_decompression_size",
   help=_("Max snappy decompression size in bytes."),
   private=True,
-  default=1024*1024*25,
+  default=1024 * 1024 * 25,
   type=int)
 
 ARCHIVE_UPLOAD_TEMPDIR = Config(
@@ -52,10 +51,12 @@
   type=int,
   help=_('Configure the maximum number of concurrent connections(chunks) for file uploads using the chunked file uploader.'))
 
+
 def get_desktop_enable_download():
   """Get desktop enable_download default"""
   return ENABLE_DOWNLOAD.get()
 
+
 SHOW_DOWNLOAD_BUTTON = Config(
   key="show_download_button",
   help=_("whether to show the download button in hdfs file browser."),
@@ -78,7 +79,7 @@ def get_desktop_enable_download():
 
 REDIRECT_DOWNLOAD = Config(
   key="redirect_download",
-  help=_("Redirect client to WebHdfs or S3 for file download. Note: Turning this on will "\
+  help=_("Redirect client to WebHdfs or S3 for file download. Note: Turning this on will "
     "override notebook/redirect_whitelist for user selected file downloads on WebHdfs & S3."),
   type=coerce_bool,
   default=False)
@@ -96,6 +97,11 @@ def get_desktop_enable_download():
   help=_('A limit on a file size (bytes) that can be uploaded to a filesystem. '
           'A value of -1 means there will be no limit.'))
 
+
+def max_file_size_upload_limit():
+  return MAX_FILE_SIZE_UPLOAD_LIMIT.get()
+
+
 FILE_DOWNLOAD_CACHE_CONTROL = Config(
   key="file_download_cache_control",
   type=str,

apps/filebrowser/src/filebrowser/tasks.py

@@ -14,26 +14,29 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import logging
-import subprocess
 import os
-
 import shutil
-import psutil
+import logging
+import subprocess
 from datetime import datetime, timedelta
 
-from celery.utils.log import get_task_logger
+import pytz
+import psutil
 from celery import states
-from django.db import transaction
+from celery.utils.log import get_task_logger
 from django.http import HttpRequest
 from django.utils import timezone
 
 from desktop.auth.backend import rewrite_user
 from desktop.celery import app
-from desktop.conf import TASK_SERVER
+from desktop.conf import TASK_SERVER_V2
 from desktop.lib import fsmanager
-from useradmin.models import User
+from filebrowser.utils import release_reserved_space_for_file_uploads, reserve_space_for_file_uploads
+
+if hasattr(TASK_SERVER_V2, 'get') and TASK_SERVER_V2.ENABLED.get():
+  from desktop.settings import TIME_ZONE, initialize_free_disk_space_in_redis, parse_broker_url
 from filebrowser.views import UPLOAD_CLASSES
+from useradmin.models import User
 
 LOG_TASK = get_task_logger(__name__)
 LOG = logging.getLogger()
@@ -52,24 +55,34 @@
   states.IGNORED: 'ignored'
 }
 
+
 @app.task
 def error_handler(request, exc, traceback):
   print('Task {0} raised exception: {1!r}\n{2!r}'.format(
     request.id, exc, traceback))
 
+
 @app.task()
 def upload_file_task(**kwargs):
   task_id = kwargs.get("qquuid")
+  file_size = kwargs.get("qqtotalfilesize")
   user_id = kwargs["user_id"]
   scheme = kwargs["scheme"]
   postdict = kwargs.get("postdict", None)
   request = _get_request(postdict=postdict, user_id=user_id, scheme=scheme)
   kwargs["username"] = request.user.username
   kwargs["task_name"] = "fileupload"
   kwargs["state"] = "STARTED"
-  now = timezone.now()
-  kwargs["task_start"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_start"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["started"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S")
   upload_file_task.update_state(task_id=task_id, state='STARTED', meta=kwargs)
+
+  # Reserve space for upload
+  if not reserve_space_for_file_uploads(task_id, file_size):
+    kwargs["state"] = "FAILURE"
+    upload_file_task.update_state(task_id=task_id, state='FAILURE', meta=kwargs)
+    raise Exception("Insufficient space for upload")
+
   try:
     upload_class = UPLOAD_CLASSES.get(kwargs["scheme"])
     _fs = upload_class(request, args=[], **kwargs)
@@ -82,11 +95,13 @@ def upload_file_task(**kwargs):
     raise Exception(f"Upload failed %s" % err)
 
   kwargs["state"] = "SUCCESS"
-  kwargs["started"] = now.strftime("%Y-%m-%dT%H:%M:%S")
-  kwargs["task_end"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+  kwargs["progress"] = "100%"
   upload_file_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
+  release_reserved_space_for_file_uploads(task_id)
   return None
 
+
 def _get_request(postdict=None, user_id=None, scheme=None):
   request = HttpRequest()
   request.POST = postdict
@@ -100,19 +115,19 @@ def _get_request(postdict=None, user_id=None, scheme=None):
 
   return request
 
+
 @app.task()
 def document_cleanup_task(**kwargs):
   keep_days = kwargs.get('keep_days')
   task_id = kwargs.get("qquuid")
-  now = timezone.now()
 
   kwargs["username"] = kwargs["username"]
   kwargs["task_name"] = "document_cleanup"
   kwargs["state"] = "STARTED"
   kwargs["parameters"] = keep_days
   kwargs["task_id"] = task_id
   kwargs["progress"] = "0%"
-  kwargs["task_start"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_start"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S")
 
   try:
     INSTALL_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..'))
@@ -121,63 +136,109 @@ def document_cleanup_task(**kwargs):
     kwargs["state"] = "SUCCESS"
     LOG.info(f"Document_cleanup_task completed successfully.")
   except Exception as err:
+    kwargs["state"] = "FAILURE"
     document_cleanup_task.update_state(task_id=task_id, state='FAILURE', meta=kwargs)
     raise Exception(f"Upload failed %s" % err)
 
   kwargs["state"] = "SUCCESS"
   kwargs["progress"] = "100%"
-  kwargs["task_end"] = now.strftime("%Y-%m-%dT%H:%M:%S")
+  kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
   document_cleanup_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
 
   return None
 
+
 @app.task()
 def check_disk_usage_and_clean_task(**kwargs):
   task_id = kwargs.get("qquuid")
-  default_cleanup_threshold = 90
-  cleanup_threshold = kwargs.get('cleanup_threshold', default_cleanup_threshold)
+  cleanup_threshold = kwargs.get('cleanup_threshold', TASK_SERVER_V2.DISK_USAGE_CLEANUP_THRESHOLD)
   username = kwargs.get("username", "celery_scheduler")
-  now = timezone.now()
+  now = timezone.now().astimezone(pytz.timezone(TIME_ZONE))
   kwargs = {
     "username": username,
     "task_name": "tmp_cleanup",
     "task_id": task_id,
     "parameters": cleanup_threshold,
     "progress": "0%",
-    "task_start": now.strftime("%Y-%m-%dT%H:%M:%S")
+    "task_start": now.strftime("%Y-%m-%dT%H:%M:%S.%f"),
   }
 
+  check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.STARTED, meta=kwargs)
+
+  def delete_old_files(directory, current_time, time_delta):
+    for root, dirs, files in os.walk(directory):
+      for file in files:
+        file_path = os.path.join(root, file)
+        try:
+          file_modified_time = datetime.fromtimestamp(os.path.getmtime(file_path))
+          file_modified_time = timezone.make_aware(file_modified_time, timezone.get_default_timezone())
+          if file_modified_time < current_time - time_delta:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+              os.unlink(file_path)
+              LOG.info(f"Deleted file {file_path}")
+        except PermissionError:
+          LOG.warning(f"Permission denied: unable to delete {file_path}")
+        except Exception as err:
+          check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.FAILURE, meta=kwargs)
+          raise Exception(f"Failed to delete {file_path}. Reason: {err}")
+        finally:
+          kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+
   disk_usage = psutil.disk_usage('/')
   if disk_usage.percent >= int(cleanup_threshold):
     LOG.info(f"Disk usage is above {cleanup_threshold}%, cleaning up /tmp directory...")
     tmp_dir = '/tmp'
-    for filename in os.listdir(tmp_dir):
-      file_path = os.path.join(tmp_dir, filename)
-      try:
-        file_modified_time = datetime.fromtimestamp(os.path.getmtime(file_path))
-        # Make file_modified_time timezone-aware
-        file_modified_time = timezone.make_aware(file_modified_time, timezone.get_default_timezone())
-        if file_modified_time < now - timedelta(minutes=15):
-          if os.path.isfile(file_path) or os.path.islink(file_path):
-            os.unlink(file_path)
-          elif os.path.isdir(file_path):
-            shutil.rmtree(file_path)
-          LOG.info(f"Deleted {file_path}")
-      except PermissionError:
-        LOG.warning(f"Permission denied: unable to delete {file_path}")
-      except Exception as err:
-        check_disk_usage_and_clean_task.update_state(task_id=task_id, state='FAILURE', meta=kwargs)
-        raise Exception(f"Failed to delete {file_path}. Reason: {err}")
-
+    delete_old_files(tmp_dir, now, timedelta(minutes=TASK_SERVER_V2.DISK_USAGE_AND_CLEAN_TASK_TIME_DELTA.get()))
     kwargs["progress"] = "100%"
-    check_disk_usage_and_clean_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
+    kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+    check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.SUCCESS, meta=kwargs)
     LOG.info("/tmp directory cleaned.")
-
   else:
     kwargs["progress"] = "100%"
-    check_disk_usage_and_clean_task.update_state(task_id=task_id, state='SUCCESS', meta=kwargs)
+    kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+    check_disk_usage_and_clean_task.update_state(task_id=task_id, state=states.SUCCESS, meta=kwargs)
     LOG.info(f"Disk usage is {disk_usage.percent}%, no need to clean up.")
 
   # Get available disk space after cleanup
   free_space = psutil.disk_usage('/tmp').free
   return {'free_space': free_space}
+
+
+@app.task()
+def cleanup_stale_uploads(**kwargs):
+  timedelta_minutes = kwargs.get("timeout_minutes", TASK_SERVER_V2.CLEANUP_STALE_UPLOADS_TASK_TIME_DELTA.get())
+  username = kwargs.get("username", "celery_scheduler")
+  task_id = kwargs.get("qquuid")
+  kwargs = {
+    "username": username,
+    "task_name": "cleanup_stale_uploads",
+    "task_id": task_id,
+    "parameters": timedelta_minutes,
+    "progress": "0%",
+    "task_start": timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S"),
+  }
+  cleanup_stale_uploads.update_state(task_id=task_id, state=states.STARTED, meta=kwargs)
+  redis_client = parse_broker_url(TASK_SERVER_V2.BROKER_URL.get())
+  try:
+    current_time = int(datetime.now().timestamp())
+
+    for key in redis_client.scan_iter('upload__*'):
+      timestamp_key = f'{key.decode()}_timestamp'
+      timestamp = redis_client.get(timestamp_key)
+
+      if timestamp:
+        timestamp = int(timestamp)
+        if current_time - timestamp > timedelta_minutes * 60:
+          file_size = int(redis_client.get(key))
+          redis_client.incrby('upload_available_space', file_size)
+          redis_client.delete(key)
+          redis_client.delete(timestamp_key)
+    initialize_free_disk_space_in_redis()
+    kwargs["progress"] = "100%"
+    kwargs["task_end"] = timezone.now().astimezone(pytz.timezone(TIME_ZONE)).strftime("%Y-%m-%dT%H:%M:%S.%f")
+    cleanup_stale_uploads.update_state(task_id=task_id, state=states.SUCCESS, meta=kwargs)
+  except Exception as e:
+    cleanup_stale_uploads.update_state(task_id=task_id, state=states.FAILURE, meta={"error": str(e)})
+    LOG.exception("Failed to cleanup stale uploads: %s", str(e))
+  finally:
+    redis_client.close()

apps/filebrowser/src/filebrowser/templates/listdir.mako

@@ -168,7 +168,7 @@ ${ fb_components.menubar() }
             <a class="btn fileToolbarBtn" title="${_('Upload files')}" data-bind="visible: !inTrash(), css: {'disabled': isS3Root()}, click: function(){ if (!isS3Root()) { uploadFile(false) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Upload')}</a>
             <!-- /ko -->
             <!-- ko if: isTaskServerEnabled -->
-            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isS3Root()}, click: function(){ if (!isS3Root()) { uploadFile(true) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
+            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isS3Root()}, click: function(){ if (!isS3Root()) { checkAndDisplayAvailableSpace(); uploadFile(true); }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
             <!-- /ko -->
           <!-- /ko -->
           <!-- ko if: isGS -->
@@ -179,7 +179,7 @@ ${ fb_components.menubar() }
             <a class="btn fileToolbarBtn" title="${_('Upload files')}" data-bind="visible: !inTrash(), css: {'disabled': isABFSRoot()}, click: function(){ if (!isABFSRoot()) { uploadFile(false) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Upload')}</a>
             <!-- /ko -->
             <!-- ko if: isTaskServerEnabled -->
-            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isABFSRoot()}, click: function(){ if (!isABFSRoot()) { uploadFile(true) }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
+            <a class="btn fileToolbarBtn" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="visible: !inTrash(), css: {'disabled': isABFSRoot()}, click: function(){ if (!isABFSRoot()) { checkAndDisplayAvailableSpace(); uploadFile(true); }}"><i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}</a>
             <!-- /ko -->
           <!-- /ko -->
           <!-- ko ifnot: isS3() || isGS() || isABFS() -->
@@ -192,7 +192,7 @@ ${ fb_components.menubar() }
           <!-- /ko -->
           <!-- ko if: isTaskServerEnabled -->
           <div id="upload-dropdown" class="btn-group" style="vertical-align: middle">
-            <a data-hue-analytics="filebrowser:upload-btn-click" href="javascript: void(0)" class="btn upload-link dropdown-toggle" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="click: function() { uploadFile(true); }, visible: !inTrash(), css: {'disabled': (isOFS() && (isOFSRoot() || isOFSServiceID() || isOFSVol()))}">
+            <a data-hue-analytics="filebrowser:upload-btn-click" href="javascript: void(0)" class="btn upload-link dropdown-toggle" title="${_('Select a file to upload. The file will first be saved locally and then automatically transferred to the designated file system (e.g., S3, Azure) in the background. The upload modal closes immediately after the file is queued, allowing you to continue working. A notification, \'File upload scheduled. Please check the task server page for progress,\' will confirm the upload has started. This feature is especially useful for large files, as it eliminates the need to wait for the upload to complete.')}" data-bind="click: function() { checkAndDisplayAvailableSpace(); uploadFile(true);}, visible: !inTrash(), css: {'disabled': (isOFS() && (isOFSRoot() || isOFSServiceID() || isOFSVol()))}">
               <i class="fa fa-arrow-circle-o-up"></i> ${_('Schedule Upload')}
             </a>
           </div>