first commit
This commit is contained in:
lendry
273
apps/media-ws/main.go
Normal file
273
apps/media-ws/main.go
Normal file
@@ -0,0 +1,273 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/rabbitmq/amqp091-go"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type Notification struct {
|
||||
UserID string `json:"userId"`
|
||||
Type string `json:"type"`
|
||||
Title string `json:"title"`
|
||||
Message string `json:"message"`
|
||||
Payload json.RawMessage `json:"payload,omitempty"`
|
||||
}
|
||||
|
||||
type Hub struct {
|
||||
mu sync.RWMutex
|
||||
connections map[string]map[*websocket.Conn]struct{}
|
||||
redis *redis.Client
|
||||
}
|
||||
|
||||
func NewHub(redisClient *redis.Client) *Hub {
|
||||
return &Hub{
|
||||
connections: make(map[string]map[*websocket.Conn]struct{}),
|
||||
redis: redisClient,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) Add(ctx context.Context, userID string, conn *websocket.Conn) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.connections[userID] == nil {
|
||||
h.connections[userID] = make(map[*websocket.Conn]struct{})
|
||||
}
|
||||
h.connections[userID][conn] = struct{}{}
|
||||
_ = h.redis.SAdd(ctx, "ws:online_users", userID).Err()
|
||||
_ = h.redis.Set(ctx, "ws:user:"+userID, "online", 90*time.Second).Err()
|
||||
}
|
||||
|
||||
func (h *Hub) Remove(ctx context.Context, userID string, conn *websocket.Conn) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
delete(h.connections[userID], conn)
|
||||
if len(h.connections[userID]) == 0 {
|
||||
delete(h.connections, userID)
|
||||
_ = h.redis.SRem(ctx, "ws:online_users", userID).Err()
|
||||
_ = h.redis.Del(ctx, "ws:user:"+userID).Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) Send(userID string, notification Notification) {
|
||||
h.mu.RLock()
|
||||
conns := h.connections[userID]
|
||||
h.mu.RUnlock()
|
||||
for conn := range conns {
|
||||
if err := conn.WriteJSON(notification); err != nil {
|
||||
log.Printf("Не удалось отправить уведомление пользователю %s: %v", userID, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) Broadcast(notification Notification) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
for userID, conns := range h.connections {
|
||||
for conn := range conns {
|
||||
if err := conn.WriteJSON(notification); err != nil {
|
||||
log.Printf("Не удалось отправить широковещательное уведомление пользователю %s: %v", userID, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func env(key, fallback string) string {
|
||||
if value := os.Getenv(key); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func resolveUserID(r *http.Request) (string, error) {
|
||||
if token := strings.TrimSpace(r.URL.Query().Get("token")); token != "" {
|
||||
secret := env("JWT_ACCESS_SECRET", "docker-access-secret")
|
||||
parsed, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
|
||||
if token.Method != jwt.SigningMethodHS256 {
|
||||
return nil, errors.New("неподдерживаемый алгоритм JWT")
|
||||
}
|
||||
return []byte(secret), nil
|
||||
})
|
||||
if err != nil || !parsed.Valid {
|
||||
return "", err
|
||||
}
|
||||
claims, ok := parsed.Claims.(jwt.MapClaims)
|
||||
if !ok {
|
||||
return "", errors.New("некорректные claims JWT")
|
||||
}
|
||||
sub, _ := claims["sub"].(string)
|
||||
if sub == "" {
|
||||
return "", errors.New("JWT не содержит sub")
|
||||
}
|
||||
return sub, nil
|
||||
}
|
||||
|
||||
userID := strings.TrimSpace(r.URL.Query().Get("user_id"))
|
||||
if userID == "" {
|
||||
return "", errors.New("не передан token или user_id")
|
||||
}
|
||||
return userID, nil
|
||||
}
|
||||
|
||||
func main() {
|
||||
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
redisClient := redis.NewClient(&redis.Options{Addr: env("REDIS_ADDR", "localhost:6379")})
|
||||
if err := redisClient.Ping(ctx).Err(); err != nil {
|
||||
log.Printf("Redis недоступен при старте, сервис продолжит переподключаться: %v", err)
|
||||
}
|
||||
|
||||
hub := NewHub(redisClient)
|
||||
upgrader := websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("ok"))
|
||||
})
|
||||
|
||||
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
|
||||
userID, err := resolveUserID(r)
|
||||
if err != nil || userID == "" {
|
||||
http.Error(w, "Не удалось авторизовать WebSocket", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
log.Printf("Не удалось открыть WebSocket: %v", err)
|
||||
return
|
||||
}
|
||||
hub.Add(r.Context(), userID, conn)
|
||||
defer func() {
|
||||
hub.Remove(context.Background(), userID, conn)
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
if err := conn.SetReadDeadline(time.Now().Add(2 * time.Minute)); err != nil {
|
||||
log.Printf("Не удалось обновить таймаут чтения: %v", err)
|
||||
return
|
||||
}
|
||||
_, _, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_ = redisClient.Set(r.Context(), "ws:user:"+userID, "online", 90*time.Second).Err()
|
||||
}
|
||||
})
|
||||
|
||||
go consumeRabbit(ctx, hub)
|
||||
|
||||
server := &http.Server{
|
||||
Addr: ":" + env("PORT", "8085"),
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
go func() {
|
||||
log.Printf("Media WS запущен на %s", server.Addr)
|
||||
if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
|
||||
log.Fatalf("Ошибка HTTP сервера: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
<-ctx.Done()
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := server.Shutdown(shutdownCtx); err != nil {
|
||||
log.Printf("Не удалось корректно остановить сервер: %v", err)
|
||||
}
|
||||
_ = redisClient.Close()
|
||||
}
|
||||
|
||||
func consumeRabbit(ctx context.Context, hub *Hub) {
|
||||
url := env("RABBITMQ_URL", "amqp://lendry:lendry_password@localhost:5672/")
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
conn, err := amqp091.Dial(url)
|
||||
if err != nil {
|
||||
log.Printf("RabbitMQ недоступен, повтор через 5 секунд: %v", err)
|
||||
time.Sleep(5 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
channel, err := conn.Channel()
|
||||
if err != nil {
|
||||
log.Printf("Не удалось открыть RabbitMQ channel: %v", err)
|
||||
_ = conn.Close()
|
||||
time.Sleep(5 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
queue, err := channel.QueueDeclare("id.notifications", true, false, false, false, nil)
|
||||
if err != nil {
|
||||
log.Printf("Не удалось объявить очередь уведомлений: %v", err)
|
||||
_ = channel.Close()
|
||||
_ = conn.Close()
|
||||
time.Sleep(5 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
deliveries, err := channel.Consume(queue.Name, "media-ws", false, false, false, false, nil)
|
||||
if err != nil {
|
||||
log.Printf("Не удалось подписаться на уведомления: %v", err)
|
||||
_ = channel.Close()
|
||||
_ = conn.Close()
|
||||
time.Sleep(5 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_ = channel.Close()
|
||||
_ = conn.Close()
|
||||
return
|
||||
case delivery, ok := <-deliveries:
|
||||
if !ok {
|
||||
_ = channel.Close()
|
||||
_ = conn.Close()
|
||||
time.Sleep(2 * time.Second)
|
||||
break
|
||||
}
|
||||
handleDelivery(hub, delivery)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func handleDelivery(hub *Hub, delivery amqp091.Delivery) {
|
||||
var notification Notification
|
||||
if err := json.Unmarshal(delivery.Body, ¬ification); err != nil {
|
||||
log.Printf("Некорректное уведомление: %v", err)
|
||||
_ = delivery.Nack(false, false)
|
||||
return
|
||||
}
|
||||
|
||||
if notification.UserID == "" {
|
||||
hub.Broadcast(notification)
|
||||
} else {
|
||||
hub.Send(notification.UserID, notification)
|
||||
}
|
||||
_ = delivery.Ack(false)
|
||||
}
|
||||
Reference in New Issue
Block a user