From 313257a1aba26dea837512ab7b01850a33cba56d Mon Sep 17 00:00:00 2001
From: speatzle <sam@soontm.de>
Date: Tue, 21 May 2024 15:43:08 +0200
Subject: [PATCH] Add Request Timeout, Cleanup Request map

---
 rpc/call.go     | 44 +++++++++++++++++++++++++++++---------------
 rpc/response.go |  3 +++
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/rpc/call.go b/rpc/call.go
index f7ad220..df653fb 100644
--- a/rpc/call.go
+++ b/rpc/call.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"time"
 
 	"github.com/google/uuid"
 	"nhooyr.io/websocket"
@@ -57,26 +58,39 @@ func (s *server) Call(ctx context.Context, c *websocket.Conn, method string, par
 		return nil, fmt.Errorf("Error Writing Request: %w", err)
 	}
 
-	// Wait for Response, TODO add Select Timeout
-	response := <-resp
+	timeout := time.NewTimer(time.Second * 10)
 
-	if response.Error != nil {
-		return response, fmt.Errorf("Call Error: %w", err)
-	}
+	select {
+	case response := <-resp:
+		// Cleanup timeout
+		if !timeout.Stop() {
+			<-timeout.C
+		}
+		if response.Error != nil {
+			return response, fmt.Errorf("Call Error: %w", err)
+		}
 
-	if result == nil {
+		if result == nil {
+			return response, nil
+		}
+
+		if response.Result == nil {
+			return response, fmt.Errorf("Got Empty Result")
+		}
+
+		err = json.Unmarshal(*response.Result, &result)
+		if err != nil {
+			return response, fmt.Errorf("Error Parsing Result: %w", err)
+		}
 		return response, nil
-	}
+	case <-timeout.C:
+		s.requestMutex.Lock()
+		defer s.requestMutex.Unlock()
 
-	if response.Result == nil {
-		return response, fmt.Errorf("Got Empty Result")
+		// remove request from map
+		delete(s.requests, id)
+		return nil, fmt.Errorf("Request timed out")
 	}
-
-	err = json.Unmarshal(*response.Result, &result)
-	if err != nil {
-		return response, fmt.Errorf("Error Parsing Result: %w", err)
-	}
-	return response, nil
 }
 
 // TODO Call with Multiple Response (Chunked file upload)
diff --git a/rpc/response.go b/rpc/response.go
index 52d2239..33f209f 100644
--- a/rpc/response.go
+++ b/rpc/response.go
@@ -22,6 +22,9 @@ func (s *server) handleResponse(ctx context.Context, data []byte) {
 		return
 	}
 
+	// remove request from map
+	delete(s.requests, response.ID)
+
 	// Send Response to Original Caller
 	r <- &response
 }