Building a Scalable Chat Server in Golang using NATS

This is an experimental project that I am working on and so far, it’s been working great. Haven’t found any major issues with the implementation so I thought I would document it.

Now, most of you will know about golang but I bet most of you haven’t even heard about NATS. That’s what I’ll be focusing on mostly here.

I wanted to build a full stack chat application in react and use golang for backend because ever since I started working with golang at my work, I have come to love it dearly. It just makes programming simpler in a way where you worry less about the code and have the time to think about what you want to build and how. But the act of writing golang itself has been utterly enjoyable.

So, what is NATS? At the core, it’s a pub/sub message delivering system mostly used in microservices architecture to communicate between distributed servers. It’s got a persistence layer of it’s own(Jetstream), a key-value store and an object store all in one singe go module which can even be embedded in a golang server as a plugin.

That’s what caught my attention. Everything that you could possibly need for a chat application is handed to you on a platter. So I went ahead with it.

Assuming you have go installed on your system, let’s just spin up a NATS instance on docker. Run the following command to start NATS locally.

docker run -d -p 4222:4222 -p 8222:8222 -p 6222:6222 --name nats-server -ti nats:latest -js

The -js flat is to enable jetstream in NATS. Jetstream is a package that comes with NATS which is the persistence layer on top of NATS to allow message replays, access to the key-value store and object storage.

Once the NATS instance is running, you can interact with it already with the help of nats-cli which is really easy and straightforward to use. Very helpful while debugging.

Now, let’s setup a project for the server.

# Create project directory
mkdir go-nats-chat
cd go-nats-chat

# Initialize Go module
go mod init your_module_name # e.g., github.com/yourusername/go-nats-chat

# Get dependencies
go get github.com/gofiber/fiber/v2
go get github.com/gofiber/contrib/websocket
go get github.com/nats-io/nats.go@latest # Ensure you get a version with JetStream support

As you can see in the packages that we installed, we are gonna be setting up a basic golang websocket server which lets the clients send and receive messages from but the messages itself will be handled by NATS internally.

Let’s setup a config file for the server.

// config/config.go
package config

import "time"

const (
 NatsURL        = "nats://localhost:4222"
 StreamName     = "CHAT_MESSAGES"
 SubjectPrefix  = "chat.conversation" // e.g., chat.conversation.general
 ServerAddr     = ":8080"
 WriteWait      = 10 * time.Second // Time allowed to write a message to the peer.
 PongWait       = 60 * time.Second // Time allowed to read the next pong message from the peer.
 PingPeriod     = (PongWait * 9) / 10 // Send pings to peer with this period. Must be less than pongWait.
 MaxMessageSize = 512
)

Before we go ahead, let’s just understand how exactly using NATS and Jetstream is going to help us here.

Like I said before, Jetstream is a persistence layer for NATS. Now, for our server, we just need one stream which keeps track of all the items. Now, in NATS, while consuming items in a stream, we can choose which particular subset of items from the whole stream that we want using wildcards. This is decided by the subject name into which the item was pushed in the first place.

For example, in our case the subject names for conversations is going to be chat.conversation.<conversation_id>. But we can get all the messages for chat.conversation irrespective of which conversation by using the following wildcard chat.conversation.>.

Our setup will be simple, each conversation gets a subject in the stream. That way, we can easily track conversations.

Here is a quick overview:

Now that we have a good understanding of what we are building. Let’s setup a message data model.

// models/message.go
package models

import (
 "time"
)

// Message represents a chat message
type Message struct {
 ID             string    `json:"id"`             // Unique message ID (e.g., UUID)
 ConversationID string    `json:"conversationId"` // ID of the chat room/conversation
 SenderID       string    `json:"senderId"`       // ID of the user sending the message
 SenderAlias    string    `json:"senderAlias"`    // Display name of the sender
 Text           string    `json:"text"`           // Message content
 CreatedAt      time.Time `json:"createdAt"`      // Timestamp of message creation
}

Now, let’s create a new package called nats_service to keep all the NATS related code in one place. But before just dumping the whole code here, let’s go method by method understanding what each does.

// nats_service/service.go
package nats_service

