From e8fd5ac2a3b47007355de001fcc8e9fba6b0927b Mon Sep 17 00:00:00 2001 From: whysman Date: Tue, 3 Dec 2024 15:23:18 -0500 Subject: [PATCH] Refactored app for SSE --- main.go | 244 ++++++++++++++++++++++---------------------------------- 1 file changed, 96 insertions(+), 148 deletions(-) diff --git a/main.go b/main.go index 06f055e..22117bb 100644 --- a/main.go +++ b/main.go @@ -23,8 +23,8 @@ var ( return true }, } - redisClient *redis.Client - ctx = context.Background() + //redisClient *redis.Client + //ctx = context.Background() ) type Message struct { @@ -42,14 +42,14 @@ type Config struct { } } -type Client struct { - conn *websocket.Conn -} - -type Server struct { - clients map[*Client]bool - broadcast chan Message - mu sync.Mutex +type App struct { + redisClient *redis.Client + sseClients map[chan string]bool + sseChannel chan string + mu sync.Mutex + ctx context.Context + flusher http.Flusher + ticker *time.Ticker } func getConfig(configPath string) (*Config, error) { @@ -87,8 +87,8 @@ func loadConfig() *Config { return cfg } -func initRedis(cfg Config) { - redisClient = redis.NewClient(&redis.Options{ +func initRedis(cfg Config, ctx context.Context) *redis.Client { + redisClient := redis.NewClient(&redis.Options{ Addr: cfg.Redis.Host + ":" + cfg.Redis.Port, }) _, err := redisClient.Do(context.Background(), "CONFIG", "SET", "notify-keyspace-events", "KEA").Result() @@ -97,139 +97,87 @@ func initRedis(cfg Config) { os.Exit(1) } redisClient.ConfigSet(ctx, "notify-keyspace-events", "Ex") + return redisClient } -func newServer() *Server { - return &Server{ - clients: make(map[*Client]bool), - broadcast: make(chan Message), +func newApp(cfg Config) *App { + return &App{ + sseClients: make(map[chan string]bool), + sseChannel: make(chan string), + redisClient: initRedis(cfg, context.Background()), + ctx: context.Background(), + ticker: time.NewTicker(5 * time.Second), } } -func (s *Server) addClient(client *Client) { - s.mu.Lock() - defer s.mu.Unlock() - s.clients[client] = true - broadcastAllRecordsToClient(client) -} +func (app *App) handleSSE(w http.ResponseWriter, r *http.Request) { + // Create a new SSE client channel + clientChan := make(chan string) + log.Printf("Clientchan: %v %v", clientChan, len(app.sseClients)) + // Register the client channel + app.mu.Lock() + app.sseClients[clientChan] = true + app.mu.Unlock() -func (s *Server) removeClient(client *Client) { - s.mu.Lock() - defer s.mu.Unlock() - if _, ok := s.clients[client]; ok { - delete(s.clients, client) - err := client.conn.Close() + // Send headers for SSE + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + + // Stream messages to the client + for msg := range clientChan { + _, err := fmt.Fprintf(w, "data: %s\n\n", msg) + log.Printf("SSE: %v", msg) if err != nil { - return + log.Fatal(err) } + if flusher, ok := w.(http.Flusher); ok { + app.flusher = flusher + app.flusher.Flush() + } + } + + // Clean up when done + app.mu.Lock() + log.Println("Deleting %v", clientChan) + delete(app.sseClients, clientChan) + close(clientChan) + app.mu.Unlock() +} + +func (app *App) broadcastSSE() { + for msg := range app.sseChannel { + app.mu.Lock() + for clientChan := range app.sseClients { + clientChan <- msg + } + app.mu.Unlock() } } -func (s *Server) handleConnections(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) - if err != nil { - fmt.Printf("Error upgrading connection: %v\n", err) - return - } - - client := &Client{conn: conn} - s.addClient(client) - - defer s.removeClient(client) - - for { - _, msgContent, err := conn.ReadMessage() - if err != nil { - fmt.Printf("Error reading message: %v\n", err) - break - } - - var message Message - err = json.Unmarshal(msgContent, &message) - if err != nil { - log.Fatalf("Unable to marshal JSON due to %s", err) - } - fmt.Printf("Broadcasting ws message to %d clients", s.clients) - s.broadcast <- message - } -} - -func broadcastAllRecords(s *Server) { - allRecords, err := fetchAllRecords() - if err != nil { - log.Println("Error fetching all records:", err) - } - if len(allRecords) > 0 { - var message Message - for _, msgContent := range allRecords { - err = json.Unmarshal([]byte(msgContent), &message) - if err != nil { - log.Println("Unable to marshal JSON due to: ", err) - } - fmt.Printf("Broadcasting %s,%s,%s,%s to %d clients\n", message.Id, message.Status, message.Name, message.Timestamp, len(s.clients)) - - for client := range s.clients { - if err := client.conn.WriteMessage(websocket.TextMessage, []byte(msgContent)); err != nil { - log.Println("Failed to broadcast update:", err) - } - } - } - } -} - -func broadcastRemovedRecords(s *Server, message Message) { +func broadcastRemovedRecords(app *App, message Message) { var msgJSON, err = json.Marshal(message) if err != nil { log.Println("Error marshalling json:", err) return } - fmt.Printf("Broadcasting removal: %s,%s,%s,%s to %d clients\n", message.Id, message.Status, message.Name, message.Timestamp, len(s.clients)) - - for client := range s.clients { - if err := client.conn.WriteMessage(websocket.TextMessage, msgJSON); err != nil { - log.Println("Failed to broadcast update:", err) - return - } - } + fmt.Printf("Broadcasting removal: %s,%s,%s,%s to %d clients\n", message.Id, message.Status, message.Name, message.Timestamp, len(app.sseClients)) + app.sseChannel <- string(msgJSON) } -func broadcastExpiredRecords(s *Server, removed string) { - var message = Message{Id: removed, Name: "", Image: "", Status: "removed", Timestamp: ""} +func broadcastExpiredRecords(app *App, expired string) { + var message = Message{Id: expired, Name: "", Image: "", Status: "removed", Timestamp: ""} var msgJSON, err = json.Marshal(message) if err != nil { log.Println("Error marshalling json:", err) return } - fmt.Printf("Broadcasting expiration: %s to %s clients\n", string(msgJSON), len(s.clients)) - for client := range s.clients { - if err := client.conn.WriteMessage(websocket.TextMessage, msgJSON); err != nil { - log.Println("Failed to broadcast update:", err) - return - } - } + fmt.Printf("Broadcasting expiration: %s to %s clients\n", string(msgJSON), len(app.sseClients)) + app.sseChannel <- string(msgJSON) } -func broadcastAllRecordsToClient(c *Client) { - allRecords, err := fetchAllRecords() - if err != nil { - log.Println("Error fetching all records:", err) - } - fmt.Printf("Broadcasting %d records to client\n", len(allRecords)) - var message Message - for _, msgContent := range allRecords { - err = json.Unmarshal([]byte(msgContent), &message) - if err != nil { - log.Println("Unable to marshal JSON due to: ", err) - } - fmt.Printf("Broadcasting %s,%s,%s,%s\n", message.Id, message.Status, message.Name, message.Timestamp) - if err := c.conn.WriteMessage(websocket.TextMessage, []byte(msgContent)); err != nil { - log.Println("Failed to broadcast update:", err) - } - } -} - -func (s *Server) listenForExpirationEvents() { - ps := redisClient.PSubscribe(ctx, "__keyevent@0__:expired") +func (app *App) listenForExpirationEvents() { + ps := app.redisClient.PSubscribe(app.ctx, "__keyevent@0__:expired") defer func(ps *redis.PubSub) { err := ps.Close() if err != nil { @@ -238,24 +186,23 @@ func (s *Server) listenForExpirationEvents() { }(ps) for msg := range ps.Channel() { - // Extract expired payload ID - expiredID := msg.Payload - fmt.Println(expiredID) + //Print expired payload ID + fmt.Println(msg.Payload) // Broadcast expiration event - broadcastExpiredRecords(s, expiredID) + broadcastExpiredRecords(app, msg.Payload) fmt.Printf("Done Broadcasting Expiration\n") } } -func fetchAllRecords() (map[string]string, error) { - allKeys, err := redisClient.Keys(ctx, "*").Result() +func (app *App) fetchAllRecords() (map[string]string, error) { + allKeys, err := app.redisClient.Keys(app.ctx, "*").Result() if err != nil { return nil, fmt.Errorf("could not fetch keys: %v", err) } records := make(map[string]string) for _, key := range allKeys { - value, err := redisClient.Get(ctx, key).Result() + value, err := app.redisClient.Get(app.ctx, key).Result() if errors.Is(err, redis.Nil) { continue // skip if key does not exist } else if err != nil { @@ -267,7 +214,7 @@ func fetchAllRecords() (map[string]string, error) { return records, nil } -func getState(w http.ResponseWriter, r *http.Request) { +func getStatus(w http.ResponseWriter, r *http.Request, app *App) { if r.Method != http.MethodPost { http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed) return @@ -290,7 +237,7 @@ func getState(w http.ResponseWriter, r *http.Request) { return } fmt.Println(request.Id) - value, err := redisClient.Get(ctx, request.Id).Result() + value, err := app.redisClient.Get(app.ctx, request.Id).Result() if errors.Is(err, redis.Nil) { message := Message{Id: request.Id, Status: "none", Timestamp: time.Now().Format(time.RFC3339)} w.Header().Set("Content-Type", "application/json") @@ -320,7 +267,7 @@ func getState(w http.ResponseWriter, r *http.Request) { } } -func setState(w http.ResponseWriter, r *http.Request, s *Server) { +func setStatus(w http.ResponseWriter, r *http.Request, app *App) { if r.Method != http.MethodPost { http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed) return @@ -334,12 +281,12 @@ func setState(w http.ResponseWriter, r *http.Request, s *Server) { } if message.Status == "none" { - if err := redisClient.Del(ctx, message.Id).Err(); err != nil { + if err := app.redisClient.Del(app.ctx, message.Id).Err(); err != nil { log.Printf("Error deleting key from Redis: %v", err) http.Error(w, "Failed to delete key from Redis", http.StatusInternalServerError) return } - broadcastRemovedRecords(s, message) + broadcastRemovedRecords(app, message) w.WriteHeader(http.StatusOK) return } @@ -351,23 +298,22 @@ func setState(w http.ResponseWriter, r *http.Request, s *Server) { } timeout := 20 * time.Second - err = redisClient.Set(ctx, message.Id, msgJSON, timeout).Err() + err = app.redisClient.Set(app.ctx, message.Id, msgJSON, timeout).Err() if err != nil { log.Printf("Failed to set key in Redis: %v", err) http.Error(w, "Failed to store data", http.StatusInternalServerError) return } - - broadcastAllRecords(s) + app.sseChannel <- string(msgJSON) w.WriteHeader(http.StatusOK) } -func updateState(w http.ResponseWriter, r *http.Request) { +func updateMessages(w http.ResponseWriter, r *http.Request, app *App) { if r.Method != http.MethodGet { http.Error(w, "Only GET method is allowed", http.StatusMethodNotAllowed) return } - allRecords, err := fetchAllRecords() + allRecords, err := app.fetchAllRecords() if err != nil { log.Printf("Error fetching records: %v", err) http.Error(w, "Failed to fetch records", http.StatusInternalServerError) @@ -411,23 +357,25 @@ func enableCORS(next http.Handler) http.Handler { } func main() { - initRedis(*loadConfig()) - - server := newServer() + app := newApp(*loadConfig()) router := mux.NewRouter() // Register routes on the mux router - router.HandleFunc("/ws", server.handleConnections).Methods("GET") - router.HandleFunc("/get", getState).Methods("POST") - router.HandleFunc("/set", func(w http.ResponseWriter, r *http.Request) { - setState(w, r, server) + router.HandleFunc("/events", app.handleSSE).Methods("GET") + router.HandleFunc("/getStatus", func(w http.ResponseWriter, r *http.Request) { + getStatus(w, r, app) }).Methods("POST") - router.HandleFunc("/update", updateState).Methods("GET") - corsRouter := enableCORS(router) + router.HandleFunc("/setStatus", func(w http.ResponseWriter, r *http.Request) { + setStatus(w, r, app) + }).Methods("POST") + router.HandleFunc("/updateMessages", func(w http.ResponseWriter, r *http.Request) { + updateMessages(w, r, app) + }).Methods("GET") // Start server and other necessary goroutines - go server.listenForExpirationEvents() - + go app.listenForExpirationEvents() + go app.broadcastSSE() + corsRouter := enableCORS(router) // Pass the mux router to ListenAndServe fmt.Println("Server started on :8080") err := http.ListenAndServe(":8080", corsRouter)