DEV Community: Egemen Göl

Building a WebSocket Chatroom using Golang and Spread the PubSub Library

Egemen Göl — Thu, 07 Mar 2024 13:15:06 +0000

We are trying to build a websocket based chatroom.

There are two excellent examples of websocket based chatrooms already, let's go over them.

The server struct holds a set of subscriber structs, each consisting of a chan []byte while knowing how to close themselves.

The server manages access to this set via a mutex and helper functions to keep it current for all of the subscribers coming and going.

gorilla/websocket example

There is a Hub structure that keeps a set of Client structs which point to the Hub and a channel, and manages access to them via exposed channels.
The clients, however, send pointers of themselves over these channels whenever they want to come and go.

The Hub and Client structures are highly coupled, and the logic is spread (pun intended) over two files.

Our Plan

We will use:

nhooyr/websocket for handling websockets.
egemengol/spread for handling the PubSub between multiple websocket connections, with clean and obvious code.

Message Struct

Let's create our Message type.

This struct will be the "message" type of the PubSub topic later.

type Message struct {
    Username string `json:"name"`
    Message  string `json:"msg"`
}

Publish HTTP Handler

Creating a http.Handler for publishing messages to the PubSub topic is easy.

The clients will use it by sending HTTP POST requests to the /publish endpoint

We read the incoming websocket message body and parse it into our Message struct.
We call the spread.Topic.Publish method with our Message.

When we call this function with its required dependencies (logger and topic), we will obtain a http.Handler. We will pass that to http.ServeMux for serving, later.

func HandlePublish(logger *slog.Logger, topic *spread.Topic[Message]) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        var msg Message
        if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
            logger.Warn("error decoding message", "err", err)
            http.Error(w, err.Error(), http.StatusBadRequest)
            return
        }
        r.Body.Close()

        logger.Info("publishing message", "msg", msg)

        if err := topic.Publish(msg); err != nil {
            logger.Error("error publishing message", "err", err)
            http.Error(w, err.Error(), http.StatusInternalServerError)
            return
        }

        w.WriteHeader(http.StatusNoContent)
    })
}

As you can see from the accepted topic argument, the topic knows and restricts the message types it broadcasts. Only Message structs are allowed into this particular topic.

Also, the handler does not know or care anything about its subscribers, it blissfully fires and forgets.

Subscribe Websocket Handler

The clients will connect to this endpoint by making a connection to ws://localhost:8000/subscribe endpoint with their WebSocket library.

Our chatroom implementation chooses to use this websocket connection only for sending messages from the server to the client, even though it could potentially use it in both ways. Makes the implementation and error handling easier. Also, the writer of our websocket library has chosen to implement the chat functionality in this way.

We will upgrade the incoming request to a WebSocket connection by calling Accept.
We will subscribe to the topic, by getting a <-chan Message from it.
We loop through the messages and write them to the client. We keep in mind that at any point, client can disconnect, or topic can be cancelled (on shutdown). We handle these cases in our loop.

func HandleSubscribe(logger *slog.Logger, topic *spread.Topic[Message]) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        conn, err := websocket.Accept(w, r, nil)
        if err != nil {
            logger.Error("error accepting websocket", "err", err)
            return
        }
        defer conn.CloseNow()

        recvChan, removeRecvChan, err := topic.GetRecvChannel(20)
        if err != nil {
            logger.Error("error getting recv channel", "err", err)
            return
        }
        defer removeRecvChan()

        ctx := conn.CloseRead(r.Context())

        for {
            select {
            case <-ctx.Done():
                logger.Info("client disconnected", "err", ctx.Err())
                return
            case msg, ok := <-recvChan:
                if !ok {
                    logger.Info("recv channel closed")
                    conn.Close(websocket.StatusGoingAway, "")
                    return
                }
                data, err := json.Marshal(msg)
                if err != nil {
                    logger.Error("error marshaling message", "err", err)
                    return
                }
                if err := conn.Write(r.Context(), websocket.MessageText, data); err != nil {
                    logger.Warn("error writing message", "err", err)
                    return
                }
                logger.Info(
                    "forwarded to listener", 
                    "fromUser", msg.Username, 
                    "msg", msg.Message
                )
            }
        }
    })
}

When we subscribe to the topic by requesting a read channel, a copy of all of the messages published to the topic start to getting sent to the newly created channel. We are responsible for closing it via the returned destructor function.

ctx := conn.CloseRead(r.Context())

Our websocket library detects the disconnected clients, only when it tries to read or write.

Since we know we do not read anything from client, we use this helper function to create a context, which gets cancelled when the client disconnects.

Construct http.Server

We create a logger to be used by both our handlers and our topic.
We create a topic. It needs:
- The type of the message it will carry, using generics.
- A context, for easy cancellation
- An optional logger, mainly for debug purposes.
- A channel size for the publishers, they will block if the channel is full.
We populate the routes, via the handler factories defined above. We give each their dependencies. We then create the http server.

func Run(ctx context.Context) error {
    logger := slog.New(slog.NewTextHandler(
        os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}
    ))

    topic := spread.NewTopic[Message](ctx, logger, 20)

    mux := http.NewServeMux()

    mux.Handle("POST /publish", HandlePublish(logger, topic))
    mux.Handle("/subscribe", HandleSubscribe(logger, topic))

    httpServer := &http.Server{
        Addr:         "localhost:8000",
        Handler:      mux,
        ReadTimeout:  time.Second * 10,
        WriteTimeout: time.Second * 10,
    }
    // 
    // We will run the server here
    //
}

func main() {
    ctx := context.Background()
    if err := Run(ctx); err != nil {
        log.Fatal(err)
    }
}

Shutdown

