From 4983d0ffad14eedb1603933bfc50e708e2e27d62 Mon Sep 17 00:00:00 2001
From: Alexander Bushnev <Alexander@Bushnev.pro>
Date: Thu, 21 Nov 2024 11:48:33 +0100
Subject: [PATCH] Fix liveliness cleanup

---
 include/zenoh-pico/api/liveliness.h |  1 +
 src/api/liveliness.c                | 18 +++++++++++++-----
 src/net/liveliness.c                |  2 ++
 src/session/liveliness.c            |  3 +--
 tests/z_api_liveliness_test.c       |  1 +
 5 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/include/zenoh-pico/api/liveliness.h b/include/zenoh-pico/api/liveliness.h
index d3751bc4f..ee2ac9d2f 100644
--- a/include/zenoh-pico/api/liveliness.h
+++ b/include/zenoh-pico/api/liveliness.h
@@ -31,6 +31,7 @@ typedef struct {
     _z_session_weak_t _zn;
 } _z_liveliness_token_t;
 
+_z_liveliness_token_t _z_liveliness_token_null(void);
 _Z_OWNED_TYPE_VALUE(_z_liveliness_token_t, liveliness_token)
 _Z_OWNED_FUNCTIONS_DEF(liveliness_token)
 
diff --git a/src/api/liveliness.c b/src/api/liveliness.c
index 3fb2c63b1..79881395b 100644
--- a/src/api/liveliness.c
+++ b/src/api/liveliness.c
@@ -37,10 +37,17 @@ _z_liveliness_token_t _z_liveliness_token_null(void) {
 }
 
 void _z_liveliness_token_clear(_z_liveliness_token_t *token) {
-    _z_session_rc_t sess_rc = _z_session_weak_upgrade_if_open(&token->_zn);
-    if (!_Z_RC_IS_NULL(&sess_rc)) {
-        _z_undeclare_liveliness_token(token);
-        _z_session_rc_drop(&sess_rc);
+    if (!_z_liveliness_token_check(token)) {
+        return;
+    }
+    // TODO(sashacmc): implement proper check
+    if (token->_zn._val != NULL) {
+        _z_session_rc_t sess_rc = _z_session_weak_upgrade_if_open(&token->_zn);
+        if (!_Z_RC_IS_NULL(&sess_rc)) {
+            _z_undeclare_liveliness_token(token);
+            _z_session_rc_drop(&sess_rc);
+        }
+        _z_session_weak_drop(&token->_zn);
     }
     _z_keyexpr_clear(&token->_key);
 }
@@ -132,7 +139,8 @@ z_result_t z_liveliness_get(const z_loaned_session_t *zs, const z_loaned_keyexpr
         opt = *options;
     }
 
-    ret = _z_liveliness_query(_Z_RC_IN_VAL(zs), *keyexpr, callback->_this._val.call, callback->_this._val.drop, ctx,
+    _z_keyexpr_t ke = _z_keyexpr_duplicate(*keyexpr);
+    ret = _z_liveliness_query(_Z_RC_IN_VAL(zs), ke, callback->_this._val.call, callback->_this._val.drop, ctx,
                               opt.timeout_ms);
 
     z_internal_closure_reply_null(
diff --git a/src/net/liveliness.c b/src/net/liveliness.c
index 45966698f..43354be30 100644
--- a/src/net/liveliness.c
+++ b/src/net/liveliness.c
@@ -60,6 +60,7 @@ z_result_t _z_undeclare_liveliness_token(_z_liveliness_token_t *token) {
     ret = _z_send_n_msg(_Z_RC_IN_VAL(&token->_zn), &n_msg, Z_RELIABILITY_RELIABLE, Z_CONGESTION_CONTROL_BLOCK);
     _z_n_msg_clear(&n_msg);
 
+    *token = _z_liveliness_token_null();
     return ret;
 }
 
@@ -163,6 +164,7 @@ z_result_t _z_liveliness_query(_z_session_t *zn, _z_keyexpr_t keyexpr, _z_closur
 
         } else {
             _z_liveliness_pending_query_clear(pq);
+            z_free(pq);
         }
     }
 
diff --git a/src/session/liveliness.c b/src/session/liveliness.c
index a02a99cc3..ac1be06b9 100644
--- a/src/session/liveliness.c
+++ b/src/session/liveliness.c
@@ -182,8 +182,7 @@ z_result_t _z_liveliness_register_pending_query(_z_session_t *zn, uint32_t id, _
         _Z_ERROR("Duplicate liveliness query id %i", (int)id);
         ret = _Z_ERR_ENTITY_DECLARATION_FAILED;
     } else {
-        _z_liveliness_pending_query_intmap_insert(&zn->_liveliness_pending_queries, id,
-                                                  _z_liveliness_pending_query_clone(pen_qry));
+        _z_liveliness_pending_query_intmap_insert(&zn->_liveliness_pending_queries, id, pen_qry);
     }
 
     _zp_session_unlock_mutex(zn);
diff --git a/tests/z_api_liveliness_test.c b/tests/z_api_liveliness_test.c
index 2cf4c350e..9b8c6c002 100644
--- a/tests/z_api_liveliness_test.c
+++ b/tests/z_api_liveliness_test.c
@@ -199,6 +199,7 @@ void test_liveliness_get(void) {
     assert(z_fifo_handler_reply_recv(z_fifo_handler_reply_loan(&handler), &reply) == Z_CHANNEL_DISCONNECTED);
 
     z_fifo_handler_reply_drop(z_fifo_handler_reply_move(&handler));
+    z_closure_reply_drop(z_closure_reply_move(&cb));
 
     assert_ok(zp_stop_read_task(z_loan_mut(s1)));
     assert_ok(zp_stop_read_task(z_loan_mut(s2)));