Skip to content

Instantly share code, notes, and snippets.

@gjtorikian
Created October 13, 2020 15:43
Show Gist options
  • Save gjtorikian/8894dec140a6e57934572f5b447f6d51 to your computer and use it in GitHub Desktop.
Save gjtorikian/8894dec140a6e57934572f5b447f6d51 to your computer and use it in GitHub Desktop.
Heroku chat sample
package main
import (
"encoding/json"
"io"
"log"
"net/http"
"os"
"github.com/go-redis/redis"
"github.com/gorilla/websocket"
"github.com/joho/godotenv"
)
type ChatMessage struct {
Username string `json:"username"`
Text string `json:"text"`
}
var (
rdb *redis.Client
)
var clients = make(map[*websocket.Conn]bool)
var broadcaster = make(chan ChatMessage)
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
func handleConnections(w http.ResponseWriter, r *http.Request) {
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Fatal(err)
}
// ensure connection close when function returns
defer ws.Close()
clients[ws] = true
// if it's zero, no messages were ever sent/saved
if rdb.Exists("chat_messages").Val() != 0 {
sendPreviousMessages(ws)
}
for {
var msg ChatMessage
// Read in a new message as JSON and map it to a Message object
err := ws.ReadJSON(&msg)
if err != nil {
delete(clients, ws)
break
}
// send new message to the channel
broadcaster <- msg
}
}
func sendPreviousMessages(ws *websocket.Conn) {
chatMessages, err := rdb.LRange("chat_messages", 0, -1).Result()
if err != nil {
panic(err)
}
// send previous messages
for _, chatMessage := range chatMessages {
var msg ChatMessage
json.Unmarshal([]byte(chatMessage), &msg)
messageClient(ws, msg)
}
}
// If a message is sent while a client is closing, ignore the error
func unsafeError(err error) bool {
return !websocket.IsCloseError(err, websocket.CloseGoingAway) && err != io.EOF
}
func handleMessages() {
for {
// grab any next message from channel
msg := <-broadcaster
storeInRedis(msg)
messageClients(msg)
}
}
func storeInRedis(msg ChatMessage) {
json, err := json.Marshal(msg)
if err != nil {
panic(err)
}
if err := rdb.RPush("chat_messages", json).Err(); err != nil {
panic(err)
}
}
func messageClients(msg ChatMessage) {
// send to every client currently connected
for client := range clients {
messageClient(client, msg)
}
}
func messageClient(client *websocket.Conn, msg ChatMessage) {
err := client.WriteJSON(msg)
if err != nil && unsafeError(err) {
log.Printf("error: %v", err)
client.Close()
delete(clients, client)
}
}
func main() {
env := os.Getenv("GO_ENV")
if "" == env {
err := godotenv.Load()
if err != nil {
log.Fatal("Error loading .env file")
}
}
port := os.Getenv("PORT")
redisURL := os.Getenv("REDIS_URL")
opt, err := redis.ParseURL(redisURL)
if err != nil {
panic(err)
}
rdb = redis.NewClient(opt)
http.Handle("/", http.FileServer(http.Dir("./public")))
http.HandleFunc("/websocket", handleConnections)
go handleMessages()
log.Print("Server starting at localhost:" + port)
if err := http.ListenAndServe(":"+port, nil); err != nil {
log.Fatal(err)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment