1 files changed, 191 insertions, 0 deletions

diff --git a/src/balutil/balretry.cpp b/src/balutil/balretry.cpp
new file mode 100644
index 00000000..d95d86b2
--- /dev/null
+++ b/src/balutil/balretry.cpp
@@ -0,0 +1,191 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+#include "precomp.h"
+
+struct BALRETRY_INFO
+{
+    LPWSTR sczId;           // package or container id.
+    LPWSTR sczPayloadId;    // optional payload id.
+    DWORD cRetries;
+    DWORD dwLastError;
+};
+
+static DWORD vdwMaxRetries = 0;
+static DWORD vdwTimeout = 0;
+static BALRETRY_INFO vrgRetryInfo[2];
+
+// prototypes
+static BOOL IsActiveRetryEntry(
+    __in BALRETRY_TYPE type,
+    __in_z LPCWSTR wzPackageId,
+    __in_z_opt LPCWSTR wzPayloadId
+    );
+
+
+DAPI_(void) BalRetryInitialize(
+    __in DWORD dwMaxRetries,
+    __in DWORD dwTimeout
+    )
+{
+    BalRetryUninitialize(); // clean everything out.
+
+    vdwMaxRetries = dwMaxRetries;
+    vdwTimeout = dwTimeout;
+}
+
+
+DAPI_(void) BalRetryUninitialize()
+{
+    for (DWORD i = 0; i < countof(vrgRetryInfo); ++i)
+    {
+        ReleaseStr(vrgRetryInfo[i].sczId);
+        ReleaseStr(vrgRetryInfo[i].sczPayloadId);
+        memset(vrgRetryInfo + i, 0, sizeof(BALRETRY_INFO));
+    }
+
+    vdwMaxRetries = 0;
+    vdwTimeout = 0;
+}
+
+
+DAPI_(void) BalRetryStartPackage(
+    __in BALRETRY_TYPE type,
+    __in_z_opt LPCWSTR wzPackageId,
+    __in_z_opt LPCWSTR wzPayloadId
+    )
+{
+    if (!wzPackageId || !*wzPackageId)
+    {
+        ReleaseNullStr(vrgRetryInfo[type].sczId);
+        ReleaseNullStr(vrgRetryInfo[type].sczPayloadId);
+    }
+    else if (IsActiveRetryEntry(type, wzPackageId, wzPayloadId))
+    {
+        ++vrgRetryInfo[type].cRetries;
+        ::Sleep(vdwTimeout);
+    }
+    else
+    {
+        StrAllocString(&vrgRetryInfo[type].sczId, wzPackageId, 0);
+        if (wzPayloadId)
+        {
+            StrAllocString(&vrgRetryInfo[type].sczPayloadId, wzPayloadId, 0);
+        }
+
+        vrgRetryInfo[type].cRetries = 0;
+    }
+
+    vrgRetryInfo[type].dwLastError = ERROR_SUCCESS;
+}
+
+
+DAPI_(void) BalRetryErrorOccurred(
+    __in_z LPCWSTR wzPackageId,
+    __in DWORD dwError
+    )
+{
+    if (IsActiveRetryEntry(BALRETRY_TYPE_CACHE, wzPackageId, NULL))
+    {
+        vrgRetryInfo[BALRETRY_TYPE_CACHE].dwLastError = dwError;
+    }
+    else if (IsActiveRetryEntry(BALRETRY_TYPE_EXECUTE, wzPackageId, NULL))
+    {
+        vrgRetryInfo[BALRETRY_TYPE_EXECUTE].dwLastError = dwError;
+    }
+}
+
+
+DAPI_(HRESULT) BalRetryEndPackage(
+    __in BALRETRY_TYPE type,
+    __in_z_opt LPCWSTR wzPackageId,
+    __in_z_opt LPCWSTR wzPayloadId,
+    __in HRESULT hrError,
+    __inout BOOL* pfRetry
+    )
+{
+    HRESULT hr = S_OK;
+
+    if (!wzPackageId || !*wzPackageId)
+    {
+        ReleaseNullStr(vrgRetryInfo[type].sczId);
+        ReleaseNullStr(vrgRetryInfo[type].sczPayloadId);
+    }
+    else if (FAILED(hrError) && vrgRetryInfo[type].cRetries < vdwMaxRetries && IsActiveRetryEntry(type, wzPackageId, wzPayloadId))
+    {
+        if (BALRETRY_TYPE_CACHE == type)
+        {
+            // Retry on all errors except the following.
+            if (HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT) != hrError &&
+                BG_E_NETWORK_DISCONNECTED != hrError &&
+                HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND) != hrError &&
+                HRESULT_FROM_WIN32(ERROR_INTERNET_NAME_NOT_RESOLVED) != hrError)
+            {
+                *pfRetry = TRUE;
+            }
+        }
+        else if (BALRETRY_TYPE_EXECUTE == type)
+        {
+            // If the service is out of whack, just try again.
+            if (HRESULT_FROM_WIN32(ERROR_INSTALL_SERVICE_FAILURE) == hrError)
+            {
+                *pfRetry = TRUE;
+            }
+            else if (HRESULT_FROM_WIN32(ERROR_INSTALL_FAILURE) == hrError)
+            {
+                DWORD dwError = vrgRetryInfo[type].dwLastError;
+
+                // If we failed with one of these specific error codes, then retry since
+                // we've seen these have a high success of succeeding on retry.
+                if (1303 == dwError ||
+                    1304 == dwError ||
+                    1306 == dwError ||
+                    1307 == dwError ||
+                    1309 == dwError ||
+                    1310 == dwError ||
+                    1311 == dwError ||
+                    1312 == dwError ||
+                    1316 == dwError ||
+                    1317 == dwError ||
+                    1321 == dwError ||
+                    1335 == dwError ||
+                    1402 == dwError ||
+                    1406 == dwError ||
+                    1606 == dwError ||
+                    1706 == dwError ||
+                    1719 == dwError ||
+                    1723 == dwError ||
+                    1923 == dwError ||
+                    1931 == dwError)
+                {
+                    *pfRetry = TRUE;
+                }
+            }
+            else if (HRESULT_FROM_WIN32(ERROR_INSTALL_ALREADY_RUNNING) == hrError)
+            {
+                *pfRetry = TRUE;
+            }
+        }
+    }
+
+    return hr;
+}
+
+
+// Internal functions.
+
+static BOOL IsActiveRetryEntry(
+    __in BALRETRY_TYPE type,
+    __in_z LPCWSTR wzPackageId,
+    __in_z_opt LPCWSTR wzPayloadId
+    )
+{
+    BOOL fActive = FALSE;
+
+    fActive = vrgRetryInfo[type].sczId && CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPackageId, -1, vrgRetryInfo[type].sczId, -1);
+    if (fActive && wzPayloadId) // if a payload id was provided ensure it matches.
+    {
+        fActive = vrgRetryInfo[type].sczPayloadId && CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPayloadId, -1, vrgRetryInfo[type].sczPayloadId, -1);
+    }
+
+    return fActive;
+}