team/server/channels/app/platform/websocket_reliable.go

// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
// See LICENSE.txt for license information.

package platform

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/hanzoteam/server/server/public/model"
)

func (ps *PlatformService) GetWSQueues(userID, connectionID string, seqNum int64) (*model.WSQueues, error) {
	hub := ps.GetHubForUserId(userID)
	if hub == nil {
		return nil, nil
	}
	connRes := hub.CheckConn(userID, connectionID)
	if connRes == nil {
		return nil, nil
	}
	aq := connRes.ActiveQueue
	dq := connRes.DeadQueue
	dqPtr := connRes.DeadQueuePointer

	// Nothing was written on this server. Early return.
	if dq[0] == nil {
		return nil, nil
	}

	// Check if seq_num-1 == last value in the dead queue.
	if perfectMatch := !_hasMsgLoss(dq, dqPtr, seqNum); perfectMatch {
		close(aq)
		aqSlice, err := ps.marshalAQ(aq, connectionID, userID)
		if err != nil {
			return nil, fmt.Errorf("failed to get from active queue: %w", err)
		}
		// send only aq
		return &model.WSQueues{
			ActiveQ:    aqSlice,
			ReuseCount: connRes.ReuseCount,
		}, nil
	}

	// Check if seq_num is somewhere else in the dead queue.
	if ok, index := _isInDeadQueue(dq, seqNum); ok {
		close(aq)
		aqSlice, err := ps.marshalAQ(aq, connectionID, userID)
		if err != nil {
			return nil, fmt.Errorf("failed to get from active queue: %w", err)
		}
		dqSlice, err := ps.marshalDQ(dq, index, dqPtr)
		if err != nil {
			return nil, fmt.Errorf("failed to get from dead queue: %w", err)
		}
		// send aq + drainedDq.
		return &model.WSQueues{
			ActiveQ:    aqSlice,
			DeadQ:      dqSlice,
			ReuseCount: connRes.ReuseCount,
		}, nil
	}

	// Nothing matched.
	return nil, nil
}

func (ps *PlatformService) marshalAQ(aq <-chan model.WebSocketMessage, connID, userID string) ([]model.ActiveQueueItem, error) {
	aqSlice := make([]model.ActiveQueueItem, 0)
	for msg := range aq {
		evtType := model.WebSocketMsgTypeResponse
		_, evtOk := msg.(*model.WebSocketEvent)
		if evtOk {
			evtType = model.WebSocketMsgTypeEvent
		}
		buf, err := msg.ToJSON()
		if err != nil {
			return nil, fmt.Errorf("failed to marshal websocket event: %w, connection_id=%s, user_id=%s", err, connID, userID)
		}
		aqSlice = append(aqSlice, model.ActiveQueueItem{
			Buf:  json.RawMessage(buf),
			Type: evtType,
		})
	}

	return aqSlice, nil
}

func (ps *PlatformService) UnmarshalAQItem(aqItem model.ActiveQueueItem) (model.WebSocketMessage, error) {
	var item model.WebSocketMessage
	var err error
	if aqItem.Type == model.WebSocketMsgTypeEvent {
		item, err = model.WebSocketEventFromJSON(bytes.NewReader(aqItem.Buf))
	} else if aqItem.Type == model.WebSocketMsgTypeResponse {
		item, err = model.WebSocketResponseFromJSON(bytes.NewReader(aqItem.Buf))
	} else {
		return nil, fmt.Errorf("unknown websocket message type: %q", aqItem.Type)
	}
	return item, err
}

// marshalDQ is the same as drainDeadQueue, except it writes to a byte slice
// instead of the network. To be refactored into a single method.
func (ps *PlatformService) marshalDQ(dq []*model.WebSocketEvent, index, dqPtr int) ([]json.RawMessage, error) {
	if len(dq) == 0 || dq[0] == nil {
		return nil, nil
	}

	dqSlice := make([]json.RawMessage, 0)
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	// This means pointer hasn't rolled over.
	if dq[dqPtr] == nil {
		// Clear till the end of queue.
		for i := index; i < dqPtr; i++ {
			buf.Reset()
			err := dq[i].Encode(enc, &buf)
			if err != nil {
				return nil, fmt.Errorf("error in encoding websocket message in dead queue: %w", err)
			}
			dqSlice = append(dqSlice, bytes.Clone(buf.Bytes()))
		}
		return dqSlice, nil
	}

	// We go on until next sequence number is smaller than previous one.
	// Which means it has rolled over.
	currPtr := index
	for {
		buf.Reset()
		err := dq[currPtr].Encode(enc, &buf)
		if err != nil {
			return nil, fmt.Errorf("error in encoding websocket message in dead queue: %w", err)
		}
		dqSlice = append(dqSlice, bytes.Clone(buf.Bytes()))
		oldSeq := dq[currPtr].GetSequence()
		currPtr = (currPtr + 1) % deadQueueSize
		newSeq := dq[currPtr].GetSequence()
		if oldSeq > newSeq {
			break
		}
	}
	return dqSlice, nil
}

func (ps *PlatformService) UnmarshalDQ(buf []json.RawMessage) ([]*model.WebSocketEvent, int, error) {
	dqPtr := 0
	dq := make([]*model.WebSocketEvent, deadQueueSize)
	for _, dqItem := range buf {
		item, err := model.WebSocketEventFromJSON(bytes.NewReader(dqItem))
		if err != nil {
			return nil, 0, err
		}

		// Same as active queue, this can never be out of bounds because all dead queues
		// are of deadQueueSize.
		dq[dqPtr] = item
		dqPtr = (dqPtr + 1) % deadQueueSize
	}
	return dq, dqPtr, nil
}