Graceful shutdown is the hardest part of a http application in my opinion, especially a websocket enabled one.

We will implement an optimist but brutal shutdown mechanism here.

Get a context that gets cancelled when an interrupt is received.
Spawn a goroutine that waits on that context.
1. Wait for the shutdown of websocket handlers, which Shutdown does not close.
2. Shutdown the http server. The timeout is for the regular connections to close, in this case serving an index.html.

About the wait of the handlers:

The topic is aware of the outer context, which gets cancelled on an interrupt.
The topic notifies any channel based subscribers by closing their receive channels.
The handlers return when they see their recv channels close.

func Run(ctx context.Context) error {
    // Connect the context to interrupts.
    ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
    defer cancel()

    //
    // Rest of the Run function
    //

    go func() {
        // Wait for the context be notified of an interrupt
        <-ctx.Done()

        // Topic is already listening to the context,
        // we know it will send close signals to the handlers
        // We wait for them to return for a bit
        time.Sleep(100 * time.Millisecond)

        // Give the server time to close all the connections
        timeoutCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
        defer cancel()
        if err := httpServer.Shutdown(timeoutCtx); err != nil {
            logger.Error("error closing http server", "err", err)
        }
    }()

    logger.Info("http server started listening on", "addr", httpServer.Addr)
    return httpServer.ListenAndServe()
}

In principle, we could share a sync.WaitGroup between subscribe handlers, and wait on it instead of sleeping a fixed amount. The drawback would be waiting while blocking, which would hang on a non-progressing websocket handler, which is not good. We could wrap that blocking wait inside a goroutine and a channel, with a timeout around them.

It is outside of the scope of this article, which is demonstrating how easy PubSub can be in Golang. The wait here is good enough, the remaining websocket connections are dropped and clients get notified via their websocket libraries anyway.

Conclusion

The application may look complicated, but if we squint hard enough we can see it is mainly boilerplate whenever one writes http servers in golang via the standard library. The /subscribe part is typical for a listening websocket.

When we compare the approach here with the official examples, how obvious the code becomes, pushing all the synchronization work to the egemengol/spread library.

You can see the fully working implementation, along with a simple UI at spread examples directory.

Introducing Spread - Ergonomic PubSub/EventBus in Golang

Egemen Göl — Mon, 26 Feb 2024 21:34:40 +0000

Why

When one wants pub/sub functionality in a Golang application, there are mainly two choices.

One: You bring a Big Boi like RabbitMQ, Kafka, Redis or some other messaging solution. They are battle-tested and scalable, however they come with some drawbacks:

Development complexity: Every developer needs to know the sdk, semantics and performance characteristics of given message broker.
Deployment complexity: If you are hosting the broker yourself, there is another service you need to manage and keep track of, as well as creating a repeatable development environment in your system for you and your colleagues.
Expensive: In terms of cloud costs or energy consumption, it is quite a bit heavier than your standalone Golang application.

Two: You implement a custom solution, like the chatroom examples of gorilla/websocket or nhooyr/websocket. It works well, you use no external dependencies in typical Gopher fashion, but:

Error prone: It is notoriously easy to miss subtle bugs in concurrent code. We use Golang which makes things so much easier, it is still a problem
Tight coupling: If you are using channels, opening and closing channels in a correct manner is hard (malloc and free vibes anyone?) and requires synchronization between producers and consumers. It makes your application hard to modify as well, since you never know if you introduced a new bug, too subtle to catch, right away.

There are a lot of problems that live in between. I will go out of a limb and will claim that %90 of startup ideas and side projects will never grow out of a singular beefy multicore server that hosts a well-written Golang monolith. In the case you do, you have wonderful problems, have at thee with the Big Bois, but only then.

There is the really nice asaskevich/EventBus library which covers a lot, and is the main inspiration of this library.

A couple minor notes:

Unmaintained in the last 4 years.
Uses mutexes, hard to grasp what is going on.
Precedes generics, interfaces are clunky to work with.

So Spread, this library here, is a

Single process and in-memory,
PubSub / EventBus / Broadcast / Fanout library with
Ergonomic, type-safe topics implemented using generics,
Based on channels.

Intended to be useful for:

In-memory pipelines with persistence ==> Event Sourcing
Decoupled aggregates for said pipelines
Soft-realtime applications that needs broadcast, like chatrooms.

Overview

You can subscribe to a Topic in three ways, simultaneously and dynamically:

var topic *spread.Topic[T]

// Channel based
recvChan, _, _ := topic.GetRecvChannel(bufSize)
for msg := range recvChan {
    ...
}

// A separate, freshly spawned goroutine per message
topic.HandleAsync(
    func(context.Context, T) {}
)

// Synchronous but blocks the topic
topic.HandleSync(
    func(T) {}
)

Performance Characteristics

Every topic has a inbound channel with a dedicated goroutine for broadcasting.
Synchronous handlers in HandleSync get executed in this goroutine.
Asynchronous handlers or receiver channels that cannot keep up (with full buffers) get eliminated from the subscribers.
Publishing is the same as sending to a buffered channel, blocks when full.

Plans

I love the excellent Phoenix.PubSub library in the Elixir ecosystem, and would love to achieve the same flavor and balance of usefulness, reliability and flexibility with Spread. I believe we can come close.

I also am very impressed by the LMAX Architecture, which prompted me to think about the PubSub pattern.

There are at least two impressive implementations of it in Golang, go-distruptor and zenq, which I will thoroughly study and will try to join their very performant ideas with an easy-to-use API, in the long run.

I also want to make it really easy to build Event Sourced applications, Spread can be a good building block for it.

Github: egemengol/spread

If you think it is worth trying out or talking about, if you have any ideas, lets meet!