import (
 "context"
 "encoding/json"
 "fmt"
 "log"
 "time"

 "your_module_name/config" // Adjust import path
 "your_module_name/models" // Adjust import path

 "github.com/nats-io/nats.go"
 "github.com/nats-io/nats.go/jetstream"
)

type NatsService struct {
 js jetstream.JetStream
 nc *nats.Conn
}

// NewNatsService connects to NATS and initializes JetStream
func NewNatsService() (*NatsService, error) {
 nc, err := nats.Connect(config.NatsURL) // can also use nats.DefaultURL for localhost connections
 if err != nil {
  return nil, fmt.Errorf("failed to connect to NATS: %w", err)
 }

 js, err := jetstream.New(nc) // creating a jetstream instance for the above created nats connection
 if err != nil {
  nc.Close()
  return nil, fmt.Errorf("failed to create jetstream context: %w", err)
 }

 // Ensure stream exists
 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 defer cancel()

 stream, err := js.Stream(ctx, config.StreamName)
 if err != nil {
  log.Printf("Stream '%s' not found, attempting to create...", config.StreamName)
  streamCfg := jetstream.StreamConfig{
   Name:        config.StreamName,
   Description: "Stores chat messages",
   Subjects:    []string{fmt.Sprintf("%s.*", config.SubjectPrefix)}, // Subject hierarchy
   MaxAge:      24 * time.Hour,                                     // Example retention policy
   Storage:     jetstream.FileStorage, // also allows memory storage
  }
  stream, err = js.CreateStream(ctx, streamCfg)
  if err != nil {
   nc.Close()
   return nil, fmt.Errorf("failed to create stream '%s': %w", config.StreamName, err)
  }
  log.Printf("Stream '%s' created successfully", config.StreamName)
 } else {
  log.Printf("Found existing stream '%s'", stream.CachedInfo().Config.Name)
 }

 return &NatsService{js: js, nc: nc}, nil
}

// Close NATS connection
func (s *NatsService) Close() {
 if s.nc != nil {
  s.nc.Close()
 }
}

The initiation is pretty straight forward. You connect to the local NATS server. Create a jetstream instance for the nats connection. Check if our stream already exists, if not, create a new one.

The error management style offered in go is truly amazing. Yes, it might feel repetitive in the start but once you get used to it, other languages ahem, ahem, javascript feels outright scary to use. I still have to work with javascript but the feeling of safety that I get in go is unmatched. On multiple occasions, I have made a change and pushed it to prod without even having to think about it. Granted, the changes were not huge and very predictable. But go just makes you incredibly confident in your code and that’s what I love about it so much.

Alright, back to business. Now that we have the NATS connection and the stream setup, let’s see how to publish messages and realise how easy it is.

// PublishMessage sends a message to the appropriate NATS subject
func (s *NatsService) PublishMessage(ctx context.Context, msg *models.Message) error {
 subject := getSubject(msg.ConversationID)
 msgData, err := json.Marshal(msg)
 if err != nil {
  return fmt.Errorf("failed to marshal message: %w", err)
 }

 // Publish message to JetStream
 _, err = s.js.Publish(ctx, subject, msgData)
 if err != nil {
  return fmt.Errorf("failed to publish message to subject '%s': %w", subject, err)
 }
 log.Printf("Published message to %s (ID: %s)", subject, msg.ID)
 return nil
}

// getSubject generates the NATS subject for a conversation
func getSubject(conversationID string) string {
 return fmt.Sprintf("%s.%s", config.SubjectPrefix, conversationID)
}

All that we need is a context, the subject to which the message needs to be passed and the message data itself. This simplicity of the NATS library is what got me hooked to NATS.

Now there are different flavours of Publish available like: PublishMsg, PublishAsync, andPublishMsgAsync . I have only been using Publish so far and it’s working great. The only notable difference is of course with the synchronous nature of Publish and the asynchronous nature of PublishAsync , but there are very minute differences between the other two methods relative to their synchronous and asynchronous counterparts.

Now, coming to the most important concept of NATS, the consumers! This one has caught me off-guard multiple times and the configuration of the consumer is very important to understand.

In our server, we are just going to have a single consumer which gets all the messages.

