package events import ( "encoding/json" "fmt" "sync" "time" "github.com/google/uuid" "github.com/pkg/errors" "go-micro.dev/v6/logger" "go-micro.dev/v6/store" ) // NewStream returns an initialized memory stream func NewStream(opts ...Option) (Stream, error) { // parse the options var options Options for _, o := range opts { o(&options) } st := options.Store if st == nil { st = store.NewMemoryStore() } return &mem{store: st}, nil } type subscriber struct { Group string Topic string Channel chan Event sync.RWMutex retryMap map[string]int retryLimit int autoAck bool ackWait time.Duration pending []Event notify chan struct{} } type mem struct { store store.Store subs []*subscriber sync.RWMutex } func (m *mem) Publish(topic string, msg interface{}, opts ...PublishOption) error { // validate the topic if len(topic) == 0 { return ErrMissingTopic } // parse the options options := PublishOptions{ Timestamp: time.Now(), } for _, o := range opts { o(&options) } // encode the message if it's not already encoded var payload []byte if p, ok := msg.([]byte); ok { payload = p } else { p, err := json.Marshal(msg) if err != nil { return ErrEncodingMessage } payload = p } // construct the event event := &Event{ ID: uuid.New().String(), Topic: topic, Timestamp: options.Timestamp, Metadata: options.Metadata, Payload: payload, } // serialize the event to bytes bytes, err := json.Marshal(event) if err != nil { return errors.Wrap(err, "Error encoding event") } // write to the store key := fmt.Sprintf("%v/%v", event.Topic, event.ID) if err := m.store.Write(&store.Record{Key: key, Value: bytes}); err != nil { return errors.Wrap(err, "Error writing event to store") } // send to the subscribers async go m.handleEvent(event) return nil } func (m *mem) Consume(topic string, opts ...ConsumeOption) (<-chan Event, error) { // validate the topic if len(topic) == 0 { return nil, ErrMissingTopic } // parse the options options := ConsumeOptions{ Group: uuid.New().String(), AutoAck: true, } for _, o := range opts { o(&options) } // Note: RetryLimit is configured but retry logic is basic for the in-memory implementation. // For production use with advanced retry capabilities, use NATS JetStream. // setup the subscriber sub := &subscriber{ Channel: make(chan Event), Topic: topic, Group: options.Group, retryMap: map[string]int{}, autoAck: true, retryLimit: options.GetRetryLimit(), notify: make(chan struct{}, 1), } if !options.AutoAck { if options.AckWait == 0 { return nil, fmt.Errorf("invalid AckWait passed, should be positive integer") } sub.autoAck = options.AutoAck sub.ackWait = options.AckWait go sub.dispatchManualAck() } // register the subscriber m.Lock() m.subs = append(m.subs, sub) m.Unlock() // lookup previous events if the start time option was passed if options.Offset.Unix() > 0 { go m.lookupPreviousEvents(sub, options.Offset) } // return the channel return sub.Channel, nil } // lookupPreviousEvents finds events for a subscriber which occurred before a given time and sends // them into the subscribers channel func (m *mem) lookupPreviousEvents(sub *subscriber, startTime time.Time) { // lookup all events which match the topic (a blank topic will return all results) recs, err := m.store.Read(sub.Topic+"/", store.ReadPrefix()) if err != nil && logger.V(logger.ErrorLevel, logger.DefaultLogger) { logger.Errorf("Error looking up previous events: %v", err) return } else if err != nil { return } // loop through the records and send it to the channel if it matches for _, r := range recs { var ev Event if err := json.Unmarshal(r.Value, &ev); err != nil { continue } if ev.Timestamp.Unix() < startTime.Unix() { continue } sendEvent(&ev, sub) } } // handleEvents sends the event to any registered subscribers. func (m *mem) handleEvent(ev *Event) { m.RLock() subs := m.subs m.RUnlock() // filteredSubs is a KV map of the queue name and subscribers. This is used to prevent a message // being sent to two subscribers with the same queue. filteredSubs := map[string]*subscriber{} // filter down to subscribers who are interested in this topic for _, sub := range subs { if len(sub.Topic) == 0 || sub.Topic == ev.Topic { filteredSubs[sub.Group] = sub } } // send the message to each channel async (since one channel might be blocked) for _, sub := range filteredSubs { sendEvent(ev, sub) } } func sendEvent(ev *Event, sub *subscriber) { evCopy := *ev if !sub.autoAck { sub.Lock() sub.pending = append(sub.pending, evCopy) sub.Unlock() sub.wake() return } go func(s *subscriber) { s.Channel <- evCopy }(sub) } func (s *subscriber) wake() { select { case s.notify <- struct{}{}: default: } } func (s *subscriber) dispatchManualAck() { for { s.Lock() for len(s.pending) == 0 { s.Unlock() <-s.notify s.Lock() } ev := s.pending[0] s.pending = s.pending[1:] s.retryMap[ev.ID] = 0 s.Unlock() s.deliverManualAck(ev) } } func (s *subscriber) deliverManualAck(ev Event) { retries := 0 for { if s.retryLimit > -1 && retries > s.retryLimit { if logger.V(logger.ErrorLevel, logger.DefaultLogger) { logger.Errorf("Message retry limit reached, discarding: %v %d %d", ev.ID, retries, s.retryLimit) } s.Lock() delete(s.retryMap, ev.ID) s.Unlock() return } result := make(chan bool, 1) evCopy := ev evCopy.SetAckFunc(func() error { select { case result <- true: default: } return nil }) evCopy.SetNackFunc(func() error { select { case result <- false: default: } return nil }) s.Channel <- evCopy timer := time.NewTimer(s.ackWait) select { case acked := <-result: if !timer.Stop() { select { case <-timer.C: default: } } if acked { s.Lock() delete(s.retryMap, ev.ID) s.Unlock() return } retries++ case <-timer.C: retries++ } s.Lock() s.retryMap[ev.ID] = retries s.Unlock() } }