From 3ef37728ba409b95fd010c45599187365d26f8f1 Mon Sep 17 00:00:00 2001
From: yahavi <yahavi@jfrog.com>
Date: Thu, 7 Mar 2024 12:26:11 +0200
Subject: [PATCH] Retry completion on status 503 in the polling

---
 artifactory/services/utils/multipartupload.go |  5 ++++
 .../services/utils/multipartupload_test.go    | 24 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/artifactory/services/utils/multipartupload.go b/artifactory/services/utils/multipartupload.go
index d691463d1..e44750045 100644
--- a/artifactory/services/utils/multipartupload.go
+++ b/artifactory/services/utils/multipartupload.go
@@ -369,6 +369,11 @@ func (mu *MultipartUpload) completeMultipartUpload(logMsgPrefix, sha1 string, mu
 func (mu *MultipartUpload) status(logMsgPrefix string, multipartUploadClientWithNodeId *httputils.HttpClientDetails) (status statusResponse, err error) {
 	url := fmt.Sprintf("%s%sstatus", mu.artifactoryUrl, uploadsApi)
 	resp, body, err := mu.client.GetHttpClient().SendPost(url, []byte{}, *multipartUploadClientWithNodeId, logMsgPrefix)
+	// If the Artifactory node returns a "Service unavailable" error (status 503), attempt to retry the upload completion process on a different node.
+	if resp != nil && resp.StatusCode == http.StatusServiceUnavailable {
+		log.Warn(logMsgPrefix + fmt.Sprintf("The Artifactory node ID %s is unavailable. Trying to finalize the multipart upload process on a different node.", multipartUploadClientWithNodeId.Headers[routeToHeader]))
+		return statusResponse{Status: retryableError, Error: string(body)}, nil
+	}
 	if err != nil {
 		return
 	}
diff --git a/artifactory/services/utils/multipartupload_test.go b/artifactory/services/utils/multipartupload_test.go
index 0ee4e552b..6669f2744 100644
--- a/artifactory/services/utils/multipartupload_test.go
+++ b/artifactory/services/utils/multipartupload_test.go
@@ -233,6 +233,30 @@ func TestStatus(t *testing.T) {
 	assert.Equal(t, statusResponse{Status: finished, Progress: utils.Pointer(100)}, status)
 }
 
+func TestStatusServiceUnavailable(t *testing.T) {
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Check method
+		assert.Equal(t, http.MethodPost, r.Method)
+
+		// Check URL
+		assert.Equal(t, "/api/v1/uploads/status", r.URL.Path)
+
+		// Send response 503 Service unavailable
+		w.WriteHeader(http.StatusServiceUnavailable)
+		_, err := w.Write([]byte("Service unavailable"))
+		assert.NoError(t, err)
+	})
+
+	// Create mock multipart upload with server
+	multipartUpload, cleanUp := createMockMultipartUpload(t, handler)
+	defer cleanUp()
+
+	// Execute status
+	status, err := multipartUpload.status("", &httputils.HttpClientDetails{})
+	assert.NoError(t, err)
+	assert.Equal(t, statusResponse{Status: retryableError, Error: "Service unavailable"}, status)
+}
+
 func TestAbort(t *testing.T) {
 	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// Check method