From 5ec344b7332391067ba38e4ee771892d481edb63 Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Wed, 18 Dec 2019 13:47:46 +0900
Subject: [PATCH] Add way to get connection queued size

Requires locking the mutex for access to conMap.
---
 src/UDPConnection.cpp       | 24 ++++++++++++++++++++++++
 src/UDPConnection.h         |  4 ++++
 src/test/UDPC_NetworkTest.c | 37 ++++++++++++++++++++++---------------
 3 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/src/UDPConnection.cpp b/src/UDPConnection.cpp
index ca419cb..40638ad 100644
--- a/src/UDPConnection.cpp
+++ b/src/UDPConnection.cpp
@@ -2129,6 +2129,30 @@ unsigned long UDPC_get_queue_send_current_size(UDPC_HContext ctx) {
     return c->cSendPkts.size();
 }
 
+unsigned long UDPC_get_queued_size(UDPC_HContext ctx, UDPC_ConnectionId id, int *exists) {
+    UDPC::Context *c = UDPC::verifyContext(ctx);
+    if(!c) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> lock(c->mutex);
+    auto iter = c->conMap.find(id);
+    if(iter != c->conMap.end()) {
+        if(exists) {
+            *exists = 1;
+        }
+        return iter->second.sendPkts.size();
+    }
+    if(exists) {
+        *exists = 0;
+    }
+    return 0;
+}
+
+unsigned long UDPC_get_max_queued_size() {
+    return UDPC_QUEUED_PKTS_MAX_SIZE;
+}
+
 int UDPC_set_accept_new_connections(UDPC_HContext ctx, int isAccepting) {
     UDPC::Context *c = UDPC::verifyContext(ctx);
     if(!c) {
diff --git a/src/UDPConnection.h b/src/UDPConnection.h
index bb9eef1..611861f 100644
--- a/src/UDPConnection.h
+++ b/src/UDPConnection.h
@@ -455,6 +455,10 @@ void UDPC_queue_send(UDPC_HContext ctx, UDPC_ConnectionId destinationId,
  */
 unsigned long UDPC_get_queue_send_current_size(UDPC_HContext ctx);
 
+unsigned long UDPC_get_queued_size(UDPC_HContext ctx, UDPC_ConnectionId id, int *exists);
+
+unsigned long UDPC_get_max_queued_size();
+
 int UDPC_set_accept_new_connections(UDPC_HContext ctx, int isAccepting);
 
 void UDPC_drop_connection(UDPC_HContext ctx, UDPC_ConnectionId connectionId, int dropAllWithAddr);
diff --git a/src/test/UDPC_NetworkTest.c b/src/test/UDPC_NetworkTest.c
index 83dd539..4add21a 100644
--- a/src/test/UDPC_NetworkTest.c
+++ b/src/test/UDPC_NetworkTest.c
@@ -225,6 +225,7 @@ int main(int argc, char **argv) {
     unsigned int tick = 0;
     unsigned int temp = 0;
     unsigned int temp2, temp3;
+    int temp4;
     unsigned long size;
     UDPC_ConnectionId *list = NULL;
     unsigned int sendIds[SEND_IDS_SIZE];
@@ -242,24 +243,30 @@ int main(int argc, char **argv) {
         }
         if(!noPayload) {
             list = UDPC_get_list_connected(context, &temp);
-            if(list) {
-                if(sendIdsSize < temp) {
-                    while(sendIdsSize < temp) {
-                        if(sendIdsSize == SEND_IDS_SIZE) {
-                            temp = SEND_IDS_SIZE;
-                            break;
-                        }
-                        sendIds[sendIdsSize++] = 0;
+
+            if(sendIdsSize < temp) {
+                while(sendIdsSize < temp) {
+                    if(sendIdsSize == SEND_IDS_SIZE) {
+                        temp = SEND_IDS_SIZE;
+                        break;
                     }
-                } else if(sendIdsSize > temp) {
-                    sendIdsSize = temp;
+                    sendIds[sendIdsSize++] = 0;
                 }
-                size = UDPC_get_queue_send_current_size(context);
-                temp2 = size < QUEUED_MAX_SIZE ? QUEUED_MAX_SIZE - size : 0;
-                for(unsigned int i = 0; i < temp2; ++i) {
-                    temp3 = htonl(sendIds[i % temp]++);
-                    UDPC_queue_send(context, list[i % temp], 0, &temp3, sizeof(unsigned int));
+            } else if(sendIdsSize > temp) {
+                sendIdsSize = temp;
+            }
+
+            for(unsigned int i = 0; i < temp; ++i) {
+                size = UDPC_get_max_queued_size() - UDPC_get_queued_size(context, list[i], &temp4);
+                if(temp4 == 0) {
+                    continue;
                 }
+                for(unsigned int j = 0; j < size; ++j) {
+                    temp2 = htonl(sendIds[i]++);
+                    UDPC_queue_send(context, list[i], 0, &temp2, sizeof(unsigned int));
+                }
+            }
+            if(list) {
                 UDPC_free_list_connected(list);
             }
             do {