Improved network code

server/core/Server.go

@@ -3,6 +3,8 @@ package core
 import (
 	"fmt"
 	"net"
+	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -12,15 +14,38 @@ type Handler func([]byte, *Client)
 // Server represents a UDP server.
 type Server struct {
 	socket   *net.UDPConn
-	clients  map[string]*Client
 	handlers [256]Handler
+	clients  sync.Map
+	count    atomic.Int64
 }
 
 // New creates a new server.
 func New() *Server {
-	return &Server{
-		clients: make(map[string]*Client),
-	}
+	timeout := 3 * time.Second
+	interval := time.Second
+	server := &Server{}
+
+	go func() {
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+
+		for range ticker.C {
+			now := time.Now()
+
+			server.clients.Range(func(key, value interface{}) bool {
+				item := value.(*Client)
+
+				if !item.lastPacket.IsZero() && now.After(item.lastPacket.Add(timeout)) {
+					server.clients.Delete(key)
+					server.count.Add(-1)
+				}
+
+				return true
+			})
+		}
+	}()
+
+	return server
 }
 
 // AddHandler adds the handler for the given byte code.
@@ -28,8 +53,14 @@ func (s *Server) AddHandler(code byte, handler Handler) {
 	s.handlers[code] = handler
 }
 
-// SendTo sends the data to a client.
-func (s *Server) SendTo(data []byte, client *Client) {
+// Count returns the number of connected clients.
+func (s *Server) Count() int {
+	return int(s.count.Load())
+}
+
+// Send sends the data prefixed with the byte code to the client.
+func (s *Server) Send(code byte, data []byte, client *Client) {
+	data = append([]byte{code}, data...)
 	_, err := s.socket.WriteToUDP(data, client.address)
 
 	if err != nil {
@@ -91,25 +122,26 @@ func (s *Server) handle(data []byte, addr *net.UDPAddr) {
 	handler := s.handlers[data[0]]
 
 	if handler == nil {
-		fmt.Println("Unknown packet type.")
+		fmt.Printf("No callback registered for packet type %d\n", data[0])
 		return
 	}
 
-	handler(data, c)
+	handler(data[1:], c)
 }
 
 // getClient either returns a new or existing client for the requested address.
 func (s *Server) getClient(addr *net.UDPAddr) *Client {
-	c, exists := s.clients[addr.String()]
+	obj, exists := s.clients.Load(addr.String())
 
 	if exists {
-		return c
+		return obj.(*Client)
 	}
 
-	c = &Client{
+	client := &Client{
 		address: addr,
 	}
 
-	s.clients[addr.String()] = c
-	return c
+	s.clients.Store(addr.String(), client)
+	s.count.Add(1)
+	return client
 }