Keep controller connection alive while uploading files (dotnet#655)

src/Microsoft.Crank.Agent/Controllers/JobsController.cs

@@ -325,7 +325,15 @@ public async Task<IActionResult> UploadAttachment(int id)
 
             var tempFilename = Path.GetTempFileName();
 
-            await SaveBodyAsync(tempFilename);
+            try
+            {
+                job.LastDriverCommunicationUtc = DateTime.UtcNow;
+                await SaveBodyAsync(job, tempFilename);
+            }
+            catch
+            {
+                return StatusCode(500, "Error while saving file");
+            }
 
             job.Attachments.Add(new Attachment
             {
@@ -367,15 +375,18 @@ public async Task<IActionResult> UploadAttachmentZip(int id)
 
             var tempFilename = Path.GetTempFileName() + ".zip";
 
-            await SaveBodyAsync(tempFilename);
-
-            job.LastDriverCommunicationUtc = DateTime.UtcNow;
-
             try
             {
+                job.LastDriverCommunicationUtc = DateTime.UtcNow;
+                await SaveBodyAsync(job, tempFilename);
+
                 // Extract the zip file in a temporary folder
                 ZipFile.ExtractToDirectory(tempFilename, destinationTempFilename);
             }
+            catch
+            {
+                return StatusCode(500, "Error while saving file");
+            }
             finally
             {
                 System.IO.File.Delete(tempFilename);
@@ -416,7 +427,16 @@ public async Task<IActionResult> UploadSource(int id, string sourceName = Source
                 return NotFound("Source does not exist on job");
 
             var tempFilename = Path.GetTempFileName();
-            await SaveBodyAsync(tempFilename);
+
+            try
+            {
+                job.LastDriverCommunicationUtc = DateTime.UtcNow;
+                await SaveBodyAsync(job, tempFilename);
+            }
+            catch
+            {
+                return StatusCode(500, "Error while saving file");
+            }
 
             source.SourceCode = new Attachment
             {
@@ -441,7 +461,15 @@ public async Task<IActionResult> UploadBuildFile(int id)
             var job = _jobs.Find(id);
             var tempFilename = Path.GetTempFileName();
 
-            await SaveBodyAsync(tempFilename);
+            try
+            {
+                job.LastDriverCommunicationUtc = DateTime.UtcNow;
+                await SaveBodyAsync(job, tempFilename);
+            }
+            catch
+            {
+                return StatusCode(500, "Error while saving file");
+            }
 
             job.BuildAttachments.Add(new Attachment
             {
@@ -454,21 +482,34 @@ public async Task<IActionResult> UploadBuildFile(int id)
             return Ok();
         }
 
-        private async Task SaveBodyAsync(string filename)
+        private async Task SaveBodyAsync(Job job, string filename)
         {
             using var outputFileStream = System.IO.File.Create(filename);
 
+            Task task = null;
+
             if (Request.Headers.TryGetValue("Content-Encoding", out var encoding) && encoding.Contains("gzip"))
             {
                 Log.Info($"Received gzipped file content");
                 using var decompressor = new GZipStream(Request.Body, CompressionMode.Decompress);
-                await decompressor.CopyToAsync(outputFileStream, Request.HttpContext.RequestAborted);
+                task = decompressor.CopyToAsync(outputFileStream, Request.HttpContext.RequestAborted);
             }
             else
             {
                 Log.Info($"Received uncompressed file content");
-                await Request.Body.CopyToAsync(outputFileStream, Request.HttpContext.RequestAborted);
+                task = Request.Body.CopyToAsync(outputFileStream, Request.HttpContext.RequestAborted);
             }
+
+            // Keeps controller's connection alive as PINGs might not be accepted
+            // by the agent while big files are uploaded.
+
+            while (!Request.HttpContext.RequestAborted.IsCancellationRequested && !task.IsCompleted)
+            {
+                await Task.Delay(500);
+                job.LastDriverCommunicationUtc = DateTime.UtcNow;
+            }
+
+            await task;
         }
 
         [HttpGet("{id}/trace")]