There are few things that are simplified in this setup which might not be feasible in an actual production server. For example, we are only creating a single consumer in NATS which uses the DeliverAllPolicy. Sending all the messages always is not a good idea.

1. In a real production server, the messages might be getting saved in a DB somewhere. Having the server return the latest 20 messages from the DB directly would be a more practical approach.
2. The consumer could be setup to only start sending messages from 5 minutes ago or so by using the DeliverByStartTimePolicy and setting OptStartTime totime.Now().Add(-5*time.Minute).

3. Periodic sync to the required DB can be setup in the server, where another consumer is setup to read the messages of all the subjects every hour or so and push them to the DB.

4. The messages delivery statuses can also be managed by the key value store that’s available with jetstream.

Here is the message consumer:

// It calls the handler function for each new message received.
func (s *NatsService) SubscribeToConversation(ctx context.Context, conversationID string, handler func(msg *models.Message)) (jetstream.ConsumeContext, error) {
 subject := getSubject(conversationID)
 // lastTimestamp := time.Now().Add(-5 * time.Minute)
 // Create a durable consumer (or ephemeral if preferred)
 // Durable consumers remember their position. Ephemeral start from now/end.
 // Using ephemeral here for simplicity, means clients only get messages sent *after* they connect.
  cons, err := s.js.CreateOrUpdateConsumer(ctx, config.StreamName, jetstream.ConsumerConfig{  FilterSubject: subject,
  DeliverPolicy: jetstream.DeliverAllPolicy,
  // DeliverPolicy: jetstream.DeliverByStartTimePolicy, // Start from point in time
  AckPolicy:     jetstream.AckNonePolicy,
  // OptStartTime:  lastTimestamp // this starts consuming messages added 5 minutes ago
 })
 if err != nil {
  return nil, fmt.Errorf("failed to create consumer for subject '%s': %w", subject, err)
 }

 log.Printf("Subscribing to %s", subject)

 // Consume messages
 consumeCtx, err := cons.Consume(func(jsMsg jetstream.Msg) {
  var msg models.Message
  if err := json.Unmarshal(jsMsg.Data(), &msg); err != nil {
   log.Printf("Error unmarshaling message from subject '%s': %v", jsMsg.Subject(), err)
   return
  }

  // Process the message using the provided handler
  handler(&msg)
 })
 if err != nil {
  return nil, fmt.Errorf("failed to start consuming from subject '%s': %w", subject, err)
 }

 // Return the context so the caller can stop it later
 return consumeCtx, nil
}

Again, the code looks very straightforward, but, we need to understand a few key terms here. In NATS jetstream, there are two kinds of consumers.

Ephemeral: These consumers don’t retain state between restarts. NATS has an internal system which clears out these consumers when no client is actively connected to it after a certain timeout.
This works great for our chat system because, we don’t want to keep the consumers alive even when the users aren’t active. Instead, we create a new ephemeral consumer when a user is online to get the job done.

Durable: This type of consumer retains the state across restarts. That’s the only primary difference when compared to ephemeral consumers. To make a consumer durable, we need to pass a consumerName property to the consumer config. To understand better, here is a durable version of the consumer config. Feel free to use any one of them.

 lastTimestamp := time.Now().Add(-5 * time.Minute)
 conversationID := "test_conversation_123"
 userID := "test_user_1"
 consumerName := fmt.Sprintf("chat_consumer_%s_%s", conversationID, userID)
 cons, err := s.js.Consumer(ctx, config.StreamName, consumerName)
 if err != nil {
  log.Printf("Consumer %s doesn't exist, creating new one: %v", consumerName, err)
  // Consumer doesn't exist, create it
  consumerConfig := jetstream.ConsumerConfig{
   FilterSubject:     subject,
   Name:              consumerName,
   Durable:           consumerName,
   AckPolicy:         jetstream.AckNonePolicy,
   DeliverPolicy:     jetstream.DeliverAllPolicy,
   //DeliverPolicy:     jetstream.DeliverByStartTimePolicy,
   InactiveThreshold: 10 * time.Minute, // consumer will be cleared after 10 minutes of no user being subscribed
   // OptStartTime:      lastTimestamp,
  }

  cons, err = s.js.CreateOrUpdateConsumer(ctx, config.StreamName, consumerConfig)
  if err != nil {
   return nil, fmt.Errorf("failed to create consumer: %w", err)
  }
  // log.Printf("Created new durable consumer %s for subject %s with DeliverNewPolicy", consumerName, subject)
 } else {
  log.Printf("Using existing durable consumer %s for subject %s", consumerName, subject)
 }

