diff --git a/include/stellar/packet.h b/include/stellar/packet.h
index a19e892..1af4dfe 100644
--- a/include/stellar/packet.h
+++ b/include/stellar/packet.h
@@ -211,11 +211,11 @@ struct packet_manager;
 struct packet_manager *module_to_packet_manager(struct module *mod);
 int packet_manager_new_packet_exdata_index(struct packet_manager *pkt_mgr, const char *name, exdata_free *func, void *arg);
 
-typedef void on_packet_callback(struct packet *pkt, void *arg);
+typedef void on_packet_callback(struct packet *pkt, struct module *mod);
 int packet_manager_register_node(struct packet_manager *pkt_mgr, const char *name, enum packet_stage stage,
                                  uint64_t interested_tag_key_bits,
                                  uint64_t interested_tag_val_bits,
-                                 on_packet_callback *cb, void *arg);
+                                 on_packet_callback *cb, struct module *mod);
 // if two modules claim the same packet at the same stage, the second 'claim' fails.
 // return  0 on success
 // return -1 on failure
diff --git a/include/stellar/session.h b/include/stellar/session.h
index 866d986..f47d6b5 100644
--- a/include/stellar/session.h
+++ b/include/stellar/session.h
@@ -146,8 +146,8 @@ int session_manager_new_session_exdata_index(struct session_manager *sess_mgr, c
 struct session *session_manager_lookup_session_by_packet(struct session_manager *sess_mgr, uint16_t thread_id, const struct packet *pkt);
 struct session *session_manager_lookup_session_by_id(struct session_manager *sess_mgr, uint16_t thread_id, uint64_t sess_id);
 
-void session_manager_on_packet_forward(struct packet *pkt, void *args);
-void session_manager_on_packet_output(struct packet *pkt, void *args);
+void session_manager_on_packet_forward(struct packet *pkt, struct module *mod);
+void session_manager_on_packet_output(struct packet *pkt, struct module *mod);
 
 struct module *session_manager_on_init(struct module_manager *mod_mgr);
 void session_manager_on_exit(struct module_manager *mod_mgr, struct module *mod);
diff --git a/infra/packet_manager/packet_manager.c b/infra/packet_manager/packet_manager.c
index 57f5f0c..cc14498 100644
--- a/infra/packet_manager/packet_manager.c
+++ b/infra/packet_manager/packet_manager.c
@@ -16,7 +16,7 @@ struct node
     uint64_t interested_tag_key_bits;
     uint64_t interested_tag_val_bits;
     on_packet_callback *node_entry;
-    void *arg;
+    struct module *mod;
 };
 
 #define MAX_NODE_PER_STAGE 128
