Driver: Optimize processes handling.

src/driver/fortdrv.c

@@ -763,7 +763,7 @@ fort_callout_timer (void)
   /* Flush traffic statistics */
   if (stat->is_dirty) {
     PCHAR out;
-    const UINT16 proc_count = stat->proc_count;
+    const UINT16 proc_count = fort_stat_proc_count(stat);
     const UINT32 len = FORT_LOG_STAT_SIZE(proc_count);
 
     /* TODO: Write by chunks */

src/driver/fortstat.c

@@ -5,27 +5,44 @@
 #define FORT_PROC_BAD_INDEX	((UINT16) -1)
 
 typedef struct fort_stat_traf {
+  union {
+    struct {
       UINT32 in_bytes;
       UINT32 out_bytes;
+    };
+
+    UINT64 v;
+  };
 } FORT_STAT_TRAF, *PFORT_STAT_TRAF;
 
 typedef struct fort_stat_group {
-  union {
-    LARGE_INTEGER traf_all;
   FORT_STAT_TRAF traf;
-  };
 } FORT_STAT_GROUP, *PFORT_STAT_GROUP;
 
+/* Synchronize with tommy_hashdyn_node! */
 typedef struct fort_stat_proc {
-  UINT16 next_index;
+  struct fort_stat_proc *next;
-  UINT16 refcount;
+  struct fort_stat_proc *prev;
-
-  UINT32 process_id;
 
   union {
-    LARGE_INTEGER traf_all;
+#if defined(_WIN64)
     FORT_STAT_TRAF traf;
+#else
+    UINT32 process_id;
+#endif
+    void *data;
   };
+
+  tommy_key_t proc_hash;
+
+  UINT16 proc_index;
+  UINT16 refcount;
+
+#if defined(_WIN64)
+  UINT32 process_id;
+#else
+  FORT_STAT_TRAF traf;
+#endif
 } FORT_STAT_PROC, *PFORT_STAT_PROC;
 
 #define FORT_STAT_FLOW_SPEED_LIMIT	0x01
@@ -69,22 +86,17 @@ typedef struct fort_stat {
 
   UINT16 limit_bits;
 
-  UINT16 proc_count;
-  UINT16 proc_top;
-  UINT16 proc_end;
-
-  UINT16 proc_head_index;
-  UINT16 proc_free_index;
-
   UINT32 stream4_id;
   UINT32 datagram4_id;
   UINT32 in_transport4_id;
   UINT32 out_transport4_id;
 
-  PFORT_STAT_PROC procs;
+  PFORT_STAT_PROC proc_free;
-
   PFORT_STAT_FLOW flow_free;
 
+  tommy_arrayof procs;
+  tommy_hashdyn procs_map;
+
   tommy_arrayof flows;
   tommy_hashdyn flows_map;
 
@@ -94,90 +106,44 @@ typedef struct fort_stat {
   KSPIN_LOCK lock;
 } FORT_STAT, *PFORT_STAT;
 
+#define fort_stat_proc_hash(process_id)	tommy_inthash_u32((UINT32) (process_id))
 #define fort_stat_flow_hash(flow_id)	tommy_inthash_u32((UINT32) (flow_id))
 
+#define fort_stat_proc_count(stat)	(UINT16) tommy_hashdyn_count(&(stat)->procs_map)
+
 #define fort_stat_group_speed_limit(stat, group_index) \
   ((stat)->limit_bits & (1 << (group_index)))
 
 
-static PVOID
+static PFORT_STAT_PROC
-fort_stat_array_new (SIZE_T size)
+fort_stat_proc_get (PFORT_STAT stat, UINT32 process_id, tommy_key_t proc_hash)
 {
-  return fort_mem_alloc(size, FORT_STAT_POOL_TAG);
+  PFORT_STAT_PROC proc = (PFORT_STAT_PROC) tommy_hashdyn_bucket(
+    &stat->procs_map, proc_hash);
+
+  while (proc != NULL) {
+    if (proc->proc_hash == proc_hash && proc->process_id == process_id)
+      return proc;
+
+    proc = proc->next;
+  }
+  return NULL;
 }
 
 static void
-fort_stat_array_del (PVOID p)
+fort_stat_proc_free (PFORT_STAT stat, PFORT_STAT_PROC proc)
 {
-  fort_mem_free(p, FORT_STAT_POOL_TAG);
+  tommy_hashdyn_remove_existing(&stat->procs_map, (tommy_hashdyn_node *) proc);
-}
 
-static UINT16
-fort_stat_proc_index (PFORT_STAT stat, UINT32 process_id)
-{
-  UINT16 proc_index = stat->proc_head_index;
-
-  while (proc_index != FORT_PROC_BAD_INDEX) {
-    PFORT_STAT_PROC proc = &stat->procs[proc_index];
-
-    if (process_id == proc->process_id)
-      break;
-
-    proc_index = proc->next_index;
-  }
-
-  return proc_index;
-}
-
-static void
-fort_stat_proc_exclude (PFORT_STAT stat, PFORT_STAT_PROC ex_proc,
-                        UINT16 ex_proc_index)
-{
-  UINT16 proc_index = stat->proc_head_index;
-
-  if (proc_index == ex_proc_index) {
-    stat->proc_head_index = ex_proc->next_index;
-  } else {
-    do {
-      PFORT_STAT_PROC proc = &stat->procs[proc_index];
-
-      proc_index = proc->next_index;
-
-      if (proc_index == ex_proc_index) {
-        proc->next_index = ex_proc->next_index;
-        return;
-      }
-    } while (proc_index != FORT_PROC_BAD_INDEX);
-  }
-}
-
-static void
-fort_stat_proc_free (PFORT_STAT stat, PFORT_STAT_PROC proc, UINT16 proc_index,
-                     PFORT_STAT_PROC prev_proc /* = NULL */)
-{
-  /* Exclude from the active chain */
-  if (prev_proc == NULL) {
-    fort_stat_proc_exclude(stat, proc, proc_index);
-  } else {
-    prev_proc->next_index = proc->next_index;
-  }
-
-  if (proc_index == stat->proc_top - 1) {
-    /* Chop from buffer */
-    stat->proc_top--;
-  } else {
   /* Add to free chain */
-    proc->next_index = stat->proc_free_index;
+  proc->next = stat->proc_free;
-    stat->proc_free_index = proc_index;
+  stat->proc_free = proc;
-  }
-
-  stat->proc_count--;
 }
 
 static void
 fort_stat_proc_inc (PFORT_STAT stat, UINT16 proc_index)
 {
-  PFORT_STAT_PROC proc = &stat->procs[proc_index];
+  PFORT_STAT_PROC proc = tommy_arrayof_ref(&stat->procs, proc_index);
 
   ++proc->refcount;
 }
@@ -185,68 +151,41 @@ fort_stat_proc_inc (PFORT_STAT stat, UINT16 proc_index)
 static void
 fort_stat_proc_dec (PFORT_STAT stat, UINT16 proc_index)
 {
-  PFORT_STAT_PROC proc = &stat->procs[proc_index];
+  PFORT_STAT_PROC proc = tommy_arrayof_ref(&stat->procs, proc_index);
 
   if (!--proc->refcount) {
     stat->is_dirty = TRUE;
   }
 }
 
-static BOOL
+static PFORT_STAT_PROC
-fort_stat_proc_realloc (PFORT_STAT stat)
-{
-  const UINT16 proc_end = stat->proc_end;
-  const UINT16 new_end = (proc_end ? proc_end : 16) * 3 / 2;
-  PFORT_STAT_PROC new_procs = fort_stat_array_new(
-    new_end * sizeof(FORT_STAT_PROC));
-
-  if (new_procs == NULL)
-    return FALSE;
-
-  if (proc_end) {
-    PFORT_STAT_PROC procs = stat->procs;
-
-    RtlCopyMemory(new_procs, procs, stat->proc_top * sizeof(FORT_STAT_PROC));
-    fort_stat_array_del(procs);
-  }
-
-  stat->proc_end = new_end;
-  stat->procs = new_procs;
-
-  return TRUE;
-}
-
-static UINT16
 fort_stat_proc_add (PFORT_STAT stat, UINT32 process_id)
 {
+  const tommy_key_t proc_hash = fort_stat_proc_hash(process_id);
   PFORT_STAT_PROC proc;
-  UINT16 proc_index = stat->proc_free_index;
 
-  if (proc_index != FORT_PROC_BAD_INDEX) {
+  if (stat->proc_free != NULL) {
-    proc = &stat->procs[proc_index];
+    proc = stat->proc_free;
-
+    stat->proc_free = proc->next;
-    stat->proc_free_index = proc->next_index;
   } else {
-    if (stat->proc_top >= stat->proc_end
+    const tommy_count_t size = tommy_arrayof_size(&stat->procs);
-        && !fort_stat_proc_realloc(stat)) {
+
-      return FORT_PROC_BAD_INDEX;
+    /* TODO: tommy_arrayof_grow(): check calloc()'s result for NULL */
+    if (tommy_arrayof_grow(&stat->procs, size + 1), 0)
+      return NULL;
+
+    proc = tommy_arrayof_ref(&stat->procs, size);
+
+    proc->proc_index = (UINT16) size;
   }
 
-    proc_index = stat->proc_top++;
+  tommy_hashdyn_insert(&stat->procs_map, (tommy_hashdyn_node *) proc, 0, proc_hash);
-    proc = &stat->procs[proc_index];
-  }
-
-  /* Prepend to active processes chain */
-  proc->next_index = stat->proc_head_index;
-  stat->proc_head_index = proc_index;
 
   proc->refcount = 0;
   proc->process_id = process_id;
-  proc->traf_all.QuadPart = 0;
+  proc->traf.v = 0;
 
-  stat->proc_count++;
+  return proc;
-
-  return proc_index;
 }
 
 static UCHAR
@@ -372,28 +311,11 @@ fort_stat_flow_add (PFORT_STAT stat, UINT64 flow_id,
   return STATUS_SUCCESS;
 }
 
-static void
-fort_stat_init (PFORT_STAT stat)
-{
-  stat->is_dirty = FALSE;
-
-  if (stat->procs != NULL) {
-    fort_stat_array_del(stat->procs);
-    stat->procs = NULL;
-
-    stat->proc_count = 0;
-    stat->proc_top = 0;
-    stat->proc_end = 0;
-  }
-
-  stat->proc_head_index = FORT_PROC_BAD_INDEX;
-  stat->proc_free_index = FORT_PROC_BAD_INDEX;
-}
-
 static void
 fort_stat_open (PFORT_STAT stat)
 {
-  fort_stat_init(stat);
+  tommy_arrayof_init(&stat->procs, sizeof(FORT_STAT_PROC));
+  tommy_hashdyn_init(&stat->procs_map);
 
   tommy_arrayof_init(&stat->flows, sizeof(FORT_STAT_FLOW));
   tommy_hashdyn_init(&stat->flows_map);
@@ -413,12 +335,24 @@ fort_stat_close (PFORT_STAT stat)
   tommy_hashdyn_foreach_node_arg(&stat->flows_map,
     fort_stat_flow_context_remove, stat);
 
+  tommy_arrayof_done(&stat->procs);
+  tommy_hashdyn_done(&stat->procs_map);
+
   tommy_arrayof_done(&stat->flows);
   tommy_hashdyn_done(&stat->flows_map);
 
   KeReleaseInStackQueuedSpinLock(&lock_queue);
 }
 
+static void
+fort_stat_clear (PFORT_STAT stat)
+{
+  stat->is_dirty = FALSE;
+
+  tommy_hashdyn_foreach_node_arg(&stat->procs_map, fort_stat_proc_free, stat);
+  tommy_hashdyn_foreach_node(&stat->flows_map, fort_stat_flow_close);
+}
+
 static void
 fort_stat_update (PFORT_STAT stat, BOOL log_stat)
 {
@@ -427,9 +361,7 @@ fort_stat_update (PFORT_STAT stat, BOOL log_stat)
   KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
 
   if (stat->log_stat && !log_stat) {
-    tommy_hashdyn_foreach_node(&stat->flows_map, fort_stat_flow_close);
+    fort_stat_clear(stat);
-
-    fort_stat_init(stat);
   }
 
   stat->log_stat = (UCHAR) log_stat;
@@ -460,8 +392,9 @@ fort_stat_flow_associate (PFORT_STAT stat, UINT64 flow_id,
                           UINT32 process_id, UCHAR group_index,
                           BOOL *is_new_proc)
 {
+  const tommy_key_t proc_hash = fort_stat_proc_hash(process_id);
   KLOCK_QUEUE_HANDLE lock_queue;
-  UINT16 proc_index;
+  PFORT_STAT_PROC proc;
   BOOL speed_limit;
   NTSTATUS status;
 
@@ -472,12 +405,12 @@ fort_stat_flow_associate (PFORT_STAT stat, UINT64 flow_id,
     goto end;
   }
 
-  proc_index = fort_stat_proc_index(stat, process_id);
+  proc = fort_stat_proc_get(stat, process_id, proc_hash);
 
-  if (proc_index == FORT_PROC_BAD_INDEX) {
+  if (proc == NULL) {
-    proc_index = fort_stat_proc_add(stat, process_id);
+    proc = fort_stat_proc_add(stat, process_id);
 
-    if (proc_index == FORT_PROC_BAD_INDEX) {
+    if (proc == NULL) {
       status = STATUS_INSUFFICIENT_RESOURCES;
       goto end;
     }
@@ -488,12 +421,10 @@ fort_stat_flow_associate (PFORT_STAT stat, UINT64 flow_id,
   speed_limit = fort_stat_group_speed_limit(stat, group_index) != 0;
 
   status = fort_stat_flow_add(stat, flow_id,
-    group_index, proc_index, speed_limit);
+    group_index, proc->proc_index, speed_limit);
 
   if (!NT_SUCCESS(status) && *is_new_proc) {
-    PFORT_STAT_PROC proc = &stat->procs[proc_index];
+    fort_stat_proc_free(stat, proc);
-
-    fort_stat_proc_free(stat, proc, proc_index, NULL);
   }
 
  end:
@@ -529,7 +460,7 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flowContext,
   KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
 
   if (stat->log_stat && flow->opt.proc_index != FORT_PROC_BAD_INDEX) {
-    PFORT_STAT_PROC proc = &stat->procs[flow->opt.proc_index];
+    PFORT_STAT_PROC proc = tommy_arrayof_ref(&stat->procs, flow->opt.proc_index);
     UINT32 *proc_bytes = inbound ? &proc->traf.in_bytes
       : &proc->traf.out_bytes;
 
@@ -590,49 +521,59 @@ fort_stat_dpc_end (PKLOCK_QUEUE_HANDLE lock_queue)
   KeReleaseInStackQueuedSpinLockFromDpcLevel(lock_queue);
 }
 
+typedef struct fort_stat_proc_traf_iter {
+  PUCHAR out_proc_bits;
+  PFORT_STAT_TRAF out_traf;
+  PFORT_STAT stat;
+} FORT_STAT_PROC_TRAF_ITER, *PFORT_STAT_PROC_TRAF_ITER;
+
 static void
-fort_stat_dpc_traf_flush (PFORT_STAT stat, PCHAR out)
+fort_stat_dpc_traf_proc_flush (PFORT_STAT_PROC_TRAF_ITER iter, PFORT_STAT_PROC proc)
 {
-  const UINT32 proc_bits_len = FORT_LOG_STAT_PROC_SIZE(stat->proc_count);
+  const UINT16 i = proc->proc_index;
-  PFORT_STAT_TRAF out_traf = (PFORT_STAT_TRAF) (out + proc_bits_len);
-  PUCHAR out_proc_bits = (PUCHAR) out;
-  UINT16 i;
-
-  PFORT_STAT_PROC prev_proc = NULL;
-  UINT16 proc_index = stat->proc_head_index;
-
-  /* Mark processes as active to start */
-  memset(out_proc_bits, 0xFF, proc_bits_len);
-
-  for (i = 0; proc_index != FORT_PROC_BAD_INDEX; ++i) {
-    PFORT_STAT_PROC proc = &stat->procs[proc_index];
-    const UINT16 next_index = proc->next_index;
 
   /* Write bytes */
-    *out_traf++ = proc->traf;
+  iter->out_traf[i] = proc->traf;
 
   if (!proc->refcount) {
     /* The process is inactive */
-      out_proc_bits[i / 8] ^= (1 << (i & 7));
+    PFORT_STAT stat = iter->stat;
 
-      fort_stat_proc_free(stat, proc, proc_index, prev_proc);
+    iter->out_proc_bits[i / 8] ^= (1 << (i & 7));
+
+    fort_stat_proc_free(stat, proc);
   } else {
     /* Clear active process's bytes */
-      proc->traf_all.QuadPart = 0;
+    proc->traf.v = 0;
-
+  }
-      prev_proc = proc;
 }
 
-    proc_index = next_index;
+static void
-  }
+fort_stat_dpc_traf_flush (PFORT_STAT stat, PCHAR out)
+{
+  FORT_STAT_PROC_TRAF_ITER iter;
+  const UINT32 proc_bits_len = FORT_LOG_STAT_PROC_SIZE(
+    fort_stat_proc_count(stat));
 
+  iter.out_proc_bits = (PUCHAR) out;
+  iter.out_traf = (PFORT_STAT_TRAF) (out + proc_bits_len);
+  iter.stat = stat;
+
+  /* Mark processes as active to start */
+  memset(iter.out_proc_bits, 0xFF, proc_bits_len);
+
+  /* Iterate over active processes */
+  tommy_hashdyn_foreach_node_arg(&stat->procs_map,
+    fort_stat_dpc_traf_proc_flush, &iter);
+
+  /* Clear process group's bytes */
   if (stat->limit_bits) {
     PFORT_STAT_GROUP group = stat->groups;
     UINT16 limit_bits = stat->limit_bits;
 
     for (; limit_bits != 0; ++group) {
       if (limit_bits & 1) {
-        group->traf_all.QuadPart = 0;
+        group->traf.v = 0;
       }
 
       limit_bits >>= 1;