Remove handled fragmented messages from the map

We should remove handled fragmented messages from fragmentedMessages
map to free it up. This PR includes this fix.

Also, FragmentedClientMessageHandlerTest is enhanced to check the
the map's size before and after merging and handling the
end fragment.

Author: mdumandag

src/network/ClientConnection.ts

@@ -244,19 +244,15 @@ export class FragmentedClientMessageHandler {
         clientMessage.dropFragmentationFrame();
         if (fragmentationFrame.hasBeginFragmentFlag()) {
             this.fragmentedMessages.set(fragmentationId, clientMessage);
-        } else if (fragmentationFrame.hasEndFragmentFlag()) {
-            const mergedMessage = this.mergeIntoExistingClientMessage(fragmentationId, clientMessage);
-            callback(mergedMessage);
         } else {
-            this.mergeIntoExistingClientMessage(fragmentationId, clientMessage);
+            const existingMessage = this.fragmentedMessages.get(fragmentationId);
+            existingMessage.merge(clientMessage);
+            if (fragmentationFrame.hasEndFragmentFlag()) {
+                callback(existingMessage);
+                this.fragmentedMessages.delete(fragmentationId);
+            }
         }
     }
-
-    private mergeIntoExistingClientMessage(fragmentationId: number, clientMessage: ClientMessage): ClientMessage {
-        const existingMessage = this.fragmentedMessages.get(fragmentationId);
-        existingMessage.merge(clientMessage);
-        return existingMessage;
-    }
 }
 
 export class ClientConnection {

test/connection/FragmentedClientMessageHandlerTest.js

@@ -24,7 +24,7 @@ const ClientMessageReader = require('../../lib/network/ClientConnection').Client
 const cm = require('../../lib/ClientMessage');
 const FixSizedTypesCodec = require('../../lib/codec/builtin/FixSizedTypesCodec').FixSizedTypesCodec;
 
-describe('FragmentedClientMessageHandler', function () {
+describe('FragmentedClientMessageHandlerTest', function () {
 
     let reader;
     let handler;
@@ -62,15 +62,19 @@ describe('FragmentedClientMessageHandler', function () {
 
         // Should merge with the above message
         handler.handleFragmentedMessage(fragment1, () => done(new Error("It should just merge.")));
+        expect(handler.fragmentedMessages.size).to.equal(1);
         compareMessages(fragments[1].startFrame.next, endFrame.next);
         endFrame = fragmentedMessage.endFrame;
 
         // Should merge with the above message and run the callback
+        let isCalled = false;
         handler.handleFragmentedMessage(fragment2, () => {
             compareMessages(fragments[2].startFrame.next, endFrame.next);
-            done();
+            isCalled = true;
         });
-
+        expect(isCalled).to.be.true;
+        expect(handler.fragmentedMessages.size).to.equal(0);
+        done();
     });
 
     const compareMessages = (frame1, frame2) => {