review suggestion: deregister handle after receiving notification

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 {

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 {

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,