From 3ecab1f31e8a5db9368f723739643a9a8a3dd58c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Mon, 2 Jun 2025 18:43:08 +0200
Subject: [PATCH] Add context

---
 client/internal/peer/status.go         | 14 ++++++++++----
 client/internal/routemanager/client.go |  2 +-
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/client/internal/peer/status.go b/client/internal/peer/status.go
index d3375f107..0c6aac372 100644
--- a/client/internal/peer/status.go
+++ b/client/internal/peer/status.go
@@ -1,6 +1,7 @@
 package peer
 
 import (
+	"context"
 	"errors"
 	"net/netip"
 	"slices"
@@ -150,10 +151,12 @@ type StatusChangeSubscription struct {
 	peerID     string
 	id         string
 	eventsChan chan struct{}
+	ctx        context.Context
 }
 
-func newStatusChangeSubscription(peerID string) *StatusChangeSubscription {
+func newStatusChangeSubscription(ctx context.Context, peerID string) *StatusChangeSubscription {
 	return &StatusChangeSubscription{
+		ctx:        ctx,
 		peerID:     peerID,
 		id:         uuid.New().String(),
 		eventsChan: make(chan struct{}, 1),
@@ -569,11 +572,11 @@ func (d *Status) FinishPeerListModifications() {
 	d.notifyPeerListChanged()
 }
 
-func (d *Status) SubscribeToPeerStateChanges(peerID string) *StatusChangeSubscription {
+func (d *Status) SubscribeToPeerStateChanges(ctx context.Context, peerID string) *StatusChangeSubscription {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 
-	sub := newStatusChangeSubscription(peerID)
+	sub := newStatusChangeSubscription(ctx, peerID)
 	if _, ok := d.changeNotify[peerID]; !ok {
 		d.changeNotify[peerID] = make(map[string]*StatusChangeSubscription)
 	}
@@ -986,7 +989,10 @@ func (d *Status) notifyPeerStateChangeListeners(peerID string) {
 		// block the write because we do not want to miss notification
 		// must have to be sure we will run the GetPeerState() on separated thread
 		go func() {
-			sub.eventsChan <- struct{}{}
+			select {
+			case sub.eventsChan <- struct{}{}:
+			case <-sub.ctx.Done():
+			}
 		}()
 	}
 }
diff --git a/client/internal/routemanager/client.go b/client/internal/routemanager/client.go
index 797ef6c48..bff954c27 100644
--- a/client/internal/routemanager/client.go
+++ b/client/internal/routemanager/client.go
@@ -224,7 +224,7 @@ func (c *clientNetwork) getBestRouteFromStatuses(routePeerStatuses map[route.ID]
 }
 
 func (c *clientNetwork) watchPeerStatusChanges(ctx context.Context, peerKey string, peerStateUpdate chan struct{}, closer chan struct{}) {
-	subscription := c.statusRecorder.SubscribeToPeerStateChanges(peerKey)
+	subscription := c.statusRecorder.SubscribeToPeerStateChanges(ctx, peerKey)
 	defer c.statusRecorder.UnsubscribePeerStateChanges(subscription)
 
 	for {