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

     1  package corev1alpha1
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  )
     7  
     8  type MsgStatus uint8
     9  
    10  const (
    11  	// MsgStatusNew represents a new message.
    12  	MsgStatusNew MsgStatus = iota
    13  	// MsgStatusDone represents a message that has left the pipeline.
    14  	MsgStatusDone
    15  	// MsgStatusPartialSend represents a partially sent message: most probably,
    16  	// some branches of the pipeline succeeded to send it and some failed.
    17  	MsgStatusPartialSend
    18  	// MsgStatusInvalid represents a message recognised as invalid by the
    19  	// pipeline components and therefore it's impossible to proceed forward.
    20  	MsgStatusInvalid
    21  	// MsgStatusFailed represents a message for which submission has failed.
    22  	MsgStatusFailed
    23  	// MsgStatusTimedOut represents a message for which one or more components
    24  	// triggered a timeout watermark.
    25  	MsgStatusTimedOut
    26  	// MsgStatusUnroutable represents a message for which the submission
    27  	// destination/branch is unknown. Most likely, a branch with the
    28  	// corresponding name does not exist.
    29  	MsgStatusUnroutable
    30  	// MsgStatusThrottled represents a message which submission process was
    31  	// cancelled due to a quota exhausting.
    32  	MsgStatusThrottled
    33  )
    34  
    35  var (
    36  	MsgCompletedBeforeErr = fmt.Errorf("message has been completed before")
    37  )
    38  
    39  type Message struct {
    40  	body   []byte
    41  	done   chan struct{}
    42  	meta   map[interface{}]interface{}
    43  	status MsgStatus
    44  	mutex  sync.Mutex
    45  }
    46  
    47  func NewMessage(body []byte) *Message {
    48  	cpbody := make([]byte, len(body))
    49  	copy(cpbody, body)
    50  	return &Message{
    51  		body:   cpbody,
    52  		done:   make(chan struct{}),
    53  		meta:   make(map[interface{}]interface{}),
    54  		status: MsgStatusNew,
    55  	}
    56  }
    57  
    58  func (msg *Message) Await() MsgStatus {
    59  	<-msg.done
    60  	return msg.status
    61  }
    62  
    63  func (msg *Message) AwaitChan() <-chan MsgStatus {
    64  	res := make(chan MsgStatus)
    65  	go func() {
    66  		<-msg.done
    67  		res <- msg.status
    68  		close(res)
    69  	}()
    70  	return res
    71  }
    72  
    73  func (msg *Message) Complete(status MsgStatus) error {
    74  	msg.mutex.Lock()
    75  	defer msg.mutex.Unlock()
    76  	if msg.status != MsgStatusNew {
    77  		return MsgCompletedBeforeErr
    78  	}
    79  	msg.status = status
    80  	close(msg.done)
    81  	return nil
    82  }
    83  
    84  func (msg *Message) Body() []byte {
    85  	return msg.body
    86  }
    87  
    88  func (msg *Message) SetBody(body []byte) {
    89  	msg.mutex.Lock()
    90  	defer msg.mutex.Unlock()
    91  	msg.body = body
    92  }
    93  
    94  func (msg *Message) MetaKeys() []interface{} {
    95  	res := make([]interface{}, 0, len(msg.meta))
    96  	for k := range msg.meta {
    97  		res = append(res, k)
    98  	}
    99  	return res
   100  }
   101  
   102  func (msg *Message) Meta(key interface{}) (interface{}, bool) {
   103  	msg.mutex.Lock()
   104  	defer msg.mutex.Unlock()
   105  	return msg.unsafeMeta(key)
   106  }
   107  
   108  func (msg *Message) unsafeMeta(key interface{}) (interface{}, bool) {
   109  	val, ok := msg.meta[key]
   110  	return val, ok
   111  }
   112  
   113  func (msg *Message) SetMeta(key, val interface{}) {
   114  	msg.mutex.Lock()
   115  	defer msg.mutex.Unlock()
   116  	msg.unsafeSetMeta(key, val)
   117  }
   118  
   119  func (msg *Message) unsafeSetMeta(key, val interface{}) {
   120  	msg.meta[key] = val
   121  }
   122  
   123  func (msg *Message) Copy() *Message {
   124  	msg.mutex.Lock()
   125  	defer msg.mutex.Unlock()
   126  	return msg.unsafeCopy()
   127  }
   128  
   129  func (msg *Message) unsafeCopy() *Message {
   130  	cpmsg := NewMessage(msg.Body())
   131  	for key, val := range msg.meta {
   132  		cpmsg.unsafeSetMeta(key, val)
   133  	}
   134  	return cpmsg
   135  }