Now, with the above durable consumer setup, each user will get their own consumer for a given conversation.

Note: We are using DeliverAllPolicy without AckPolicy set to None is because if we acknowledge the messages, they won’t be replayed by the consumer. In our case, since the stream is the only source of the messages, we need the messages to be replayed even if they were already consumed so that the client can show the messages.
Read more about AckPolicy and DeliverPolicy here.

Once, the consumer is available, we can use different methods to consume messages of which one of the methods are called Consume itself, which accepts a callback function with the jetstream message as the argument. Now, this function processes messages real-time as and when the messages arrive.

The other method is Messages() which gives us access to an iterator which we can loop over.

// Messages() example
iter, _ := cons.Messages()
for {
    msg, err := iter.Next()
    // Next can return error, e.g. when iterator is closed or no heartbeats were received
    if err != nil {
        //handle error
    }
    fmt.Printf("Received a JetStream message: %s\n", string(msg.Data()))
    msg.Ack()
}
iter.Stop()

The third method is Fetch() which can accept the number of messages that we want to fetch at a time. This is more useful when we need pagination approach.

// Fetch() example
// receive up to 10 messages from the stream
msgs, err := c.Fetch(10)
if err != nil {
    // handle error
}

for msg := range msgs.Messages() {
    fmt.Printf("Received a JetStream message: %s\n", string(msg.Data()))
}

Consume and Messages are both useful for continuous polling while the Fetch method can be useful for batch processing.

At this point, we are done with the NATS related setup. Now, we can move on to the handling websockets.

We need three important functions here:

1. HandleWebsocket: This will take care of managing the lifecycle of the websocket connections.

2. HandleRead: This will take care of accepting the incoming messages and send them over to NATS.

3. HandleWrite: This will take care of sending back messages received from NATS.

Before getting into implementing the above functions, we need a helper function to create a new client.

// handlers/websocket.go
package handlers
import (
 "context"
 "log"
 "time"

 "your_module_name/config" // Adjust import path
 "your_module_name/models" // Adjust import path
 "your_module_name/nats_service" // Adjust import path

 "github.com/gofiber/contrib/websocket"
 "github.com/gofiber/fiber/v2"
 "github.com/google/uuid"
)

// Represents a connected client
type Client struct {
 Conn           *websocket.Conn
 NatsService    *nats_service.NatsService
 ConversationID string
 UserID         string // Should come from authentication
 UserAlias      string // Should come from authentication
 MessageChan    chan *models.Message // Channel for messages from NATS subscription
 DoneChan       chan struct{}        // Channel to signal closure
}

func NewClient(conn *websocket.Conn, natsSvc *nats_service.NatsService, convoID, userID, userAlias string) *Client {
 return &Client{
  Conn:           conn,
  NatsService:    natsSvc,
  ConversationID: convoID,
  UserID:         userID,
  UserAlias:      userAlias,
  MessageChan:    make(chan *models.Message, 256), // Buffered channel
  DoneChan:       make(chan struct{}),
 }
}

Now, let’s setup the HandleRead function.

