Fix latency stats

tcp_client

@@ -53,7 +53,7 @@ if __name__ == '__main__':
         sock.connect(addr)
 
         n = 0
-        latency_stats = np.zeros((timeout,))
+        latency_stats = np.zeros((timeout * 10,))
 
         while time.monotonic() - start < duration:
             req_start = time.monotonic()
@@ -64,12 +64,45 @@ if __name__ == '__main__':
                 if not resp:
                     raise SystemExit()
                 nrecv += len(resp)
-            req_time = round((time.monotonic() - req_start) * 1000)
+            req_time = round((time.monotonic() - req_start) * 10000)
             latency_stats[req_time] += 1
             n += 1
 
         return n, latency_stats
 
+    def weighted_quantile(values, quantiles, sample_weight=None,
+                          values_sorted=False, old_style=False):
+        """ Very close to np.percentile, but supports weights.
+        NOTE: quantiles should be in [0, 1]!
+        :param values: np.array with data
+        :param quantiles: array-like with many quantiles needed
+        :param sample_weight: array-like of the same length as `array`
+        :param values_sorted: bool, if True, then will avoid sorting of initial array
+        :param old_style: if True, will correct output to be consistent with np.percentile.
+        :return: np.array with computed quantiles.
+        """
+        values = np.array(values)
+        quantiles = np.array(quantiles)
+        if sample_weight is None:
+            sample_weight = np.ones(len(values))
+        sample_weight = np.array(sample_weight)
+        assert np.all(quantiles >= 0) and np.all(quantiles <= 1), \
+                     'quantiles should be in [0, 1]'
+
+        if not values_sorted:
+            sorter = np.argsort(values)
+            values = values[sorter]
+            sample_weight = sample_weight[sorter]
+
+        weighted_quantiles = np.cumsum(sample_weight) - 0.5 * sample_weight
+        if old_style:
+            # To be convenient with np.percentile
+            weighted_quantiles -= weighted_quantiles[0]
+            weighted_quantiles /= weighted_quantiles[-1]
+        else:
+            weighted_quantiles /= np.sum(sample_weight)
+        return np.interp(quantiles, weighted_quantiles, values)
+
     TIMES = args.times
     N = args.concurrency
     DURATION = args.duration
@@ -95,26 +128,25 @@ if __name__ == '__main__':
     end = time.monotonic()
     duration = end - start
 
-    weighted_latency = np.multiply(latency_stats, np.arange(timeout))
-
-    mean_latency = (np.sum(weighted_latency) / timeout)
+    arange = np.arange(len(latency_stats))
 
-    trimmed_latency = np.trim_zeros(latency_stats, 'b')
+    weighted_latency = np.multiply(latency_stats, arange)
+    mean_latency = np.sum(weighted_latency) / messages
 
     percentiles = [50, 75, 90, 99]
     percentile_data = []
 
-    latency_chart = np.stack((np.arange(len(trimmed_latency)),
-                              trimmed_latency), axis=-1)
-    percentile_values = np.percentile(latency_chart, percentiles, axis=0)
+    quantiles = weighted_quantile(arange, [p / 100 for p in percentiles],
+                                  sample_weight=latency_stats,
+                                  values_sorted=True)
 
     for i, percentile in enumerate(percentiles):
         percentile_data.append('{}%: {}ms'.format(
-            percentile, round(percentile_values[i][0], 2)))
+            percentile, round(quantiles[i], 2)))
 
     print(messages, 'in', round(duration, 2))
     print('Latency avg: {}ms'.format(round(mean_latency, 2)))
-    # print('Latency distribution: {}'.format('; '.join(percentile_data)))
+    print('Latency distribution: {}'.format('; '.join(percentile_data)))
     print('Requests/sec: {}'.format(round(messages / duration, 2)))
     transfer = (messages * MSGSIZE / (1024 * 1024)) / duration
     print('Transfer/sec: {}MiB'.format(round(transfer, 2)))