package protocol import ( "sort" "sync" "github.com/status-im/status-go/eth-node/types" "github.com/status-im/status-go/protocol/transport" ) type messagesOrderType int const ( messagesOrderRandom messagesOrderType = iota messagesOrderAsPosted messagesOrderReversed ) type MessagesOrderController struct { order messagesOrderType messagesInPostOrder [][]byte mutex sync.RWMutex quit chan struct{} quitOnce sync.Once } func NewMessagesOrderController(order messagesOrderType) *MessagesOrderController { return &MessagesOrderController{ order: order, messagesInPostOrder: [][]byte{}, mutex: sync.RWMutex{}, quit: make(chan struct{}), } } func (m *MessagesOrderController) Start(c chan *PostMessageSubscription) { go func() { for { select { case sub, more := <-c: if !more { return } m.mutex.Lock() m.messagesInPostOrder = append(m.messagesInPostOrder, sub.id) m.mutex.Unlock() case <-m.quit: return } } }() } func (m *MessagesOrderController) Stop() { m.quitOnce.Do(func() { close(m.quit) }) } func (m *MessagesOrderController) newMessagesIterator(chatWithMessages map[transport.Filter][]*types.Message) MessagesIterator { switch m.order { case messagesOrderAsPosted, messagesOrderReversed: return &messagesIterator{chatWithMessages: m.sort(chatWithMessages, m.order)} } return NewDefaultMessagesIterator(chatWithMessages) } func buildIndexMap(messages [][]byte) map[string]int { indexMap := make(map[string]int) for i, hash := range messages { hashStr := string(hash) indexMap[hashStr] = i } return indexMap } func (m *MessagesOrderController) sort(chatWithMessages map[transport.Filter][]*types.Message, order messagesOrderType) []*chatWithMessage { allMessages := make([]*chatWithMessage, 0) for chat, messages := range chatWithMessages { for _, message := range messages { allMessages = append(allMessages, &chatWithMessage{chat: chat, message: message}) } } m.mutex.RLock() indexMap := buildIndexMap(m.messagesInPostOrder) m.mutex.RUnlock() sort.SliceStable(allMessages, func(i, j int) bool { indexI, okI := indexMap[string(allMessages[i].message.Hash)] indexJ, okJ := indexMap[string(allMessages[j].message.Hash)] if okI && okJ { if order == messagesOrderReversed { return indexI > indexJ } return indexI < indexJ } return !okI && okJ // keep messages with unknown hashes at the end }) return allMessages } type chatWithMessage struct { chat transport.Filter message *types.Message } type messagesIterator struct { chatWithMessages []*chatWithMessage currentIndex int } func (it *messagesIterator) HasNext() bool { return it.currentIndex < len(it.chatWithMessages) } func (it *messagesIterator) Next() (transport.Filter, []*types.Message) { if it.HasNext() { m := it.chatWithMessages[it.currentIndex] it.currentIndex++ return m.chat, []*types.Message{m.message} } return transport.Filter{}, nil }