From 231fbf1bdfd547347b21b834d64f94b3eb26aa53 Mon Sep 17 00:00:00 2001
From: Michal Sojka <sojkam1@fel.cvut.cz>
Date: Tue, 20 Jan 2009 11:33:03 +0100
Subject: [PATCH] Improved reference counting of peer structures

---
 src/iop.c  |  6 ++----
 src/peer.c | 16 +++++++++-------
 src/peer.h |  3 ++-
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/src/iop.c b/src/iop.c
index 262b950..d7d3a53 100644
--- a/src/iop.c
+++ b/src/iop.c
@@ -412,11 +412,9 @@ process_hello(forb_port_t *port, FORB_CDR_Codec *codec)
 				if (notify_waiters) {
 					fosa_cond_broadcast(&peer->cond);
 				} else {
-					forb_peer_get(peer);
-					forb_peer_nolock_insert(forb, peer);
+					forb_peer_nolock_insert(forb, forb_peer_get(peer));
 				}
-				forb_peer_get(peer);
-				forb_port_peer_ins_tail(port, peer);
+				forb_port_peer_ins_tail(port, forb_peer_get(peer));
 				fosa_mutex_unlock(&forb->peer_mutex);
 				forb_peer_put(peer);
 			}
diff --git a/src/peer.c b/src/peer.c
index 8384ef6..ea04410 100644
--- a/src/peer.c
+++ b/src/peer.c
@@ -123,21 +123,21 @@ forb_peer_find_timed(forb_t *forb, forb_server_id *server_id,
 		     fosa_abs_time_t *timeout)
 {
 	forb_peer_t *peer;
+	bool peer_allocated = false;
 	
 	fosa_mutex_lock(&forb->peer_mutex);
 	peer = forb_peer_nolock_find(forb, server_id);
 	if (!peer && !timeout) goto unlock;
 	if (peer) {
 		forb_peer_get(peer);
-	}
-	if (!peer) {
+	} else {
 		/* We are the first who want to contact this peer */
 		peer = forb_peer_new();
 		if (!peer) goto unlock;
 		peer->server_id = *server_id;
 		peer->state = FORB_PEER_WANTED;
-		forb_peer_get(peer);
-		forb_peer_nolock_insert(forb, peer);
+		forb_peer_nolock_insert(forb, forb_peer_get(peer));
+		peer_allocated = true;
 	}
 	while (peer->state == FORB_PEER_WANTED) {
 		int ret;
@@ -146,11 +146,13 @@ forb_peer_find_timed(forb_t *forb, forb_server_id *server_id,
 					  &forb->peer_mutex,
 					  timeout);
 		if (ret == FOSA_ETIMEDOUT) {
+			/* No peer discovered within timeout */
+			if (peer_allocated) {
+				forb_peer_nolock_delete(forb, peer);
+				forb_peer_put(peer);
+			}
 			forb_peer_put(peer);
 			peer = NULL;
-			/* TODO: How to remove the peer from the tree?
-			 * We should check refcnt and if equal to 1,
-			 * we should remove it.  */
 			break;
 		}
 	}
diff --git a/src/peer.h b/src/peer.h
index 778eb24..caa30eb 100644
--- a/src/peer.h
+++ b/src/peer.h
@@ -114,10 +114,11 @@ forb_peer_new(void);
 void
 forb_peer_release(forb_ref_t *ref);
 
-static inline void
+static inline forb_peer_t *
 forb_peer_get(forb_peer_t *peer)
 {
 	forb_ref_get(&peer->ref);
+	return peer;
 }
 	
 static inline void
-- 
2.39.2