Refactored app for SSE

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)