Pass weak pointer of HeadBhv to callback

2 files changed, 75 insertions, 7 deletions

diff --git a/include/erebos/storage.h b/include/erebos/storage.h
index 01aeada..c506dfd 100644
--- a/include/erebos/storage.h
+++ b/include/erebos/storage.h
@@ -528,6 +528,7 @@ void filterAncestors(std::vector<Stored<T>> & xs)
 }
 
 template<class T> class WatchedHead;
+template<class T> class HeadBhv;
 
 template<class T>
 class Head
@@ -562,8 +563,13 @@ template<class T>
 class WatchedHead : public Head<T>
 {
 	friend class Head<T>;
+	friend class HeadBhv<T>;
+
+	WatchedHead(const Head<T> & h):
+		Head<T>(h), watcherId(-1) {}
 	WatchedHead(const Head<T> & h, int watcherId):
 		Head<T>(h), watcherId(watcherId) {}
+
 	int watcherId;
 
 public:
@@ -588,15 +594,26 @@ class HeadBhv : public BhvSource<T>
 {
 public:
 	HeadBhv(const Head<T> & head):
-		whead(head.watch([this] (const Head<T> & cur) {
-			BhvCurTime ctime;
-			whead = cur;
-			BhvImplBase::updated(ctime);
-		})) {}
+		whead(head)
+	{}
 
 	T get(const BhvCurTime &, const std::monostate &) const { return *whead; }
 
 private:
+	friend class Head<T>;
+
+	void init()
+	{
+		whead = whead.watch([wp = weak_ptr<BhvImplBase>(BhvImplBase::shared_from_this()), this] (const Head<T> & cur) {
+			// make sure this object still exists
+			if (auto ptr = wp.lock()) {
+				BhvCurTime ctime;
+				whead = cur;
+				BhvImplBase::updated(ctime);
+			}
+		});
+	}
+
 	WatchedHead<T> whead;
 };
 
@@ -665,7 +682,9 @@ template<typename T>
 Bhv<T> Head<T>::behavior() const
 {
 	auto cur = reload();
-	return make_shared<HeadBhv<T>>(cur ? *cur : *this);
+	auto ret = make_shared<HeadBhv<T>>(cur ? *cur : *this);
+	ret->init();
+	return ret;
 }
 
 template<class T>
diff --git a/test/storage.test b/test/storage.test
index 17859cc..6267beb 100644
--- a/test/storage.test
+++ b/test/storage.test
@@ -1,4 +1,4 @@
-test:
+test Storage:
 	spawn as p1
 
 	# Root finding
@@ -150,3 +150,52 @@ test:
 		/stored-set-item $r2_3/
 		/stored-set-(.*)/ capture done6
 	guard done6 == "done"
+
+
+test StorageWatcher:
+	spawn as p1
+	spawn as p2
+	send "create-identity Device1 Owner" to p1
+	send "create-identity Device2" to p2
+	send "watch-local-identity" to p1
+	send "watch-local-identity" to p2
+	send "start-server" to p1
+	send "start-server" to p2
+	expect from p1:
+		/local-identity Device1 Owner/
+		/peer 1 addr ${p2.node.ip} 29665/
+		/peer 1 id Device2/
+	expect from p2:
+		/local-identity Device2/
+		/peer 1 addr ${p1.node.ip} 29665/
+		/peer 1 id Device1 Owner/
+
+	local:
+		send "attach-to 1" to p2
+		expect /attach-request 1 ([0-9]*)/ from p1 capture code1
+		expect /attach-response 1 ([0-9]*)/ from p2 capture code2
+		guard code1 == code2
+
+	send "attach-accept 1" to p1
+	send "attach-accept 1" to p2
+	expect /attach-request-done 1/ from p1
+	expect /attach-response-done 1/ from p2
+	expect /local-identity Device2 Owner/ from p2
+	expect /peer 1 id Device2 Owner/ from p1
+
+	for i in [1..5]:
+		send "update-local-identity Owner2" to p1
+		send "shared-state-get" to p1
+		expect /shared-state-get (.*)/ from p1 capture s1
+
+		send "shared-state-wait $s1" to p2
+		expect /shared-state-wait $s1/ from p2
+
+		send "update-local-identity Owner1" to p1
+		send "shared-state-get" to p1
+		expect /shared-state-get (.*)/ from p1 capture s2
+
+		send "shared-state-wait $s1" to p2
+		send "shared-state-wait $s2" to p2
+		expect /shared-state-wait $s1/ from p2
+		expect /shared-state-wait $s2/ from p2