// HandleRead reads messages from the WebSocket connection and sends them to NATS.
func (c *Client) HandleRead(ctx context.Context) {
 defer func() {
  log.Printf("Reader closed for %s in %s", c.UserID, c.ConversationID)
  close(c.DoneChan) // Signal writer to stop
  // NATS subscription cleanup happens in HandleWebSocket
 }()
 c.Conn.SetReadLimit(config.MaxMessageSize)
 c.Conn.SetReadDeadline(time.Now().Add(config.PongWait))
 c.Conn.SetPongHandler(func(string) error {
  c.Conn.SetReadDeadline(time.Now().Add(config.PongWait))
  return nil
 })

 for {
  // Read message from WebSocket client
  var clientMsg struct { // Expecting simple text messages from client
   Text string `json:"text"`
  }
  err := c.Conn.ReadJSON(&clientMsg)
  if err != nil {
   if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
    log.Printf("WebSocket read error for %s: %v", c.UserID, err)
   } else {
    log.Printf("WebSocket closed for %s: %v", c.UserID, err)
   }
   break // Exit loop on error or close
  }

  if clientMsg.Text == "" {
   continue // Ignore empty messages
  }

  // Create a message object
  msg := &models.Message{
   ID:             uuid.NewString(), // Generate unique ID
   ConversationID: c.ConversationID,
   SenderID:       c.UserID,
   SenderAlias:    c.UserAlias,
   Text:           clientMsg.Text,
   CreatedAt:      time.Now().UTC(),
  }

  // Publish the message to NATS
  pubCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
  if err := c.NatsService.PublishMessage(pubCtx, msg); err != nil {
   log.Printf("Failed to publish message from %s: %v", c.UserID, err)
   // Optionally notify the client of the failure
  }
  cancel()
 }
}

Here, we receive the message from the webocket connection, then, construct the message object as needed. And then, we use the NATS’ PublishMessage function to publish the message to the stream.

Now, that we have a message in the stream, let’s see how we get that back to the client through the HandleWrite function.

// HandleWrite writes messages from the NATS subscription (via MessageChan) to the WebSocket connection.
func (c *Client) HandleWrite() {
 ticker := time.NewTicker(config.PingPeriod)
 defer func() {
  ticker.Stop()
  log.Printf("Writer closed for %s in %s", c.UserID, c.ConversationID)
  // Connection closing is handled in HandleWebSocket's defer
 }()

 for {
  select {
  case message, ok := <-c.MessageChan:
   c.Conn.SetWriteDeadline(time.Now().Add(config.WriteWait))
   if !ok {
    // The message channel was closed.
    log.Printf("Message channel closed for %s", c.UserID)
    c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
    return
   }

   // Write message received from NATS to WebSocket client
   if err := c.Conn.WriteJSON(message); err != nil {
    log.Printf("WebSocket write error for %s: %v", c.UserID, err)
    return // Exit loop on write error
   }

  case <-ticker.C:
   // Send ping message periodically
   c.Conn.SetWriteDeadline(time.Now().Add(config.WriteWait))
   if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
    log.Printf("WebSocket ping error for %s: %v", c.UserID, err)
    return // Exit loop on ping error
   }

  case <-c.DoneChan:
   // HandleWrite closed, stop writing
   log.Printf("Received done signal for %s", c.UserID)
   return
  }
 }
}

This is simple function which just processes all the messages in the MessageChan.

A ticker is maintained to send ping messages to the client which helps retaining the connection for longer periods and also track the user’s presence.

With both the HandleWrite and HandleRead functions setup, let’s complete the puzzle with HandleWebsocket.

// HandleWebSocket manages the lifecycle of a WebSocket connection
func HandleWebSocket(c *websocket.Conn, natsSvc *nats_service.NatsService) {
 // ** IMPORTANT: Add Authentication/Authorization here! **
 // Extract UserID, UserAlias from token/session passed during upgrade
 userID := "user_" + uuid.NewString()[:6] // Placeholder
 userAlias := "Guest " + userID[5:]      // Placeholder

 conversationID := c.Params("conversationID")
 if conversationID == "" {
  log.Println("Missing conversationID parameter")
  c.WriteJSON(fiber.Map{"error": "Missing conversationID"}) // Use Fiber context if available, or handle error differently
  c.Close()
  return
 }

 client := NewClient(c, natsSvc, conversationID, userID, userAlias)
 log.Printf("Client %s connected to conversation %s", client.UserID, client.ConversationID)

 // Use a context for managing the NATS subscription lifecycle tied to the WS connection
 subCtx, cancelSub := context.WithCancel(context.Background())
 defer cancelSub() // Cancel the context when the handler exits

 // Subscribe to NATS messages for this conversation
 consumeCtx, err := client.NatsService.SubscribeToConversation(subCtx, client.ConversationID, func(msg *models.Message) {
  // This handler runs in the NATS message delivery goroutine
  // Send message to the client's buffered channel
  select {
  case client.MessageChan <- msg:
  case <-time.After(1 * time.Second): // Timeout to prevent blocking NATS delivery
   log.Printf("Timeout sending message to client %s channel", client.UserID)
  case <-client.DoneChan: // Check if client disconnected
   log.Printf("Client %s disconnected before message could be sent", client.UserID)
  }
 })
 if err != nil {
  log.Printf("Failed to subscribe client %s to %s: %v", client.UserID, client.ConversationID, err)
  c.Close()
  return
 }

 // Cleanup subscription and connection when handler exits
 defer func() {
  log.Printf("Cleaning up for client %s in %s", client.UserID, client.ConversationID)
  if consumeCtx != nil {
   consumeCtx.Stop() // Stop the NATS consumer
  }
  close(client.MessageChan) // Close channel *after* stopping consumer
  c.Close()                 // Close WebSocket connection
 }()

 // Start the write in a separate goroutine
 go client.HandleWrite()

 // Start the read pump (blocking call) in the main handler goroutine
 // It will exit when the connection closes or errors, triggering defers.
 client.HandleRead(subCtx) // Pass context for potential cancellation signals

 log.Printf("HandleWebSocket finished for %s", client.UserID)
}

