github.com/whiteboxio/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/corev1alpha1/actor/buffer.go (about)

     1  package actor
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	core "github.com/awesome-flow/flow/pkg/corev1alpha1"
     8  )
     9  
    10  const (
    11  	DefaultBufCapacity    = 65536
    12  	DefaultBufMaxAttempts = 16
    13  )
    14  
    15  type MsgCnt struct {
    16  	msg *core.Message
    17  	cnt uint32
    18  }
    19  
    20  func NewMsgCnt(msg *core.Message) *MsgCnt {
    21  	return &MsgCnt{msg: msg}
    22  }
    23  
    24  type Buffer struct {
    25  	name  string
    26  	ctx   *core.Context
    27  	queue chan *MsgCnt
    28  	wg    sync.WaitGroup
    29  }
    30  
    31  var _ core.Actor = (*Buffer)(nil)
    32  
    33  func NewBuffer(name string, ctx *core.Context, params core.Params) (core.Actor, error) {
    34  	return &Buffer{
    35  		name:  name,
    36  		ctx:   ctx,
    37  		queue: make(chan *MsgCnt, DefaultBufCapacity),
    38  	}, nil
    39  }
    40  
    41  func (b *Buffer) Name() string {
    42  	return b.name
    43  }
    44  
    45  func (b *Buffer) Start() error {
    46  	return nil
    47  }
    48  
    49  func (b *Buffer) Stop() error {
    50  	close(b.queue)
    51  	b.wg.Wait()
    52  
    53  	return nil
    54  }
    55  
    56  func (b *Buffer) Connect(nthreads int, peer core.Receiver) error {
    57  	for i := 0; i < nthreads; i++ {
    58  		b.wg.Add(1)
    59  		go func() {
    60  			var sts core.MsgStatus
    61  			for msgcnt := range b.queue {
    62  				msgcp := msgcnt.msg.Copy()
    63  				err := peer.Receive(msgcp)
    64  				sts = 0
    65  				if err == nil {
    66  					sts = msgcp.Await()
    67  					switch sts {
    68  					case core.MsgStatusDone, core.MsgStatusPartialSend:
    69  					default:
    70  						err = fmt.Errorf("failed to send message: code(%d)", sts)
    71  					}
    72  				}
    73  				if err != nil {
    74  					msgcnt.cnt++
    75  					if msgcnt.cnt < DefaultBufMaxAttempts {
    76  						b.queue <- msgcnt
    77  						continue
    78  					}
    79  					sts = core.MsgStatusFailed
    80  				}
    81  				msgcnt.msg.Complete(sts)
    82  			}
    83  			b.wg.Done()
    84  		}()
    85  	}
    86  
    87  	return nil
    88  }
    89  
    90  func (b *Buffer) Receive(msg *core.Message) error {
    91  	b.queue <- NewMsgCnt(msg)
    92  
    93  	return nil
    94  }