CODEWITHSUNDEEP
javascripthtmlgo
atye
atye

Reputation: 109

How to gracefully re-connect websocket?

I'm building a singe-page application using HTTP and Websockets. The user submits a form and I stream a response to the client. Below is a snippet client.

var html = `<!DOCTYPE html>

<meta charset="utf-8">
<head>
</head>
<body>

<script>
var ws = new WebSocket("ws://localhost:8000/ws")
ws.onmessage = function(e) {
  document.getElementById("output").innerHTML += e.data + "<br>"
}

function submitFunction() {
  document.getElementById("output").innerHTML += ""
  return false
}
</script>

<form
enctype="multipart/x-www-form-urlencoded"
action="http://localhost:8000/"
method="post"
>`

This is the server. If the request is not a POST, I write/render the html (parseAndExecute) which establishes a new websocket connection. If the request is a POST (from the form), then I start processing and eventually write to the websocket.

func (c *Config) ServeHtml(w http.ResponseWriter, r *http.Request) {
    if r.Method == http.MethodPost {
        //process
        channel <- data
    }

    c.parseAndExecute(w)
}

func (sh *SocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        w.Write([]byte(fmt.Sprintf("", err)))
        return
    }
    //defer ws.Close()

    // discard received messages
    go func(c *websocket.Conn) {
        for {
            if _, _, err := c.NextReader(); err != nil {
                c.Close()
                break
            }
        }
    }(ws)

    data <- channel

Everything works as I expect only if I do not refresh the page. If I don't refresh, I can keep submitting forms and see the different outputs come in line by line. To clarify, it actually only works if the page is already up so that parseAndExecute is never called. This function parses and executes html/template creating a new websocket client.

Any refresh of the page or initially browsing localhost:8000 would cause websocket: close sent on the server.

I'm not sure how to resolve this. Does the server to need to gracefully handle disconnections and allow re-connects? Or does the client need to do something? It seems like the server should upgrade any connection at /ws so it shouldn't matter how many new websocket clients are made but obviously my understanding is wrong.

I'm not closing the websocket connection on the server because it should be up for as long as the program is running. When the user stops the program, I assume it'll be automatically closed.

Full SocketHandler code:

func (sh *SocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        w.Write([]byte(fmt.Sprintf("", err)))
        return
    }

    // discard received messages
    go func(c *websocket.Conn) {
        for {
            if _, _, err := c.NextReader(); err != nil {
                c.Close()
                break
            }
        }
    }(ws)

    cmd := <-sh.cmdCh

    log.Printf("Executing")
    stdout, err := cmd.StdoutPipe()
    if err != nil {
        w.Write([]byte(err.Error()))
        return
    }
    defer stdout.Close()

    stderr, err := cmd.StderrPipe()
    if err != nil {
        w.Write([]byte(err.Error()))
        return
    }
    defer stderr.Close()

    if err := cmd.Start(); err != nil {
        w.Write([]byte(err.Error()))
        return
    }

    s := bufio.NewScanner(io.MultiReader(stdout, stderr))
    for s.Scan() {
        err := ws.WriteMessage(1, s.Bytes())
        if err != nil {
            log.Printf("Error writing to client: %v", err)
            ws.Close()
        }

    }

    if err := cmd.Wait(); err != nil {
        w.Write([]byte(err.Error()))
        return
    }

    log.Println("Done")
}

Upvotes: 0

Views: 2501

Answers (1)

Thundercat
Thundercat

Reputation: 120951

Websocket server applications must handle errors on the connection by closing the connection and releasing resources associated with the connection.

Browsers close websocket connections when the containing page is refreshed. The server will eventually get a read or write error after the browser closes the connection.

Connection close is one of several possible errors that the server might encounter. The server application should handle all errors on the connection the same way: close the connection and release resources associated with the connection.

The typical application design is for the client to connect on page load and to reconnect (with backoff) after an error. The server assumes that clients will connect and disconnect over time.

The JS code can be improved by adding an onerror handler that reconnects with backoff. Depending on the application, you may also want to display UI indicating the connection status.

The Go code does not close the connection in all scenarios. The running command is the resource associated with the connection. The application does not kill this program on connection error. Here are some fixes:

Add defer ws.Close() after successful upgrade. Remove other direct calls to ws.Close() from SocketHandler.ServeHTTP. This ensures that ws.Close() is called in all scenarios.

Kill the command on exit from the read and write pumps. Move the read pump to after the command is started. Kill on return.

go func(c *websocket.Conn, cmd *exec.Command) {
    defer c.Close()
    defer cmd.Process.Kill()
    for {
        if _, _, err := c.NextReader(); err != nil {
            break
        }
    }
}(ws, cmd)

Kill the command on exit from the write pump:

s := bufio.NewScanner(io.MultiReader(stdout, stderr))
for s.Scan() {
    err := ws.WriteMessage(1, s.Bytes())
    if err != nil {
        break
    }
}
cmd.Process.Kill()

I have not run or tested this code. It's possible that some of the details are wrong, but this outlines the general approach of closing the connection and releasing the resource.

Take a look at the Gorilla command example. The example shows how to pump stdin and stdout to a websocket. The example also handles more advanced features like checking the health of the connection with PING/PONG.

Upvotes: 1

Related Questions