Support for SecNetPerf Variables to be in Units of CPU Count (#4588)

src/perf/lib/PerfClient.cpp

@@ -19,18 +19,21 @@ const char* TimeUnits[] = { "m", "ms", "us", "s" };
 const uint64_t TimeMult[] = { 60 * 1000 * 1000, 1000, 1, 1000 * 1000 };
 const char* SizeUnits[] = { "gb", "mb", "kb", "b" };
 const uint64_t SizeMult[] = { 1000 * 1000 * 1000, 1000 * 1000, 1000, 1 };
+const char* CountUnits[] = { "cpu" };
+uint64_t CountMult[] = { 1 };
 
 _Success_(return != false)
+template <typename T>
 bool
 TryGetVariableUnitValue(
     _In_ int argc,
     _In_reads_(argc) _Null_terminated_ char* argv[],
     _In_z_ const char** names,
-    _Out_ uint64_t* pValue,
-    _Out_ bool* isTimed
+    _Out_ T* pValue,
+    _Out_opt_ bool* isTimed = nullptr
     )
 {
-    *isTimed = false; // Default
+    if (isTimed) *isTimed = false; // Default
 
     // Search for the first matching name.
     char* value = nullptr;
@@ -44,9 +47,9 @@ TryGetVariableUnitValue(
         size_t len = strlen(TimeUnits[i]);
         if (len < strlen(value) &&
             _strnicmp(value + strlen(value) - len, TimeUnits[i], len) == 0) {
-            *isTimed = true;
+            if (isTimed) *isTimed = true;
             value[strlen(value) - len] = '\0';
-            *pValue = (uint64_t)atoi(value) * TimeMult[i];
+            *pValue = (T)(atoi(value) * TimeMult[i]);
             return true;
         }
     }
@@ -57,16 +60,42 @@ TryGetVariableUnitValue(
         if (len < strlen(value) &&
             _strnicmp(value + strlen(value) - len, SizeUnits[i], len) == 0) {
             value[strlen(value) - len] = '\0';
-            *pValue = (uint64_t)atoi(value) * SizeMult[i];
+            *pValue = (T)(atoi(value) * SizeMult[i]);
+            return true;
+        }
+    }
+
+    // Search to see if the value has a count unit specified at the end.
+    for (uint32_t i = 0; i < ARRAYSIZE(CountUnits); ++i) {
+        size_t len = strlen(CountUnits[i]);
+        if (len < strlen(value) &&
+            _strnicmp(value + strlen(value) - len, CountUnits[i], len) == 0) {
+            value[strlen(value) - len] = '\0';
+            *pValue = (T)(atoi(value) * CountMult[i]);
             return true;
         }
     }
 
     // Default to bytes if no unit is specified.
-    *pValue = (uint64_t)atoi(value);
+    *pValue = (T)atoi(value);
     return true;
 }
 
+_Success_(return != false)
+template <typename T>
+bool
+TryGetVariableUnitValue(
+    _In_ int argc,
+    _In_reads_(argc) _Null_terminated_ char* argv[],
+    _In_z_ const char* name,
+    _Out_ T* pValue,
+    _Out_opt_ bool* isTimed = nullptr
+    )
+{
+    const char* names[] = { name, nullptr };
+    return TryGetVariableUnitValue(argc, argv, names, pValue, isTimed);
+}
+
 QUIC_STATUS
 PerfClient::Init(
     _In_ int argc,
@@ -77,6 +106,8 @@ PerfClient::Init(
         return Configuration.GetInitStatus();
     }
 
+    CountMult[0] = CxPlatProcCount();
+
     //
     // Remote target/server options
     //
@@ -113,8 +144,8 @@ PerfClient::Init(
     //
 
     WorkerCount = CxPlatProcCount();
-    TryGetValue(argc, argv, "threads", &WorkerCount);
-    TryGetValue(argc, argv, "workers", &WorkerCount);
+    TryGetVariableUnitValue(argc, argv, "threads", &WorkerCount);
+    TryGetVariableUnitValue(argc, argv, "workers", &WorkerCount);
     TryGetValue(argc, argv, "affinitize", &AffinitizeWorkers);
 
 #ifdef QUIC_COMPARTMENT_ID
@@ -165,9 +196,9 @@ PerfClient::Init(
     // Scenario options
     //
 
-    TryGetValue(argc, argv, "conns", &ConnectionCount);
-    TryGetValue(argc, argv, "requests", &StreamCount);
-    TryGetValue(argc, argv, "streams", &StreamCount);
+    TryGetVariableUnitValue(argc, argv, "conns", &ConnectionCount);
+    TryGetVariableUnitValue(argc, argv, "requests", &StreamCount);
+    TryGetVariableUnitValue(argc, argv, "streams", &StreamCount);
     TryGetValue(argc, argv, "iosize", &IoSize);
     if (IoSize < 256) {
         WriteOutput("'iosize' too small'!\n");