Feat: room lazy init

internal/op/hub.go

@@ -5,9 +5,11 @@ import (
 	"fmt"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
+	pb "github.com/synctv-org/synctv/proto"
 	"github.com/synctv-org/synctv/utils"
 	"github.com/zijiren233/gencontainer/rwmap"
 )
@@ -19,6 +21,8 @@ type Hub struct {
 	exit      chan struct{}
 	closed    uint32
 	wg        sync.WaitGroup
+
+	once utils.Once
 }
 
 type broadcastMessage struct {
@@ -56,22 +60,15 @@ func newHub(id uint) *Hub {
 	}
 }
 
-func (h *Hub) Closed() bool {
-	return atomic.LoadUint32(&h.closed) == 1
-}
-
-var (
-	ErrAlreadyClosed = fmt.Errorf("already closed")
-)
-
-func (h *Hub) Start() {
-	go h.Serve()
+func (h *Hub) Start() error {
+	h.once.Do(func() {
+		go h.serve()
+		go h.ping()
+	})
+	return nil
 }
 
-func (h *Hub) Serve() error {
-	if h.Closed() {
-		return ErrAlreadyClosed
-	}
+func (h *Hub) serve() error {
 	for {
 		select {
 		case message := <-h.broadcast:
@@ -98,6 +95,35 @@ func (h *Hub) Serve() error {
 	}
 }
 
+func (h *Hub) ping() {
+	ticker := time.NewTicker(time.Second * 5)
+	defer ticker.Stop()
+	var pre int64 = 0
+	for {
+		select {
+		case <-ticker.C:
+			current := h.ClientNum()
+			if current != pre {
+				if err := h.Broadcast(&ElementMessage{
+					ElementMessage: &pb.ElementMessage{
+						Type:      pb.ElementMessageType_CHANGE_PEOPLE,
+						PeopleNum: current,
+					},
+				}); err != nil {
+					continue
+				}
+				pre = current
+			} else {
+				if err := h.Broadcast(&PingMessage{}); err != nil {
+					continue
+				}
+			}
+		case <-h.exit:
+			return
+		}
+	}
+}
+
 func (h *Hub) devMessage(msg Message) {
 	switch msg.MessageType() {
 	case websocket.TextMessage:
@@ -105,12 +131,21 @@ func (h *Hub) devMessage(msg Message) {
 	}
 }
 
+func (h *Hub) Closed() bool {
+	return atomic.LoadUint32(&h.closed) == 1
+}
+
+var (
+	ErrAlreadyClosed = fmt.Errorf("already closed")
+)
+
 func (h *Hub) Close() error {
 	if !atomic.CompareAndSwapUint32(&h.closed, 0, 1) {
 		return ErrAlreadyClosed
 	}
 	close(h.exit)
 	h.clients.Range(func(_ uint, client *Client) bool {
+		h.clients.Delete(client.u.ID)
 		client.Close()
 		return true
 	})
@@ -125,6 +160,7 @@ func (h *Hub) Broadcast(data Message, conf ...BroadcastConf) error {
 	if h.Closed() {
 		return ErrAlreadyClosed
 	}
+	h.once.Done()
 	msg := &broadcastMessage{data: data}
 	for _, c := range conf {
 		c(msg)
@@ -141,6 +177,10 @@ func (h *Hub) RegClient(cli *Client) (*Client, error) {
 	if h.Closed() {
 		return nil, ErrAlreadyClosed
 	}
+	err := h.Start()
+	if err != nil {
+		return nil, err
+	}
 	c, loaded := h.clients.LoadOrStore(cli.u.ID, cli)
 	if loaded {
 		return nil, errors.New("client already registered")

internal/op/room.go

@@ -27,17 +27,17 @@ import (
 type Room struct {
 	model.Room
 	version    uint32
-	current    current
+	current    *current
 	rtmpa      *rtmps.App
 	initOnce   utils.Once
 	lastActive int64
 	hub        *Hub
 }
 
-func (r *Room) Init() {
+func (r *Room) lazyInit() {
 	r.initOnce.Do(func() {
 		atomic.CompareAndSwapUint32(&r.version, 0, 1)
-		r.current = *newCurrent()
+		r.current = newCurrent()
 		r.hub = newHub(r.ID)
 		a, err := rtmp.RtmpServer().NewApp(r.Name)
 		if err != nil {
@@ -48,10 +48,12 @@ func (r *Room) Init() {
 }
 
 func (r *Room) Hub() *Hub {
+	r.lazyInit()
 	return r.hub
 }
 
 func (r *Room) App() *rtmps.App {
+	r.lazyInit()
 	return r.rtmpa
 }
 
@@ -77,6 +79,7 @@ func (r *Room) UpdateMovie(movieId uint, movie model.BaseMovieInfo) error {
 	}
 	switch {
 	case (m.RtmpSource && !movie.RtmpSource) || (m.Live && m.Proxy && !movie.Proxy):
+		r.lazyInit()
 		r.rtmpa.DelChannel(m.PullKey)
 		m.PullKey = ""
 	case m.Proxy && !movie.Proxy:
@@ -96,6 +99,7 @@ func (r *Room) InitMovie(movie *model.Movie) error {
 			return errors.New("hls player is not enabled")
 		}
 		movie.PullKey = uuid.New().String()
+		r.lazyInit()
 		_, err := r.rtmpa.NewChannel(movie.PullKey)
 		if err != nil {
 			return err
@@ -111,6 +115,7 @@ func (r *Room) InitMovie(movie *model.Movie) error {
 		switch u.Scheme {
 		case "rtmp":
 			PullKey := uuid.New().String()
+			r.lazyInit()
 			c, err := r.rtmpa.NewChannel(PullKey)
 			if err != nil {
 				return err
@@ -135,6 +140,7 @@ func (r *Room) InitMovie(movie *model.Movie) error {
 			}()
 		case "http", "https":
 			PullKey := uuid.New().String()
+			r.lazyInit()
 			c, err := r.rtmpa.NewChannel(PullKey)
 			if err != nil {
 				return err
@@ -267,6 +273,7 @@ func (r *Room) DeleteMovieByID(id uint) error {
 		return err
 	}
 	if m.PullKey != "" {
+		r.lazyInit()
 		r.rtmpa.DelChannel(m.PullKey)
 	}
 	return nil
@@ -277,6 +284,7 @@ func (r *Room) ClearMovies() error {
 	if err != nil {
 		return err
 	}
+	r.lazyInit()
 	for _, m := range ms {
 		if m.PullKey != "" {
 			r.rtmpa.DelChannel(m.PullKey)
@@ -286,6 +294,7 @@ func (r *Room) ClearMovies() error {
 }
 
 func (r *Room) Current() *Current {
+	r.lazyInit()
 	c := r.current.Current()
 	return &c
 }
@@ -295,6 +304,7 @@ func (r *Room) ChangeCurrentMovie(id uint) error {
 	if err != nil {
 		return err
 	}
+	r.lazyInit()
 	r.current.SetMovie(*m)
 	return nil
 }
@@ -308,17 +318,21 @@ func (r *Room) GetMovieWithPullKey(pullKey string) (*model.Movie, error) {
 }
 
 func (r *Room) RegClient(user *User, conn *websocket.Conn) (*Client, error) {
+	r.lazyInit()
 	return r.hub.RegClient(newClient(user, r, conn))
 }
 
 func (r *Room) UnregisterClient(user *User) error {
+	r.lazyInit()
 	return r.hub.UnRegClient(user)
 }
 
 func (r *Room) SetStatus(playing bool, seek float64, rate float64, timeDiff float64) Status {
+	r.lazyInit()
 	return r.current.SetStatus(playing, seek, rate, timeDiff)
 }
 
 func (r *Room) SetSeekRate(seek float64, rate float64, timeDiff float64) Status {
+	r.lazyInit()
 	return r.current.SetSeekRate(seek, rate, timeDiff)
 }

server/handlers/room.go

@@ -5,16 +5,12 @@ import (
 	"fmt"
 	"net/http"
 	"strconv"
-	"time"
-
-	log "github.com/sirupsen/logrus"
 
 	"github.com/gin-gonic/gin"
 	"github.com/maruel/natural"
 	"github.com/synctv-org/synctv/internal/db"
 	dbModel "github.com/synctv-org/synctv/internal/model"
 	"github.com/synctv-org/synctv/internal/op"
-	pb "github.com/synctv-org/synctv/proto"
 	"github.com/synctv-org/synctv/server/middlewares"
 	"github.com/synctv-org/synctv/server/model"
 	"github.com/zijiren233/gencontainer/vec"
@@ -46,51 +42,18 @@ func CreateRoom(ctx *gin.Context) {
 		return
 	}
 
-	room, err := op.LoadRoom(r)
+	_, err = op.LoadRoom(r)
 	if err != nil {
 		ctx.AbortWithStatusJSON(http.StatusInternalServerError, model.NewApiErrorResp(err))
 		return
 	}
 
-	room.Init()
-	room.Hub().Start()
-
 	token, err := middlewares.NewAuthUserToken(user)
 	if err != nil {
 		ctx.AbortWithStatusJSON(http.StatusInternalServerError, model.NewApiErrorResp(err))
 		return
 	}
 
-	go func() {
-		ticker := time.NewTicker(time.Second * 5)
-		defer ticker.Stop()
-		var pre int64 = 0
-		for range ticker.C {
-			if room.Hub().Closed() {
-				log.Debugf("ws: room %s closed, stop broadcast people num", room.Name)
-				return
-			}
-			current := room.Hub().ClientNum()
-			if current != pre {
-				if err := room.Hub().Broadcast(&op.ElementMessage{
-					ElementMessage: &pb.ElementMessage{
-						Type:      pb.ElementMessageType_CHANGE_PEOPLE,
-						PeopleNum: current,
-					},
-				}); err != nil {
-					log.Errorf("ws: room %s broadcast people num error: %v", room.Name, err)
-					continue
-				}
-				pre = current
-			} else {
-				if err := room.Hub().Broadcast(&op.PingMessage{}); err != nil {
-					log.Errorf("ws: room %s broadcast ping error: %v", room.Name, err)
-					continue
-				}
-			}
-		}
-	}()
-
 	ctx.JSON(http.StatusCreated, model.NewApiDataResp(gin.H{
 		"token": token,
 	}))

utils/utils.go

@@ -142,13 +142,24 @@ type Once struct {
 }
 
 func (o *Once) Done() (doned bool) {
+	done := atomic.LoadUint32(&o.done)
+	if done == 1 {
+		return true
+	} else if done == 2 {
+		return false
+	}
+
 	o.m.Lock()
 	defer o.m.Unlock()
 	if o.done == 0 {
-		o.done = 1
-		return false
-	}
-	return true
+		doned = false
+		atomic.StoreUint32(&o.done, 2)
+	} else if o.done == 1 {
+		doned = true
+	} else {
+		doned = false
+	}
+	return
 }
 
 func (o *Once) Do(f func()) {