@@ -174,7 +174,7 @@ int packet_manager_new_packet_exdata_index(struct packet_manager *pkt_mgr, const
 int packet_manager_register_node(struct packet_manager *pkt_mgr, const char *name, enum packet_stage stage,
                                  uint64_t interested_tag_key_bits,
                                  uint64_t interested_tag_val_bits,
-                                 on_packet_callback *node_entry, void *arg)
+                                 on_packet_callback *node_entry, struct module *mod)
 {
     assert(pkt_mgr);
     assert(stage < PACKET_STAGE_MAX);
@@ -192,7 +192,7 @@ int packet_manager_register_node(struct packet_manager *pkt_mgr, const char *nam
     node->interested_tag_key_bits = interested_tag_key_bits;
     node->interested_tag_val_bits = interested_tag_val_bits;
     node->node_entry = node_entry;
-    node->arg = arg;
+    node->mod = mod;
 
     nodes->used++;
     return 0;
@@ -302,7 +302,7 @@ void packet_manager_dispatch(struct packet_manager *pkt_mgr, uint16_t thread_id)
                 if ((pkt_tag_key_bits & node->interested_tag_key_bits) &&
                     (pkt_tag_val_bits & node->interested_tag_val_bits))
                 {
-                    node->node_entry(pkt, node->arg);
+                    node->node_entry(pkt, node->mod);
                 }
             }
 
diff --git a/infra/packet_manager/test/gtest_packet_manager.cpp b/infra/packet_manager/test/gtest_packet_manager.cpp
index cef52e4..7c027d2 100644
--- a/infra/packet_manager/test/gtest_packet_manager.cpp
+++ b/infra/packet_manager/test/gtest_packet_manager.cpp
@@ -109,7 +109,7 @@ TEST(PACKET_MANAGER, NEW_FREE)
 #endif
 
 #if 1
-static void on_packet(struct packet *pkt, void *args)
+static void on_packet(struct packet *pkt, struct module *mod)
 {
 	enum packet_stage stage = packet_get_stage(pkt);
 	printf("on_packet_stage: %s\n", packet_stage_to_str(stage));
@@ -123,7 +123,7 @@ static void on_packet(struct packet *pkt, void *args)
 	static int count = 0;
 	EXPECT_TRUE(count == stage);
 	EXPECT_TRUE(packet_get_type(pkt) == PACKET_TYPE_PSEUDO);
-	EXPECT_TRUE(args == NULL);
+	EXPECT_TRUE(mod == NULL);
 	count++;
 }
 
@@ -176,7 +176,7 @@ TEST(PACKET_MANAGER, REGISTER)
 #endif
 
 #if 1
-static void drop_packet(struct packet *pkt, void *args)
+static void drop_packet(struct packet *pkt, struct module *mod)
 {
 	enum packet_stage stage = packet_get_stage(pkt);
 	printf("on_packet_stage: %s\n", packet_stage_to_str(stage));
@@ -190,7 +190,7 @@ static void drop_packet(struct packet *pkt, void *args)
 	static int count = 0;
 	EXPECT_TRUE(count == stage);
 	EXPECT_TRUE(packet_get_type(pkt) == PACKET_TYPE_PSEUDO);
-	EXPECT_TRUE(args == NULL);
+	EXPECT_TRUE(mod == NULL);
 	count++;
 
 	if (stage == PACKET_STAGE_FORWARD)
@@ -247,9 +247,9 @@ TEST(PACKET_MANAGER, DROP_PACKET)
 #endif
 
 #if 1
-static void packet_claimed(struct packet *pkt, void *args)
+static void packet_claimed(struct packet *pkt, struct module *mod)
 {
-	char *str = (char *)args;
+	char *str = (char *)mod;
 	EXPECT_STREQ(str, "hello");
 	printf("packet_claimed: with ctx %s\n", str);
 	EXPECT_TRUE(packet_get_type(pkt) == PACKET_TYPE_PSEUDO);
@@ -257,9 +257,9 @@ static void packet_claimed(struct packet *pkt, void *args)
 	free(str);
 }
 
-static void claim_packet_success(struct packet *pkt, void *args)
+static void claim_packet_success(struct packet *pkt, struct module *mod)
 {
-	struct packet_manager *pkt_mgr = (struct packet_manager *)args;
+	struct packet_manager *pkt_mgr = (struct packet_manager *)module_get_ctx(mod);
 	enum packet_stage stage = packet_get_stage(pkt);
 	printf("on_packet_stage: %s\n", packet_stage_to_str(stage));
 
@@ -278,9 +278,9 @@ static void claim_packet_success(struct packet *pkt, void *args)
 	count++;
 }
 
-static void claim_packet_failed(struct packet *pkt, void *args)
+static void claim_packet_failed(struct packet *pkt, struct module *mod)
 {
-	struct packet_manager *pkt_mgr = (struct packet_manager *)args;
+	struct packet_manager *pkt_mgr = (struct packet_manager *)module_get_ctx(mod);
 	enum packet_stage stage = packet_get_stage(pkt);
 	printf("on_packet_stage: %s\n", packet_stage_to_str(stage));
 
@@ -303,18 +303,19 @@ TEST(PACKET_MANAGER, CLAIM_PACKET)
 {
 	// module init
 	struct packet_manager *pkt_mgr = packet_manager_new(1);
+	struct module *pkt_mgr_mod = module_new("packet_manager", pkt_mgr);
 	EXPECT_TRUE(pkt_mgr);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_success, pkt_mgr_mod) == 0);
 
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_failed, pkt_mgr_mod) == 0);
 
 	// per-thread init
 	packet_manager_init(pkt_mgr, thread_id);
@@ -349,13 +350,14 @@ TEST(PACKET_MANAGER, CLAIM_PACKET)
 
 	// module free
 	packet_manager_free(pkt_mgr);
+	module_free(pkt_mgr_mod);
 }
 #endif
 
 #if 1
-static void schedule_packet(struct packet *pkt, void *args)
+static void schedule_packet(struct packet *pkt, struct module *mod)
 {
-	struct packet_manager *pkt_mgr = (struct packet_manager *)args;
+	struct packet_manager *pkt_mgr = (struct packet_manager *)module_get_ctx(mod);
 	enum packet_stage stage = packet_get_stage(pkt);
 	printf("on_packet_stage: %s\n", packet_stage_to_str(stage));
 
@@ -380,12 +382,13 @@ TEST(PACKET_MANAGER, SCHEDULE_PACKET)
 {
 	// module init
 	struct packet_manager *pkt_mgr = packet_manager_new(1);
+	struct module *pkt_mgr_mod = module_new("packet_manager", pkt_mgr);
 	EXPECT_TRUE(pkt_mgr);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, schedule_packet, pkt_mgr_mod) == 0);
 
 	// per-thread init
 	packet_manager_init(pkt_mgr, thread_id);
