github.com/status-im/status-go@v1.1.0/protocol/messenger_messages_order_controller_test.go (about)

     1  package protocol
     2  
     3  import (
     4  	"sort"
     5  	"sync"
     6  
     7  	"github.com/status-im/status-go/eth-node/types"
     8  	"github.com/status-im/status-go/protocol/transport"
     9  )
    10  
    11  type messagesOrderType int
    12  
    13  const (
    14  	messagesOrderRandom messagesOrderType = iota
    15  	messagesOrderAsPosted
    16  	messagesOrderReversed
    17  )
    18  
    19  type MessagesOrderController struct {
    20  	order               messagesOrderType
    21  	messagesInPostOrder [][]byte
    22  	mutex               sync.RWMutex
    23  	quit                chan struct{}
    24  	quitOnce            sync.Once
    25  }
    26  
    27  func NewMessagesOrderController(order messagesOrderType) *MessagesOrderController {
    28  	return &MessagesOrderController{
    29  		order:               order,
    30  		messagesInPostOrder: [][]byte{},
    31  		mutex:               sync.RWMutex{},
    32  		quit:                make(chan struct{}),
    33  	}
    34  }
    35  
    36  func (m *MessagesOrderController) Start(c chan *PostMessageSubscription) {
    37  	go func() {
    38  		for {
    39  			select {
    40  			case sub, more := <-c:
    41  				if !more {
    42  					return
    43  				}
    44  				m.mutex.Lock()
    45  				m.messagesInPostOrder = append(m.messagesInPostOrder, sub.id)
    46  				m.mutex.Unlock()
    47  
    48  			case <-m.quit:
    49  				return
    50  			}
    51  		}
    52  	}()
    53  }
    54  
    55  func (m *MessagesOrderController) Stop() {
    56  	m.quitOnce.Do(func() {
    57  		close(m.quit)
    58  	})
    59  }
    60  
    61  func (m *MessagesOrderController) newMessagesIterator(chatWithMessages map[transport.Filter][]*types.Message) MessagesIterator {
    62  	switch m.order {
    63  	case messagesOrderAsPosted, messagesOrderReversed:
    64  		return &messagesIterator{chatWithMessages: m.sort(chatWithMessages, m.order)}
    65  	}
    66  
    67  	return NewDefaultMessagesIterator(chatWithMessages)
    68  }
    69  
    70  func buildIndexMap(messages [][]byte) map[string]int {
    71  	indexMap := make(map[string]int)
    72  	for i, hash := range messages {
    73  		hashStr := string(hash)
    74  		indexMap[hashStr] = i
    75  	}
    76  	return indexMap
    77  }
    78  
    79  func (m *MessagesOrderController) sort(chatWithMessages map[transport.Filter][]*types.Message, order messagesOrderType) []*chatWithMessage {
    80  	allMessages := make([]*chatWithMessage, 0)
    81  	for chat, messages := range chatWithMessages {
    82  		for _, message := range messages {
    83  			allMessages = append(allMessages, &chatWithMessage{chat: chat, message: message})
    84  		}
    85  	}
    86  
    87  	m.mutex.RLock()
    88  	indexMap := buildIndexMap(m.messagesInPostOrder)
    89  	m.mutex.RUnlock()
    90  
    91  	sort.SliceStable(allMessages, func(i, j int) bool {
    92  		indexI, okI := indexMap[string(allMessages[i].message.Hash)]
    93  		indexJ, okJ := indexMap[string(allMessages[j].message.Hash)]
    94  
    95  		if okI && okJ {
    96  			if order == messagesOrderReversed {
    97  				return indexI > indexJ
    98  			}
    99  			return indexI < indexJ
   100  		}
   101  
   102  		return !okI && okJ // keep messages with unknown hashes at the end
   103  	})
   104  
   105  	return allMessages
   106  }
   107  
   108  type chatWithMessage struct {
   109  	chat    transport.Filter
   110  	message *types.Message
   111  }
   112  
   113  type messagesIterator struct {
   114  	chatWithMessages []*chatWithMessage
   115  	currentIndex     int
   116  }
   117  
   118  func (it *messagesIterator) HasNext() bool {
   119  	return it.currentIndex < len(it.chatWithMessages)
   120  }
   121  
   122  func (it *messagesIterator) Next() (transport.Filter, []*types.Message) {
   123  	if it.HasNext() {
   124  		m := it.chatWithMessages[it.currentIndex]
   125  		it.currentIndex++
   126  		return m.chat, []*types.Message{m.message}
   127  	}
   128  
   129  	return transport.Filter{}, nil
   130  }