We have almost got our chat server working now. But, let’s understand what the HandleWebsocket function is doing.

The most important part here is how we are using the SubscribeToConversation function that we setup earlier. As you can see, we are passing the handler function which feeds the messages to the MessageChan which we consume from again in the HandleWrite function.

Now to get the server running, let’s setup a main function.

// main.go
package main
import (
 "log"
 "os"
 "os/signal"
 "syscall"

 "your_module_name/config" // Adjust import path
 "your_module_name/handlers" // Adjust import path
 "your_module_name/nats_service" // Adjust import path

 "github.com/gofiber/fiber/v2"
 "github.com/gofiber/fiber/v2/middleware/logger"
 "github.com/gofiber/contrib/websocket"
)

func main() {
 // --- Initialize NATS Service ---
 natsSvc, err := nats_service.NewNatsService()
 if err != nil {
  log.Fatalf("Failed to initialize NATS Service: %v", err)
 }
 defer natsSvc.Close()
 log.Println("NATS Service Initialized")

 // --- Initialize Fiber App ---
 app := fiber.New()
 app.Use(logger.New()) // Basic request logging

 // --- Setup WebSocket Route ---
 app.Use("/ws", func(c *fiber.Ctx) error {
  // Check if the request is a WebSocket upgrade request
  if websocket.IsWebSocketUpgrade(c) {
   c.Locals("allowed", true)
   return c.Next()
  }
  return fiber.ErrUpgradeRequired
 })

 // WebSocket endpoint: /chat/:conversationID
 app.Get("/chat/:conversationID", websocket.New(func(c *websocket.Conn) {
  // Pass the NATS service instance to the handler
  handlers.HandleWebSocket(c, natsSvc)
 }))

 // --- Start Server ---
 go func() {
  log.Printf("Starting server on %s", config.ServerAddr)
  if err := app.Listen(config.ServerAddr); err != nil {
   log.Fatalf("Server failed to start: %v", err)
  }
 }()

 // --- Graceful Shutdown ---
 quit := make(chan os.Signal, 1)
 signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
 <-quit // Block until signal received

 log.Println("Shutting down server...")

 // Shutdown Fiber app
 if err := app.Shutdown(); err != nil {
  log.Printf("Error shutting down Fiber: %v", err)
 }

 // NATS connection is closed by defer in main

 log.Println("Server gracefully stopped")
}

Here, we are just creating a new NATS service, setup an endpoint for the websocket connections and also setup a graceful shutdown process. With this, we are now done with everything.

To connect to the server, use http clients like Postman and create a new websocket connection to the following url: ws://localhost:8080/chat/general and you should be seeing the logs in our server confirming that a user was connected to the chat.

And now, you can start sending messages in the JSON format like:

{
  "text": "hello!"
}

Now, if you connect to the same conversation from another tab in Postman, you’ll be receiving the above message as well!

All that’s left is to start chatting, all the best!

The complete code for the project is available here.