improve websocket service

routers/web/websocket/websocket.go

@@ -4,13 +4,12 @@
 package websocket
 
 import (
-	"time"
+	"github.com/olahol/melody"
 
 	"code.gitea.io/gitea/modules/context"
 	"code.gitea.io/gitea/modules/web"
+	notify_service "code.gitea.io/gitea/services/notify"
 	"code.gitea.io/gitea/services/websocket"
-
-	"github.com/olahol/melody"
 )
 
 var m *melody.Melody
@@ -21,17 +20,7 @@ func Init(r *web.Route) {
 	m.HandleConnect(websocket.HandleConnect)
 	m.HandleMessage(websocket.HandleMessage)
 	m.HandleDisconnect(websocket.HandleDisconnect)
-
+	notify_service.RegisterNotifier(websocket.NewNotifier(m))
-	go func() {
-		for {
-			// TODO: send proper updated html
-			err := m.Broadcast([]byte("<div hx-swap-oob=\"beforebegin:.timeline-item.comment.form\"><div class=\"hello\">hello world!</div></div>"))
-			if err != nil {
-				break
-			}
-			time.Sleep(5 * time.Second)
-		}
-	}()
 }
 
 func webSocket(ctx *context.Context) {

services/websocket/notifier.go

@@ -10,19 +10,43 @@ import (
 	repo_model "code.gitea.io/gitea/models/repo"
 	user_model "code.gitea.io/gitea/models/user"
 	notify_service "code.gitea.io/gitea/services/notify"
+	"github.com/olahol/melody"
 )
 
 type webhookNotifier struct {
 	notify_service.NullNotifier
+	m *melody.Melody
 }
 
 var _ notify_service.Notifier = &webhookNotifier{}
 
 // NewNotifier create a new webhooksNotifier notifier
-func NewNotifier() notify_service.Notifier {
+func NewNotifier(m *melody.Melody) notify_service.Notifier {
-	return &webhookNotifier{}
+	return &webhookNotifier{
+		m: m,
+	}
 }
 
 func (n *webhookNotifier) CreateIssueComment(ctx context.Context, doer *user_model.User, repo *repo_model.Repository, issue *issues_model.Issue, comment *issues_model.Comment, mentions []*user_model.User) {
-	// TODO
+	// TODO: use proper message
+	msg := []byte("<div hx-swap-oob=\"beforebegin:.timeline-item.comment.form\"><div class=\"hello\">hello world!</div></div>")
+
+	n.m.BroadcastFilter(msg, func(s *melody.Session) bool {
+		sessionData, err := getSessionData(s)
+		if err != nil {
+			return false
+		}
+
+		if sessionData.uid == doer.ID {
+			return true
+		}
+
+		for _, mention := range mentions {
+			if mention.ID == sessionData.uid {
+				return true
+			}
+		}
+
+		return false
+	})
 }

services/websocket/session.go

@@ -0,0 +1,26 @@
+package websocket
+
+import (
+	"fmt"
+
+	"github.com/olahol/melody"
+)
+
+type sessionData struct {
+	uid        int64
+	unregister func()
+}
+
+func getSessionData(s *melody.Session) (*sessionData, error) {
+	_data, ok := s.Get("data")
+	if !ok {
+		return nil, fmt.Errorf("no session data")
+	}
+
+	data, ok := _data.(*sessionData)
+	if !ok {
+		return nil, fmt.Errorf("invalid session data")
+	}
+
+	return data, nil
+}

services/websocket/websocket.go

@@ -4,7 +4,6 @@
 package websocket
 
 import (
-	"fmt"
 	"time"
 
 	"code.gitea.io/gitea/modules/context"
@@ -15,10 +14,6 @@ import (
 	"github.com/olahol/melody"
 )
 
-type SessionData struct {
-	unregister func()
-}
-
 func HandleConnect(s *melody.Session) {
 	ctx := context.GetWebContext(s.Request)
 
@@ -43,7 +38,7 @@ func HandleConnect(s *melody.Session) {
 
 	messageChan := eventsource.GetManager().Register(uid)
 
-	sessionData := &SessionData{
+	sessionData := &sessionData{
 		unregister: func() {
 			eventsource.GetManager().Unregister(uid, messageChan)
 			// ensure the messageChan is closed
@@ -109,20 +104,6 @@ func HandleMessage(s *melody.Session, msg []byte) {
 	// TODO: Handle incoming messages
 }
 
-func getSessionData(s *melody.Session) (*SessionData, error) {
-	_data, ok := s.Get("data")
-	if !ok {
-		return nil, fmt.Errorf("no session data")
-	}
-
-	data, ok := _data.(*SessionData)
-	if !ok {
-		return nil, fmt.Errorf("invalid session data")
-	}
-
-	return data, nil
-}
-
 func HandleDisconnect(s *melody.Session) {
 	data, err := getSessionData(s)
 	if err != nil {