Fix concurrent access

Author: Antonio Cheong

.gitignore

@@ -13,3 +13,4 @@
 
 # Dependency directories (remove the comment below to include it)
 # vendor/
+Data/auth.db

handlers/api.go

@@ -66,17 +66,21 @@ func API_ask(c *gin.Context) {
 		})
 		return
 	}
+	connectionPool.Mu.RUnlock()
 	if request.ConversationId == "" {
 		// Retry 3 times before giving up
+		var succeeded bool = false
 		for i := 0; i < 3; i++ {
 			// Find connection with the lowest load and where heartbeat is after last message time
+			connectionPool.Mu.RLock()
 			for _, conn := range connectionPool.Connections {
 				if connection == nil || conn.LastMessageTime.Before(connection.LastMessageTime) {
 					if conn.Heartbeat.After(conn.LastMessageTime) {
 						connection = conn
 					}
 				}
 			}
+			connectionPool.Mu.RUnlock()
 			// Check if connection was found
 			if connection == nil {
 				c.JSON(503, gin.H{
@@ -87,11 +91,21 @@ func API_ask(c *gin.Context) {
 			// Ping before sending request
 			if !ping(connection.Id) {
 				// Ping failed. Try again
+				println("Ping failed")
 				continue
+			} else {
+				println("Ping succeeded")
 			}
 			// Ping succeeded. Break the loop
+			succeeded = true
 			break
 		}
+		if !succeeded {
+			c.JSON(503, gin.H{
+				"error": "Ping failed",
+			})
+			return
+		}
 	} else {
 		// Check if conversation exists
 		conversation, ok := conversationPool.Get(request.ConversationId)
@@ -122,7 +136,6 @@ func API_ask(c *gin.Context) {
 			return
 		}
 	}
-	connectionPool.Mu.RUnlock()
 	message := types.Message{
 		Id:      utils.GenerateId(),
 		Message: "ChatGptRequest",
@@ -205,6 +218,7 @@ func API_getConnections(c *gin.Context) {
 }
 
 func ping(connection_id string) bool {
+	println("Starting ping")
 	// Get connection
 	connection, ok := connectionPool.Get(connection_id)
 	// Send "ping" to the connection
@@ -214,7 +228,9 @@ func ping(connection_id string) bool {
 			Id:      id,
 			Message: "ping",
 		}
+		connection.Ws.SetReadDeadline(time.Now().Add(5 * time.Second))
 		err := connection.Ws.WriteJSON(send)
+		println("Sent ping")
 		if err != nil {
 			// Delete connection
 			connectionPool.Delete(connection_id)
@@ -224,9 +240,10 @@ func ping(connection_id string) bool {
 		for {
 			// Read message
 			var receive types.Message
-			connection.Ws.SetReadDeadline(time.Now().Add(5 * time.Second))
 			err = connection.Ws.ReadJSON(&receive)
+			println("Received ping response")
 			if err != nil {
+				println("There was an error")
 				// Delete connection
 				connectionPool.Delete(connection_id)
 				return false

types/types.go

@@ -60,10 +60,11 @@ func (p *ConnectionPool) Set(conn *Connection) {
 	p.Connections[conn.Id] = conn
 }
 
-func (p *ConnectionPool) Delete(id string) {
+func (p *ConnectionPool) Delete(id string) error {
 	p.Mu.Lock()
 	defer p.Mu.Unlock()
 	delete(p.Connections, id)
+	return nil
 }
 
 func NewConnectionPool() *ConnectionPool {