From 9b1a6607ced0e47b3bca221e9db6e0e0c3bf4b64 Mon Sep 17 00:00:00 2001
From: Bernardo <68619889+bchuecos@users.noreply.github.com>
Date: Mon, 7 Mar 2022 22:00:16 +0000
Subject: [PATCH] review suggestion: deregister handle after receiving
 notification

---
 pkg/liveshare/ports.go   |  6 ++++--
 pkg/liveshare/rpc.go     | 25 ++++++++++++++++++++++++-
 pkg/liveshare/session.go |  6 +++---
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/pkg/liveshare/ports.go b/pkg/liveshare/ports.go
index b20e326a4..a1a9ffa39 100644
--- a/pkg/liveshare/ports.go
+++ b/pkg/liveshare/ports.go
@@ -91,8 +91,10 @@ func (s *Session) WaitForPortNotification(ctx context.Context, port int, notifTy
 			notificationUpdate <- notification
 		}
 	}
-	s.registerRequestHandler("serverSharing.sharingSucceeded", h(true))
-	s.registerRequestHandler("serverSharing.sharingFailed", h(false))
+	deregisterSuccess := s.registerRequestHandler("serverSharing.sharingSucceeded", h(true))
+	deregisterFailure := s.registerRequestHandler("serverSharing.sharingFailed", h(false))
+	defer deregisterSuccess()
+	defer deregisterFailure()
 
 	for {
 		select {
diff --git a/pkg/liveshare/rpc.go b/pkg/liveshare/rpc.go
index a32d8507a..5972656aa 100644
--- a/pkg/liveshare/rpc.go
+++ b/pkg/liveshare/rpc.go
@@ -53,7 +53,7 @@ func newRequestHandler() *requestHandler {
 	return &requestHandler{handlers: make(map[string][]handlerFn)}
 }
 
-func (r *requestHandler) register(requestType string, handler handlerFn) {
+func (r *requestHandler) register(requestType string, handler handlerFn) func() {
 	r.handlersMu.Lock()
 	defer r.handlersMu.Unlock()
 
@@ -62,6 +62,29 @@ func (r *requestHandler) register(requestType string, handler handlerFn) {
 	}
 
 	r.handlers[requestType] = append(r.handlers[requestType], handler)
+
+	return func() {
+		r.deregister(requestType, handler)
+	}
+}
+
+func (r *requestHandler) deregister(requestType string, handler handlerFn) {
+	r.handlersMu.Lock()
+	defer r.handlersMu.Unlock()
+
+	if handlers, ok := r.handlers[requestType]; ok {
+		newHandlers := []handlerFn{}
+		for _, h := range handlers {
+			if &h != &handler {
+				newHandlers = append(newHandlers, h)
+			}
+		}
+		r.handlers[requestType] = newHandlers
+
+		if len(r.handlers[requestType]) == 0 {
+			delete(r.handlers, requestType)
+		}
+	}
 }
 
 func (r *requestHandler) handlerFn(requestType string) []handlerFn {
diff --git a/pkg/liveshare/session.go b/pkg/liveshare/session.go
index 6249acb23..c7e38225f 100644
--- a/pkg/liveshare/session.go
+++ b/pkg/liveshare/session.go
@@ -31,9 +31,9 @@ func (s *Session) Close() error {
 }
 
 // registerRequestHandler registers a handler for the given request type with the RPC
-// server.
-func (s *Session) registerRequestHandler(requestType string, h handlerFn) {
-	s.rpc.requestHandler.register(requestType, h)
+// server and returns a callback function to deregister the handler
+func (s *Session) registerRequestHandler(requestType string, h handlerFn) func() {
+	return s.rpc.requestHandler.register(requestType, h)
 }
 
 // StartsSSHServer starts an SSH server in the container, installing sshd if necessary,