From bcf0dff302a8416cb7528ee44c2f0fb3138ff13e Mon Sep 17 00:00:00 2001
From: yangchen <yangchen145@huawei.com>
Date: Tue, 25 Jun 2024 17:53:36 +0800
Subject: [PATCH] fix function call error

---
 src/lstack/core/lstack_protocol_stack.c | 8 ++++----
 src/lstack/core/lstack_stack_stat.c     | 2 +-
 src/lstack/include/lstack_stack_stat.h  | 2 --
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/lstack/core/lstack_protocol_stack.c b/src/lstack/core/lstack_protocol_stack.c
index d130c91..d6e3c86 100644
--- a/src/lstack/core/lstack_protocol_stack.c
+++ b/src/lstack/core/lstack_protocol_stack.c
@@ -921,10 +921,6 @@ void stack_udp_send(struct rpc_msg *msg)
     int replenish_again;
     uint32_t call_num;
 
-    if (get_protocol_stack_group()->latency_start) {
-        calculate_rpcmsg_latency(&stack->latency, msg, GAZELLE_LATENCY_WRITE_RPC_MSG);
-    }
-
     struct lwip_sock *sock = get_socket(fd);
     if (sock == NULL) {
         msg->result = -1;
@@ -932,6 +928,10 @@ void stack_udp_send(struct rpc_msg *msg)
         return;
     }
 
+    if (get_protocol_stack_group()->latency_start) {
+        calculate_sock_latency(&stack->latency, sock, GAZELLE_LATENCY_WRITE_RPC_MSG);
+    }
+
     replenish_again = do_lwip_send(stack, sock->conn->callback_arg.socket, sock, len, 0);
     call_num = __sync_fetch_and_sub(&sock->call_num, 1);
     if (replenish_again < 0) {
diff --git a/src/lstack/core/lstack_stack_stat.c b/src/lstack/core/lstack_stack_stat.c
index a1bd44d..0f1f693 100644
--- a/src/lstack/core/lstack_stack_stat.c
+++ b/src/lstack/core/lstack_stack_stat.c
@@ -82,7 +82,7 @@ void time_stamp_record(int fd, struct pbuf *pbuf)
 {
     struct lwip_sock *sock = get_socket_by_fd(fd);
 
-    if (get_protocol_stack_group()->latency_start && pbuf != NULL) {
+    if (get_protocol_stack_group()->latency_start && sock && pbuf) {
         calculate_lstack_latency(&sock->stack->latency, pbuf, GAZELLE_LATENCY_INTO_MBOX, 0);
         time_stamp_into_recvmbox(sock);
     }
diff --git a/src/lstack/include/lstack_stack_stat.h b/src/lstack/include/lstack_stack_stat.h
index f0d3d5d..9fca04a 100644
--- a/src/lstack/include/lstack_stack_stat.h
+++ b/src/lstack/include/lstack_stack_stat.h
@@ -26,8 +26,6 @@ struct lwip_sock;
 
 void calculate_lstack_latency(struct gazelle_stack_latency *stack_latency, const struct pbuf *pbuf,
     enum GAZELLE_LATENCY_TYPE type, uint64_t time_record);
-void calculate_rpcmsg_latency(struct gazelle_stack_latency *stack_latency, struct rpc_msg *msg,
-    enum GAZELLE_LATENCY_TYPE type);
 void calculate_sock_latency(struct gazelle_stack_latency *stack_latency, struct lwip_sock *sock,
     enum GAZELLE_LATENCY_TYPE type);
 void stack_stat_init(void);
-- 
2.33.0