Repository: DiceDB/dice
File: internal/server/websocketServer.go (Go)

Code before:
			}
		}

		if val, ok := responseValue.(clientio.RespType); ok {
			responseValue = respArr[val]
		}

		if bt, ok := responseValue.([]byte); ok {
			responseValue = string(bt)
		}

		respBytes, err := json.Marshal(responseValue)
		if err != nil {
			writeResponse(conn, []byte("error: marshaling json response"))
			continue
		}

		// Write response with retries for transient errors
		if err := writeResponseWithRetries(conn, respBytes, 3); err != nil {
			s.logger.Error(fmt.Sprintf("Error reading message: %v", err))
			break // Exit the loop on write error
		}
	}
}

func writeResponseWithRetries(conn *websocket.Conn, text []byte, maxRetries int) error {
	for attempts := 0; attempts < maxRetries; attempts++ {
		// Set a write deadline
		if err := conn.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil {
			slog.Error(fmt.Sprintf("Error setting write deadline: %v", err))
			return err
		}

		// Write message back to client
		if err := conn.WriteMessage(websocket.TextMessage, text); err != nil {
			var netErr *net.OpError
			if errors.As(err, &netErr) {
				if opErr, ok := netErr.Err.(*os.SyscallError); ok {
					switch opErr.Syscall {
					case "write":
						switch opErr.Err {
						case syscall.EPIPE:
							return fmt.Errorf("broken pipe: %w", err)
						case syscall.ECONNRESET:
							return fmt.Errorf("connection reset by peer: %w", err)
						case syscall.ENOBUFS:
							return fmt.Errorf("no buffer space available: %w", err)
						case syscall.EAGAIN:
							time.Sleep(time.Duration(attempts+1)*100*time.Millisecond + time.Duration(rand.Intn(50))*time.Millisecond) // Exponential; backoff with jitter
							continue
						default:

Review comment (suggestion):
@ayushsatyam146 please add test cases around this. May be we can mock the websocket conn for different scenarios using `testify/mock`.



Code after:
			}
		}

		if val, ok := responseValue.(clientio.RespType); ok {
			responseValue = respArr[val]
		}

		if bt, ok := responseValue.([]byte); ok {
			responseValue = string(bt)
		}

		respBytes, err := json.Marshal(responseValue)
		if err != nil {
			writeResponse(conn, []byte("error: marshaling json response"))
			continue
		}

		// Write response with retries for transient errors
		if err := WriteResponseWithRetries(conn, respBytes, config.DiceConfig.WebSocket.MaxWriteResponseRetries); err != nil {
			s.logger.Error(fmt.Sprintf("Error reading message: %v", err))
			break // Exit the loop on write error
		}
	}
}

func WriteResponseWithRetries(conn *websocket.Conn, text []byte, maxRetries int) error {
	for attempts := 0; attempts < maxRetries; attempts++ {
		// Set a write deadline
		if err := conn.SetWriteDeadline(time.Now().Add(config.DiceConfig.WebSocket.WriteResponseTimeout)); err != nil {
			slog.Error(fmt.Sprintf("Error setting write deadline: %v", err))
			return err
		}

		// Attempt to write message
		err := conn.WriteMessage(websocket.TextMessage, text)
		if err == nil {
			break // Exit loop if write succeeds
		}

		// Handle network errors
		var netErr *net.OpError
		if !errors.As(err, &netErr) {
			return fmt.Errorf("error writing message: %w", err)
		}

		opErr, ok := netErr.Err.(*os.SyscallError)
		if !ok {
			return fmt.Errorf("network operation error: %w", err)
		}

		if opErr.Syscall != "write" {