From 79f9f60a78e90970bd33da36e2c1bffc6d2ce867 Mon Sep 17 00:00:00 2001
From: Kyle Spearrin <kyle.spearrin@gmail.com>
Date: Wed, 30 Dec 2015 22:38:00 -0500
Subject: [PATCH] Added `GetTakeCount` helper to break down large arrays into
 manageable sizes for documentdb requests (current max limit of 512kb per
 request).

---
 .../DocumentDB/CipherRepository.cs            | 40 +++++++++----------
 .../DocumentDB/Utilities/DocumentDBHelpers.cs | 22 ++++++++++
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/src/Core/Repositories/DocumentDB/CipherRepository.cs b/src/Core/Repositories/DocumentDB/CipherRepository.cs
index 41d2670ed2..cd698a759d 100644
--- a/src/Core/Repositories/DocumentDB/CipherRepository.cs
+++ b/src/Core/Repositories/DocumentDB/CipherRepository.cs
@@ -34,21 +34,21 @@ namespace Bit.Core.Repositories.DocumentDB
 
         public async Task UpdateDirtyCiphersAsync(IEnumerable<dynamic> ciphers)
         {
+            // Make sure we are dealing with cipher types since we accept any via dynamic.
+            var cleanedCiphers = ciphers.Where(c => c is Cipher);
+            if(cleanedCiphers.Count() == 0)
+            {
+                return;
+            }
+
+            var takeCount = DocumentDBHelpers.GetTakeCount(ciphers, 500);
+
             await DocumentDBHelpers.ExecuteWithRetryAsync(async () =>
             {
-                // Make sure we are dealing with cipher types since we accept any via dynamic.
-                var cleanedCiphers = ciphers.Where(c => c is Cipher);
-                if(cleanedCiphers.Count() == 0)
-                {
-                    return;
-                }
-
                 var userId = ((Cipher)cleanedCiphers.First()).UserId;
                 StoredProcedureResponse<int> sprocResponse = await Client.ExecuteStoredProcedureAsync<int>(
                     ResolveSprocIdLink(userId, "updateDirtyCiphers"),
-                    // TODO: Figure out how to better determine the max number of document to send without
-                    // going over 512kb limit for DocumentDB. 50 could still be too large in some cases.
-                    cleanedCiphers.Take(50),
+                    cleanedCiphers.Take(takeCount),
                     userId);
 
                 var replacedCount = sprocResponse.Response;
@@ -61,21 +61,21 @@ namespace Bit.Core.Repositories.DocumentDB
 
         public async Task CreateAsync(IEnumerable<dynamic> ciphers)
         {
+            // Make sure we are dealing with cipher types since we accept any via dynamic.
+            var cleanedCiphers = ciphers.Where(c => c is Cipher);
+            if(cleanedCiphers.Count() == 0)
+            {
+                return;
+            }
+
+            var takeCount = DocumentDBHelpers.GetTakeCount(ciphers, 500);
+
             await DocumentDBHelpers.ExecuteWithRetryAsync(async () =>
             {
-                // Make sure we are dealing with cipher types since we accept any via dynamic.
-                var cleanedCiphers = ciphers.Where(c => c is Cipher);
-                if(cleanedCiphers.Count() == 0)
-                {
-                    return;
-                }
-
                 var userId = ((Cipher)cleanedCiphers.First()).UserId;
                 StoredProcedureResponse<int> sprocResponse = await Client.ExecuteStoredProcedureAsync<int>(
                     ResolveSprocIdLink(userId, "bulkCreate"),
-                    // TODO: Figure out how to better determine the max number of document to send without
-                    // going over 512kb limit for DocumentDB. 50 could still be too large in some cases.
-                    cleanedCiphers.Take(50));
+                    cleanedCiphers.Take(takeCount));
 
                 var createdCount = sprocResponse.Response;
                 if(createdCount != cleanedCiphers.Count())
diff --git a/src/Core/Repositories/DocumentDB/Utilities/DocumentDBHelpers.cs b/src/Core/Repositories/DocumentDB/Utilities/DocumentDBHelpers.cs
index 5bc8ca9928..db8abbb8a1 100644
--- a/src/Core/Repositories/DocumentDB/Utilities/DocumentDBHelpers.cs
+++ b/src/Core/Repositories/DocumentDB/Utilities/DocumentDBHelpers.cs
@@ -1,7 +1,11 @@
 using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
 using System.Threading.Tasks;
 using Microsoft.Azure.Documents;
 using Microsoft.Azure.Documents.Client;
+using Newtonsoft.Json;
 
 namespace Bit.Core.Repositories.DocumentDB.Utilities
 {
@@ -58,6 +62,24 @@ namespace Bit.Core.Repositories.DocumentDB.Utilities
             }
         }
 
+        public static int GetTakeCount(IEnumerable<object> docs, int maxSizeKb = 500)
+        {
+            var takeCount = docs.Count();
+            while(takeCount > 1)
+            {
+                var bytes = Encoding.UTF8.GetBytes(JsonConvert.SerializeObject(docs.Take(takeCount)));
+                if((bytes.Length / 1000) <= maxSizeKb)
+                {
+                    // array is is small enough
+                    break;
+                }
+
+                takeCount = Convert.ToInt32(Math.Ceiling((double)takeCount / 2));
+            }
+
+            return takeCount;
+        }
+
         private static async Task HandleDocumentClientExceptionAsync(DocumentClientException e, int retryCount, int? retryMax)
         {
             if(retryMax.HasValue && retryCount >= retryMax.Value)