streams.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. package xstream
  2. import (
  3. "context"
  4. "engine/common"
  5. "engine/xstream/operators"
  6. )
  7. type TopologyNew struct {
  8. sources []Source
  9. sinks []Sink
  10. ctx context.Context
  11. cancel context.CancelFunc
  12. drain chan error
  13. ops []Operator
  14. name string
  15. }
  16. func NewWithName(name string) *TopologyNew {
  17. tp := &TopologyNew{name: name}
  18. return tp
  19. }
  20. func (s *TopologyNew) GetContext() context.Context {
  21. return s.ctx
  22. }
  23. func (s *TopologyNew) Cancel(){
  24. s.cancel()
  25. }
  26. func (s *TopologyNew) AddSrc(src Source) *TopologyNew {
  27. s.sources = append(s.sources, src)
  28. return s
  29. }
  30. func (s *TopologyNew) AddSink(inputs []Emitter, snk Sink) *TopologyNew {
  31. for _, input := range inputs{
  32. input.AddOutput(snk.GetInput())
  33. }
  34. s.sinks = append(s.sinks, snk)
  35. return s
  36. }
  37. func (s *TopologyNew) AddOperator(inputs []Emitter, operator Operator) *TopologyNew {
  38. for _, input := range inputs{
  39. input.AddOutput(operator.GetInput())
  40. }
  41. s.ops = append(s.ops, operator)
  42. return s
  43. }
  44. func Transform(op operators.UnOperation, name string) *operators.UnaryOperator {
  45. operator := operators.New(name)
  46. operator.SetOperation(op)
  47. return operator
  48. }
  49. func (s *TopologyNew) Map(f interface{}) *TopologyNew {
  50. log := common.GetLogger(s.ctx)
  51. op, err := MapFunc(f)
  52. if err != nil {
  53. log.Println(err)
  54. }
  55. return s.Transform(op)
  56. }
  57. // Filter takes a predicate user-defined func that filters the stream.
  58. // The specified function must be of type:
  59. // func (T) bool
  60. // If the func returns true, current item continues downstream.
  61. func (s *TopologyNew) Filter(f interface{}) *TopologyNew {
  62. op, err := FilterFunc(f)
  63. if err != nil {
  64. s.drainErr(err)
  65. }
  66. return s.Transform(op)
  67. }
  68. // Transform is the base method used to apply transfomrmative
  69. // unary operations to streamed elements (i.e. filter, map, etc)
  70. // It is exposed here for completeness, use the other more specific methods.
  71. func (s *TopologyNew) Transform(op operators.UnOperation) *TopologyNew {
  72. operator := operators.New("default")
  73. operator.SetOperation(op)
  74. s.ops = append(s.ops, operator)
  75. return s
  76. }
  77. // prepareContext setups internal context before
  78. // stream starts execution.
  79. func (s *TopologyNew) prepareContext() {
  80. if s.ctx == nil || s.ctx.Err() != nil {
  81. s.ctx, s.cancel = context.WithCancel(context.Background())
  82. contextLogger := common.Log.WithField("rule", s.name)
  83. s.ctx = context.WithValue(s.ctx, common.LoggerKey, contextLogger)
  84. }
  85. }
  86. func (s *TopologyNew) drainErr(err error) {
  87. go func() { s.drain <- err }()
  88. }
  89. func (s *TopologyNew) Open() <-chan error {
  90. s.prepareContext() // ensure context is set
  91. log := common.GetLogger(s.ctx)
  92. log.Println("Opening stream")
  93. // open stream
  94. go func() {
  95. // open source, if err bail
  96. for _, src := range s.sources{
  97. if err := src.Open(s.ctx); err != nil {
  98. s.drainErr(err)
  99. log.Println("Closing stream")
  100. return
  101. }
  102. }
  103. //apply operators, if err bail
  104. for _, op := range s.ops {
  105. if err := op.Exec(s.ctx); err != nil {
  106. s.drainErr(err)
  107. log.Println("Closing stream")
  108. return
  109. }
  110. }
  111. sinkErr := make(chan error)
  112. defer func() {
  113. log.Println("Closing sinkErr channel")
  114. close(sinkErr)
  115. }()
  116. // open stream sink, after log sink is ready.
  117. for _, snk := range s.sinks{
  118. snk.Open(s.ctx, sinkErr)
  119. }
  120. select {
  121. case err := <- sinkErr:
  122. log.Println("Closing stream")
  123. s.drain <- err
  124. }
  125. }()
  126. return s.drain
  127. }