package main
import (
"fmt"
"log"
"net/http"
"sync"
"time"
"github.com/google/uuid"
"github.com/gorilla/websocket"
)
var m sync.Mutex
type Socket struct {
Conn *websocket.Conn
WriteChan chan []byte
ID string
}
type Repo struct {
Member []*Socket
N int
}
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
var repo *Repo
func init() {
repo = &Repo{
Member: make([]*Socket, 0),
N: 0,
}
upgrader.CheckOrigin = func(r *http.Request) bool {
return true
}
}
func (socket *Socket) writeMessage(p []byte) {
socket.WriteChan <- p
}
func (socket *Socket) processTextMessage(p []byte) {
socket.writeMessage(p)
}
func (socket *Socket) writerListener() {
for {
data := <-socket.WriteChan
socket.Conn.WriteMessage(websocket.TextMessage, data)
}
}
func (socket *Socket) SocketHandler() {
go socket.writerListener()
for {
messageType, p, err := socket.Conn.ReadMessage()
if err != nil {
return
}
if messageType == websocket.TextMessage {
socket.processTextMessage(p)
}
}
}
func (socket *Socket) Delete() {
m.Lock()
for i, v := range repo.Member {
if v.ID == socket.ID {
if repo.N == 1 {
repo.N = 0
break
}
repo.N--
repo.Member[i] = repo.Member[repo.N-1]
break
}
}
m.Unlock()
}
func (socket *Socket) Register() {
m.Lock()
repo.N++
repo.Member = append(repo.Member, socket)
m.Unlock()
}
func handler(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
defer conn.Close()
if err != nil {
w.Write([]byte(err.Error()))
return
}
s := &Socket{
Conn: conn,
WriteChan: make(chan []byte),
ID: uuid.New().String(),
}
defer s.Delete()
s.Register()
s.SocketHandler()
}
func main() {
// test broadcast
go func() {
time.Sleep(4 * time.Second)
fmt.Println("=========")
for {
time.Sleep(2 * time.Second)
for _, v := range repo.Member {
v.writeMessage([]byte(fmt.Sprintf("%v", time.Now().UTC().String())))
}
}
}()
http.HandleFunc("/ws", handler)
log.Fatal(http.ListenAndServe(":8081", nil))
}