@@ -429,13 +432,15 @@ TEST(PACKET_MANAGER, SCHEDULE_PACKET)
 
 	// module free
 	packet_manager_free(pkt_mgr);
+
+	module_free(pkt_mgr_mod);
 }
 #endif
 
 #if 1
-static void schedule_claimed_packet(struct packet *pkt, void *args)
+static void schedule_claimed_packet(struct packet *pkt, struct module *mod)
 {
-	struct packet_manager *pkt_mgr = (struct packet_manager *)args;
+	struct packet_manager *pkt_mgr = (struct packet_manager *)module_get_ctx(mod);
 
 	printf("schedule_claimed_packet:  %p\n", pkt);
 	EXPECT_TRUE(packet_get_type(pkt) == PACKET_TYPE_PSEUDO);
@@ -444,9 +449,9 @@ static void schedule_claimed_packet(struct packet *pkt, void *args)
 	packet_manager_schedule_packet(pkt_mgr, thread_id, pkt, PACKET_STAGE_POSTROUTING);
 }
 
-static void claim_packet_to_schedule(struct packet *pkt, void *args)
+static void claim_packet_to_schedule(struct packet *pkt, struct module *mod)
 {
-	struct packet_manager *pkt_mgr = (struct packet_manager *)args;
+	struct packet_manager *pkt_mgr = (struct packet_manager *)module_get_ctx(mod);
 	enum packet_stage stage = packet_get_stage(pkt);
 	printf("on_packet_stage: %s\n", packet_stage_to_str(stage));
 
@@ -462,7 +467,7 @@ static void claim_packet_to_schedule(struct packet *pkt, void *args)
 	if (stage == PACKET_STAGE_PREROUTING)
 	{
 		EXPECT_TRUE(count == 0);																				  // packet not claim
-		EXPECT_TRUE(packet_manager_claim_packet(pkt_mgr, thread_id, pkt, schedule_claimed_packet, pkt_mgr) == 0); // claim packet success
+		EXPECT_TRUE(packet_manager_claim_packet(pkt_mgr, thread_id, pkt, schedule_claimed_packet, mod) == 0); // claim packet success
 	}
 	else if (stage == PACKET_STAGE_POSTROUTING)
 	{
@@ -480,12 +485,13 @@ TEST(PACKET_MANAGER, SCHEDULE_CLAIMED_PACKET)
 {
 	// module init
 	struct packet_manager *pkt_mgr = packet_manager_new(1);
+	struct module *pkt_mgr_mod=module_new("packet_manager", pkt_mgr);
 	EXPECT_TRUE(pkt_mgr);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr) == 0);
-	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_PREROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_INPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr_mod) == 0);
+	EXPECT_TRUE(packet_manager_register_node(pkt_mgr, "name", PACKET_STAGE_POSTROUTING, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP, claim_packet_to_schedule, pkt_mgr_mod) == 0);
 
 	// per-thread init
 	packet_manager_init(pkt_mgr, thread_id);
@@ -520,6 +526,8 @@ TEST(PACKET_MANAGER, SCHEDULE_CLAIMED_PACKET)
 
 	// module free
 	packet_manager_free(pkt_mgr);
+
+	module_free(pkt_mgr_mod);
 }
 #endif
 
diff --git a/infra/session_manager/session_manager.c b/infra/session_manager/session_manager.c
index 1b9e5f4..c69ebe6 100644
--- a/infra/session_manager/session_manager.c
+++ b/infra/session_manager/session_manager.c
@@ -131,9 +131,9 @@ static void on_polling(struct module_manager *mod_mgr, void *args)
  * session manager
  ******************************************************************************/
 
-void session_manager_on_packet_forward(struct packet *pkt, void *args)
+void session_manager_on_packet_forward(struct packet *pkt, struct module *mod)
 {
-     struct session_manager *sess_mgr = (struct session_manager *)args;
+     struct session_manager *sess_mgr = (struct session_manager *)module_get_ctx(mod);
      int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
      struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, thread_id);
 
@@ -223,9 +223,9 @@ void session_manager_on_packet_forward(struct packet *pkt, void *args)
      }
 }
 
-void session_manager_on_packet_output(struct packet *pkt, void *args)
+void session_manager_on_packet_output(struct packet *pkt, struct module *mod)
 {
-     struct session_manager *sess_mgr = (struct session_manager *)args;
+     struct session_manager *sess_mgr = (struct session_manager *)module_get_ctx(mod);
      int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
      struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, thread_id);