project: the server part of the chat is implemented

Author: GopherSuslik

connection.go

@@ -0,0 +1,38 @@
+package main
+
+import "github.com/gorilla/websocket"
+
+type connection struct {
+	ws   *websocket.Conn
+	send chan []byte
+	hub  *hub
+}
+
+func newConnection(ws *websocket.Conn, hub *hub) connection {
+	return connection{
+		ws:   ws,
+		send: make(chan []byte, 256),
+		hub:  hub,
+	}
+}
+
+func (c *connection) read() {
+	for {
+		_, message, err := c.ws.ReadMessage()
+		if err != nil {
+			break
+		}
+
+		c.hub.broadcast <- message
+	}
+	c.ws.Close()
+}
+
+func (c *connection) write() {
+	for message := range c.send {
+		if err := c.ws.WriteMessage(websocket.TextMessage, message); err != nil {
+			break
+		}
+	}
+	c.ws.Close()
+}

go.mod

@@ -1,3 +1,8 @@
-module github.com/Lapp-coder/websocket-server
+module github.com/Lapp-coder/websocket-chat
 
 go 1.16
+
+require (
+	github.com/gorilla/websocket v1.4.2
+	github.com/sirupsen/logrus v1.8.1
+)

go.sum

@@ -1,2 +1,9 @@
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
+github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

handler.go

@@ -0,0 +1,45 @@
+package main
+
+import (
+	"fmt"
+	"github.com/gorilla/websocket"
+	"net/http"
+)
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
+type handler struct {
+	hub *hub
+}
+
+func newHandler(hub *hub) *handler {
+	return &handler{hub: hub}
+}
+
+func (h *handler) initRoutes() {
+	http.HandleFunc("/", h.index)
+	http.HandleFunc("/chat", h.chat)
+}
+
+func (h *handler) index(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprint(w, "<h1>Chat on websockets</h1>")
+}
+
+func (h *handler) chat(w http.ResponseWriter, r *http.Request) {
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return
+	}
+
+	conn := newConnection(ws, h.hub)
+	conn.hub.register <- &conn
+	defer func() {
+		conn.hub.unregister <- &conn
+	}()
+
+	go conn.write()
+	conn.read()
+}

hub.go

@@ -0,0 +1,35 @@
+package main
+
+type hub struct {
+	connections map[*connection]struct{}
+	broadcast   chan []byte
+	register    chan *connection
+	unregister  chan *connection
+}
+
+func newHub() hub {
+	return hub{
+		connections: make(map[*connection]struct{}),
+		broadcast:   make(chan []byte),
+		register:    make(chan *connection),
+		unregister:  make(chan *connection),
+	}
+}
+
+func (h *hub) run() {
+	for {
+		select {
+		case conn := <-h.register:
+			h.connections[conn] = struct{}{}
+		case conn := <-h.unregister:
+			if _, ok := h.connections[conn]; ok {
+				delete(h.connections, conn)
+				close(conn.send)
+			}
+		case message := <-h.broadcast:
+			for conn := range h.connections {
+				conn.send <- message
+			}
+		}
+	}
+}

main.go

@@ -1,5 +1,38 @@
 package main
 
+import (
+	"context"
+	"github.com/sirupsen/logrus"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
 func main() {
+	hub := newHub()
+	handler := newHandler(&hub)
+	handler.initRoutes()
+	go hub.run()
+
+	host := os.Getenv("CHAT_HOST")
+	port := os.Getenv("CHAT_PORT")
+	srv := newServer(host, port)
+
+	go func() {
+		if err := srv.start(); err != nil {
+			logrus.Errorf("failed to start server: %s", err.Error())
+		}
+	}()
+
+	logrus.Info("server started")
+
+	shutdown := make(chan os.Signal)
+	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)
+	<-shutdown
+
+	logrus.Info("server shutdown")
 
+	if err := srv.shutdown(context.Background()); err != nil {
+		logrus.Errorf("failed to graceful shutdown server: %s", err.Error())
+	}
 }

server.go

@@ -0,0 +1,26 @@
+package main
+
+import (
+	"context"
+	"net/http"
+)
+
+type server struct {
+	httpServer *http.Server
+}
+
+func newServer(host, port string) *server {
+	return &server{
+		httpServer: &http.Server{
+			Addr: host + ":" + port,
+		},
+	}
+}
+
+func (s *server) start() error {
+	return s.httpServer.ListenAndServe()
+}
+
+func (s *server) shutdown(ctx context.Context) error {
+	return s.httpServer.Shutdown(ctx)
+}