How GIGO uses NATS to talk across its cluster

// format create workspace request and marshall it with gob
buf := bytes.NewBuffer(nil)
encoder := gob.NewEncoder(buf)
err = encoder.Encode(models2.CreateWorkspaceMsg{
  WorkspaceID: workspace.ID,
  OwnerID:     callingUser.ID,
  OwnerEmail:  callingUser.Email,
  OwnerName:   callingUser.UserName,
  Disk:        wsConfig.Resources.Disk,
  CPU:         wsConfig.Resources.CPU,
  Memory:      wsConfig.Resources.Mem,
  Container:   wsConfig.BaseContainer,
  AccessUrl:   accessUrl,
})
if err != nil {
  _ = tx.Rollback()
  return nil, fmt.Errorf("failed to encode workspace: %v", err)
}

// send workspace create message to jetstream so a follower will
// create the workspace
_, err = js.PublishAsync(streams.SubjectWorkspaceCreate, buf.Bytes())
if err != nil {
  _ = tx.Rollback()
  return nil, fmt.Errorf("failed to send workspace create message to jetstream: %v", err)

func asyncCreateWorkspace(nodeId int64, tidb *ti.Database, wsClient *ws.WorkspaceClient, wsStatusUpdater *utils.WorkspaceStatusUpdater,
 js *mq.JetstreamClient, msg *nats.Msg, logger logging.Logger) {
  ctx, span := otel.Tracer("gigo-core").Start(context.TODO(), "async-create-workspace-routine")
  callerName := "asyncCreateWorkspace"

  // unmarshall create workspace request message
  var createWsMsg models2.CreateWorkspaceMsg
  decoder := gob.NewDecoder(bytes.NewBuffer(msg.Data))
  err := decoder.Decode(&createWsMsg)
  if err != nil {
    logger.Errorf("(workspace: %d) failed to decode create workspace message: %v", nodeId, err)
    return
  }

  // log that we started provisioning
  logger.Infof("(workspace: %d) creating workspace %d", nodeId, createWsMsg.WorkspaceID)

  // more code
}

// more code

// process workspace create stream
processStream(
  nodeId,
  js,
  workerPool,
  streams.StreamWorkspace,
  streams.SubjectWorkspaceCreate,
  "gigo-core-follower-ws-create",
  time.Minute*10,
  "workspace",
  logger,
  func(msg *nats.Msg) {
    asyncCreateWorkspace(nodeId, tidb, wsClient, wsStatusUpdater, js, msg, logger)
  },
)

func WorkspaceInitializationStepCompleted(ctx context.Context, tidb *ti.Database, wsStatusUpdater *utils.WorkspaceStatusUpdater, wsId int64,
 step models.WorkspaceInitState) error {
 ctx, span := otel.Tracer("gigo-core").Start(ctx, "workspace-initialization-step-completed-core")
 defer span.End()
 callerName := "WorkspaceInitializationStepCompleted"

 // default workspace state to starting but set to Active if this is the VSCodeLaunch init step
 state := models.WorkspaceStarting
 if step == models.WorkspaceInitCompleted {
  state = models.WorkspaceActive
 }

 // perform update directly on the workspace row
 _, err := tidb.ExecContext(ctx, &span, &callerName,
  "update workspaces set init_state = ?, state = ?, last_state_update = ? where _id = ?",
  step, state, time.Now(), wsId,
 )
 if err != nil {
  // handle not found
  if err == sql.ErrNoRows {
   return fmt.Errorf("workspace not found")
  }
  // handle error
  return fmt.Errorf("failed to update workspace: %v", err)
 }

 // push workspace status update to subscribers
 wsStatusUpdater.PushStatus(ctx, wsId, nil)

 return nil
}

func (u *WorkspaceStatusUpdater) PushStatus(ctx context.Context, id int64, ws *models.Workspace) error {
 // retrieve the workspace if we did not receive one
 if ws == nil {
  var err error
  ws, err = u.getWorkspace(ctx, id)
  if err != nil {
   return err
  }
 }

 // increment init state before writing the update to represent the step that we
 // are working on unless we are finished then we return the true code
 if ws.InitState != models.WorkspaceInitCompleted {
  ws.InitState++
 }

 // gob encode the status update
 buf := bytes.NewBuffer(nil)
 encoder := gob.NewEncoder(buf)
 err := encoder.Encode(models2.WorkspaceStatusUpdateMsg{
  Workspace: ws.ToFrontend(u.Hostname, u.Tls),
 })
 if err != nil {
  return fmt.Errorf("failed to serialize workspace status update: %v", err)
 }

 // send workspace status update message to jetstream so any listening websocket
 // connections can properly update the client for the current state of the workspace
 _, err = u.Js.PublishAsync(fmt.Sprintf(streams.SubjectWorkspaceStatusUpdateDynamic, ws.ID), buf.Bytes())
 if err != nil {
  return fmt.Errorf("failed to send workspace status update message to jetstream: %v", err)
 }

 return nil
}

 // format the message and send it to the client
 p.outputChan <- ws.PrepMessage[any](
  msg.SequenceID,
  ws.MessageTypeWorkspaceStatusUpdate,
  status,
 )

 // exit if we already have a subscription for this workspace
 p.mu.Lock()
 defer p.mu.Unlock()
 if _, ok := p.wsSubs[workspaceId]; ok {
  return
 }

 // create a subscriber to workspace status events
 sub, err := p.s.jetstreamClient.Subscribe(
  fmt.Sprintf(streams.SubjectWorkspaceStatusUpdateDynamic, workspaceId),
  p.workspaceStatusCallback,
 )
 if err != nil {
  p.socket.logger.Errorf("failed to create workspace status subscriber: %v", err)
  // handle internal server error via websocket
  p.outputChan <- ws.PrepMessage[any](
   msg.SequenceID,
   ws.MessageTypeGenericError,
   ws.GenericErrorPayload{
    Code:  ws.ResponseCodeServerError,
    Error: "internal server error occurred",
   },
  )
  return
 }

 // save the subscription
 p.wsSubs[workspaceId] = sub

func (p *WebSocketPluginChat) handleNewChatMessage(msg *ws.Message[any]) {
 // boring validation stuff

 // execute core function logic
 chat, event, err := core.CreateChat(p.ctx, p.s.tiDB, p.s.sf, p.socket.user.Load(), createChatParams)
 if err != nil {
  p.socket.logger.Errorf("failed to execute create chat internal: %v", err)
  // handle internal server error via websocket
  p.outputChan <- ws.PrepMessage[any](
   msg.SequenceID,
   ws.MessageTypeGenericError,
   ws.GenericErrorPayload{
    Code:  ws.ResponseCodeServerError,
    Error: "internal server error occurred",
   },
  )
  return
 }

 // some code...

 // subscribe to the new chat stream

 // create an async subscriber to broadcast init messages
 sub, err := p.s.jetstreamClient.Subscribe(
  fmt.Sprintf(streams.SubjectChatMessagesDynamic, chat.ID),
  p.chatMsgHandler,
  nats.AckExplicit(),
  nats.DeliverNew(),
 )
 if err != nil {
  p.socket.logger.Errorf("failed to subscribe to new chat stream: %v", err)
  return
 }

 // save the chat and subscription to the map
 p.mu.Lock()
 p.subs[chat.ID] = sub
 p.chats[chat.ID] = chat
 p.mu.Unlock()

 // marshall message using gob
 var buf bytes.Buffer
 encoder := gob.NewEncoder(&buf)
 err = encoder.Encode(models2.NewChatMsg{
  Chat: *chat,
 })
 if err != nil {
  p.socket.logger.Errorf("failed to encode message: %v", err)
  return
 }

 // save the buffer bytes to a variable
 chatBytes := buf.Bytes()

 // forward the message to all the active users so they can subscribe to the chat
 for _, user := range chat.Users {
  // skip if the user is the calling user since me just manually subscribed
  if user == p.socket.user.Load().ID {
   continue
  }

  _, err = p.s.jetstreamClient.Publish(
   fmt.Sprintf(streams.SubjectChatNewChatDynamic, user),
   chatBytes,
  )
  if err != nil {
   p.socket.logger.Errorf("failed to publish message to chat stream: %v", err)
  }
 }

 // if the event is nil exit
 if event == nil {
  return
 }

 // iterate over all the users including and send the event

 // marshall message using gob
 var eventBuf bytes.Buffer
 encoder = gob.NewEncoder(&eventBuf)
 err = encoder.Encode(event)
 if err != nil {
  p.socket.logger.Errorf("failed to encode message: %v", err)
  return
 }

 // save the buffer bytes to a variable
 eventBytes := eventBuf.Bytes()

 // forward the message to all the active users so the new chat is populated in their UI
 for _, user := range chat.Users {
  // skip if the user is the calling user since me just manually subscribed
  if user == p.socket.user.Load().ID {
   continue
  }

  _, err = p.s.jetstreamClient.Publish(
   fmt.Sprintf(streams.SubjectChatUpdatedDynamic, user),
   eventBytes,
  )
  if err != nil {
   p.socket.logger.Errorf("failed to publish message to chat stream: %v", err)
  }
 }
}

// create function to handle stream messages for new chats
  newChatHandler := func(m *nats.Msg) {
   // we always ack the message
   _ = m.Ack()

   // parse the message
   var newChat models2.NewChatMsg
   decoder := gob.NewDecoder(bytes.NewBuffer(m.Data))
   err := decoder.Decode(&newChat)
   if err != nil {
    s.logger.Errorf("failed to decode new chat message: %v", err)
    return
   }

   // skip if we already have a subscription for this chat
   lock.Lock()
   defer lock.Unlock()
   if _, ok := subs[newChat.Chat.ID]; ok {
    return
   }

   // create an async subscriber to broadcast init messages
   sub, err := s.jetstreamClient.Subscribe(
    fmt.Sprintf(streams.SubjectChatMessagesDynamic, newChat.Chat.ID),
    chatMessageHandler,
    nats.AckExplicit(),
    nats.DeliverNew(),
   )
   if err != nil {
    s.logger.Errorf("failed to subscribe to new chat stream: %v", err)
    return
   }

   subs[newChat.Chat.ID] = sub
   chatMap[newChat.Chat.ID] = &newChat.Chat

   // send the message to the client to add the chat
   outputChan <- ws.PrepMessage[any](
    "",
    ws.MessageTypeNewChatBroadcast,
    newChat.Chat.ToFrontend(),
   )
  }

  // create an async subscriber to broadcast init messages
  activeChatStream, err = s.jetstreamClient.Subscribe(
   fmt.Sprintf(streams.SubjectChatNewChatDynamic, callingUser.ID),
   newChatHandler,
   nats.AckExplicit(),
   nats.DeliverNew(),
  )
  if err != nil {
   return nil, fmt.Errorf("failed to subscribe to new chat stream: %w", err)
  }