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

     1  package corev1alpha1
     2  
     3  import (
     4  	"sync"
     5  
     6  	core "github.com/awesome-flow/flow/pkg/corev1alpha1"
     7  )
     8  
     9  type TestActorState string
    10  
    11  const (
    12  	TestActorStateNew     TestActorState = "new"
    13  	TestActorStateStarted                = "started"
    14  	TestActorStateStopped                = "stopped"
    15  )
    16  
    17  type TestActor struct {
    18  	name      string
    19  	peerscnt  int
    20  	ctx       *core.Context
    21  	lock      sync.Mutex
    22  	params    core.Params
    23  	state     TestActorState
    24  	queue     chan *core.Message
    25  	done      chan struct{}
    26  	onstart   []func()
    27  	onstop    []func()
    28  	onconnect []func(int, core.Receiver)
    29  	onreceive []func(*core.Message)
    30  }
    31  
    32  func NewTestActor(name string, ctx *core.Context, params core.Params) (core.Actor, error) {
    33  	return &TestActor{
    34  		name:      name,
    35  		ctx:       ctx,
    36  		params:    params,
    37  		queue:     make(chan *core.Message, 1),
    38  		done:      make(chan struct{}),
    39  		state:     TestActorStateNew,
    40  		onstart:   make([]func(), 0, 1),
    41  		onstop:    make([]func(), 0, 1),
    42  		onconnect: make([]func(int, core.Receiver), 0, 1),
    43  		onreceive: make([]func(*core.Message), 0, 1),
    44  	}, nil
    45  }
    46  
    47  func (t *TestActor) Name() string {
    48  	return t.name
    49  }
    50  
    51  func (t *TestActor) Start() error {
    52  	t.lock.Lock()
    53  	defer t.lock.Unlock()
    54  	t.state = TestActorStateStarted
    55  	for _, h := range t.onstart {
    56  		h()
    57  	}
    58  
    59  	return nil
    60  }
    61  
    62  func (t *TestActor) OnStart(h func()) {
    63  	t.onstart = append(t.onstart, h)
    64  }
    65  
    66  func (t *TestActor) Stop() error {
    67  	t.lock.Lock()
    68  	defer t.lock.Unlock()
    69  	close(t.queue)
    70  	if t.peerscnt > 0 {
    71  		<-t.done
    72  	}
    73  	t.state = TestActorStateStopped
    74  	for _, h := range t.onstop {
    75  		h()
    76  	}
    77  
    78  	return nil
    79  }
    80  
    81  func (t *TestActor) OnStop(h func()) {
    82  	t.onstop = append(t.onstop, h)
    83  }
    84  
    85  func (t *TestActor) Connect(nthreads int, peer core.Receiver) error {
    86  	t.lock.Lock()
    87  	defer t.lock.Unlock()
    88  	t.peerscnt++
    89  	go func() {
    90  		for msg := range t.queue {
    91  			if err := peer.Receive(msg); err != nil {
    92  				t.ctx.Logger().Error(err.Error())
    93  			}
    94  		}
    95  		close(t.done)
    96  	}()
    97  	for _, h := range t.onconnect {
    98  		h(nthreads, peer)
    99  	}
   100  
   101  	return nil
   102  }
   103  
   104  func (t *TestActor) OnConnect(h func(int, core.Receiver)) {
   105  	t.onconnect = append(t.onconnect, h)
   106  }
   107  
   108  func (t *TestActor) Receive(msg *core.Message) error {
   109  	t.queue <- msg
   110  	for _, h := range t.onreceive {
   111  		h(msg)
   112  	}
   113  
   114  	return nil
   115  }
   116  
   117  func (t *TestActor) OnReceive(h func(*core.Message)) {
   118  	t.onreceive = append(t.onreceive, h)
   119  }
   120  
   121  func (t *TestActor) State() TestActorState {
   122  	return t.state
   123  }
   124  
   125  func (t *TestActor) Flush() {
   126  	for len(t.queue) > 0 {
   127  		<-t.queue
   128  	}
   129  }