github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/remote/endpoint_writer_mailbox.go (about)

     1  package remote
     2  
     3  import (
     4  	"runtime"
     5  	"sync/atomic"
     6  
     7  	"github.com/asynkron/protoactor-go/actor"
     8  
     9  	"github.com/asynkron/protoactor-go/internal/queue/goring"
    10  	"github.com/asynkron/protoactor-go/internal/queue/mpsc"
    11  )
    12  
    13  const (
    14  	mailboxIdle    int32 = iota
    15  	mailboxRunning int32 = iota
    16  )
    17  
    18  const (
    19  	mailboxHasNoMessages   int32 = iota
    20  	mailboxHasMoreMessages int32 = iota
    21  )
    22  
    23  type endpointWriterMailbox struct {
    24  	userMailbox     *goring.Queue
    25  	systemMailbox   *mpsc.Queue
    26  	schedulerStatus int32
    27  	hasMoreMessages int32
    28  	invoker         actor.MessageInvoker
    29  	batchSize       int
    30  	dispatcher      actor.Dispatcher
    31  	suspended       bool
    32  }
    33  
    34  func (m *endpointWriterMailbox) PostUserMessage(message interface{}) {
    35  	// batching mailbox only use the message part
    36  	m.userMailbox.Push(message)
    37  	m.schedule()
    38  }
    39  
    40  func (m *endpointWriterMailbox) PostSystemMessage(message interface{}) {
    41  	m.systemMailbox.Push(message)
    42  	m.schedule()
    43  }
    44  
    45  func (m *endpointWriterMailbox) RegisterHandlers(invoker actor.MessageInvoker, dispatcher actor.Dispatcher) {
    46  	m.invoker = invoker
    47  	m.dispatcher = dispatcher
    48  }
    49  
    50  func (m *endpointWriterMailbox) Start() {
    51  }
    52  
    53  func (m *endpointWriterMailbox) schedule() {
    54  	atomic.StoreInt32(&m.hasMoreMessages, mailboxHasMoreMessages) // we have more messages to process
    55  	if atomic.CompareAndSwapInt32(&m.schedulerStatus, mailboxIdle, mailboxRunning) {
    56  		m.dispatcher.Schedule(m.processMessages)
    57  	}
    58  }
    59  
    60  func (m *endpointWriterMailbox) processMessages() {
    61  	// we are about to start processing messages, we can safely reset the message flag of the mailbox
    62  	atomic.StoreInt32(&m.hasMoreMessages, mailboxHasNoMessages)
    63  process:
    64  	m.run()
    65  
    66  	// set mailbox to idle
    67  	atomic.StoreInt32(&m.schedulerStatus, mailboxIdle)
    68  
    69  	// check if there are still messages to process (sent after the message loop ended)
    70  	if atomic.SwapInt32(&m.hasMoreMessages, mailboxHasNoMessages) == mailboxHasMoreMessages {
    71  		// try setting the mailbox back to running
    72  		if atomic.CompareAndSwapInt32(&m.schedulerStatus, mailboxIdle, mailboxRunning) {
    73  			goto process
    74  		}
    75  	}
    76  }
    77  
    78  func (m *endpointWriterMailbox) run() {
    79  	var msg interface{}
    80  	defer func() {
    81  		if r := recover(); r != nil {
    82  			m.invoker.EscalateFailure(r, msg)
    83  		}
    84  	}()
    85  
    86  	for {
    87  		// keep processing system messages until queue is empty
    88  		if msg = m.systemMailbox.Pop(); msg != nil {
    89  			switch msg.(type) {
    90  			case *actor.SuspendMailbox:
    91  				m.suspended = true
    92  			case *actor.ResumeMailbox:
    93  				m.suspended = false
    94  			default:
    95  				m.invoker.InvokeSystemMessage(msg)
    96  			}
    97  
    98  			continue
    99  		}
   100  
   101  		// didn't process a system message, so break until we are resumed
   102  		if m.suspended {
   103  			return
   104  		}
   105  
   106  		var ok bool
   107  		if msg, ok = m.userMailbox.PopMany(int64(m.batchSize)); ok {
   108  			m.invoker.InvokeUserMessage(msg)
   109  		} else {
   110  			return
   111  		}
   112  
   113  		runtime.Gosched()
   114  	}
   115  }
   116  
   117  func (m *endpointWriterMailbox) UserMessageCount() int {
   118  	return int(m.userMailbox.Length())
   119  }
   120  
   121  func endpointWriterMailboxProducer(batchSize, initialSize int) actor.MailboxProducer {
   122  	return func() actor.Mailbox {
   123  		userMailbox := goring.New(int64(initialSize))
   124  		systemMailbox := mpsc.New()
   125  		return &endpointWriterMailbox{
   126  			userMailbox:     userMailbox,
   127  			systemMailbox:   systemMailbox,
   128  			hasMoreMessages: mailboxHasNoMessages,
   129  			schedulerStatus: mailboxIdle,
   130  			batchSize:       batchSize,
   131  		}
   132  	}
   133  }