1 files changed, 37 insertions, 13 deletions

diff --git a/src/Common/RequestManager.h b/src/Common/RequestManager.h
index 07a7c05..e66b0c1 100644
--- a/src/Common/RequestManager.h
+++ b/src/Common/RequestManager.h
@@ -21,6 +21,7 @@
 #define MAD_COMMON_REQUESTMANAGER_H_
 
 #include "Request.h"
+#include "RequestHandlerGroup.h"
 
 #include <map>
 #include <memory>
@@ -69,18 +70,20 @@ class RequestManager : boost::noncopyable {
 
 
 
-    class RequestHandlerFactory {
-      protected:
-        RequestHandlerFactory() {}
+    template<class T> class RequestHandlerFactory : public RequestHandlerGroup {
+      private:
+        std::set<std::string> types;
 
       public:
-        virtual boost::shared_ptr<RequestHandler> createRequestHandler() = 0;
-        virtual ~RequestHandlerFactory() {}
-    };
+        RequestHandlerFactory(const std::string &type) {
+          types.insert(type);
+        }
 
-    template<class T> class SpecificRequestHandlerFactory : public RequestHandlerFactory {
-      public:
-        virtual boost::shared_ptr<RequestHandler> createRequestHandler() {
+        virtual const std::set<std::string>& getPacketTypes() {
+          return types;
+        }
+
+        virtual boost::shared_ptr<RequestHandler> createRequestHandler(const std::string &type _UNUSED_PARAMETER_) {
           return boost::shared_ptr<RequestHandler>(new T());
         }
     };
@@ -94,7 +97,7 @@ class RequestManager : boost::noncopyable {
     bool server;
     boost::uint16_t lastRequestId;
 
-    std::map<std::string, boost::shared_ptr<RequestHandlerFactory> > requestHandlerFactories;
+    std::map<std::string, boost::shared_ptr<RequestHandlerGroup> > requestHandlerGroups;
 
     boost::uint16_t _getRequestId() {
       return lastRequestId+=2;
@@ -141,14 +144,35 @@ class RequestManager : boost::noncopyable {
       requestMap.unregisterConnection(connection);
     }
 
-    template <class T> void registerPacketType(const std::string &type) {
+    void registerRequestHandlerGroup(boost::shared_ptr<RequestHandlerGroup> requestHandlerGroup) {
       boost::lock_guard<boost::shared_mutex> lock(mutex);
-      requestHandlerFactories.insert(std::make_pair(type, boost::shared_ptr<SpecificRequestHandlerFactory<T> >(new SpecificRequestHandlerFactory<T>())));
+
+      const std::set<std::string> &types = requestHandlerGroup->getPacketTypes();
+      for(std::set<std::string>::const_iterator type = types.begin(); type != types.end(); ++type)
+        requestHandlerGroups.insert(std::make_pair(*type, requestHandlerGroup));
+    }
+
+    void unregisterRequestHandlerGroup(boost::shared_ptr<RequestHandlerGroup> requestHandlerGroup) {
+      boost::lock_guard<boost::shared_mutex> lock(mutex);
+
+      const std::set<std::string> &types = requestHandlerGroup->getPacketTypes();
+      for(std::set<std::string>::const_iterator type = types.begin(); type != types.end(); ++type) {
+        std::map<std::string, boost::shared_ptr<RequestHandlerGroup> >::iterator it = requestHandlerGroups.find(*type);
+
+        if(it == requestHandlerGroups.end() || it->second != requestHandlerGroup)
+          continue;
+
+        requestHandlerGroups.erase(it);
+      }
+    }
+
+    template <class T> void registerPacketType(const std::string &type) {
+      registerRequestHandlerGroup(boost::shared_ptr<RequestHandlerFactory<T> >(new RequestHandlerFactory<T>(type)));
     }
 
     void unregisterPacketType(const std::string &type) {
       boost::lock_guard<boost::shared_mutex> lock(mutex);
-      requestHandlerFactories.erase(type);
+      requestHandlerGroups.erase(type);
     }
 
     bool sendRequest(Connection *connection, boost::shared_ptr<Request> request);