go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/simutil/event_source.go (about)

     1  /*
     2  
     3  Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package simutil
     9  
    10  import (
    11  	"context"
    12  	"time"
    13  )
    14  
    15  var (
    16  	_ Waiter = (*EventSource[any])(nil)
    17  )
    18  
    19  // EventSource helps produce events given an interval.
    20  //
    21  // You must provide both an `Interval` delegate AND an `Event` delegate.
    22  type EventSource[T any] struct {
    23  	Count     uint64
    24  	StopAfter time.Duration
    25  
    26  	Interval     func(context.Context) time.Duration
    27  	EventFactory func(context.Context) T
    28  
    29  	stopped chan struct{}
    30  }
    31  
    32  // Waiter is an interface that EventSource[T] implements.
    33  type Waiter interface {
    34  	Wait()
    35  }
    36  
    37  // Start begins the event source in either "simulation" or "realtime" mode.
    38  //
    39  // "Simulation" mode is enabled by passing `false` a the `realtime` flag and
    40  // produces messages with no real-world delay.
    41  //
    42  // "Realtime" mode waits the given intervals between events.
    43  //
    44  // In this mode it will publish the events to a channel
    45  // based on poisson distributed arrival times.
    46  //
    47  // The call does not block, and instead starts a background goroutine
    48  // which you can cancel with the given context.
    49  func (es *EventSource[T]) Start(ctx context.Context, realtime bool) <-chan Event[T] {
    50  	output := make(chan Event[T])
    51  	es.stopped = make(chan struct{})
    52  	go func() {
    53  
    54  		var started = time.Now()
    55  		var ts time.Time
    56  		var produced uint64
    57  		var t *time.Timer
    58  
    59  		defer func() {
    60  			if t != nil {
    61  				t.Stop()
    62  			}
    63  			close(es.stopped)
    64  		}()
    65  
    66  		for {
    67  			select {
    68  			case <-ctx.Done():
    69  				return
    70  			default:
    71  			}
    72  
    73  			interval := es.Interval(ctx)
    74  			if realtime {
    75  				if t == nil {
    76  					t = time.NewTimer(interval)
    77  				} else {
    78  					t.Reset(interval)
    79  				}
    80  				select {
    81  				case <-ctx.Done():
    82  					return
    83  				case <-t.C:
    84  					m := es.EventFactory(ctx)
    85  					e := Event[T]{
    86  						Timestamp: time.Now(),
    87  						Event:     m,
    88  					}
    89  					select {
    90  					case <-ctx.Done():
    91  						return
    92  					case output <- e:
    93  					}
    94  				}
    95  			} else {
    96  				if ts.IsZero() {
    97  					ts = started
    98  				}
    99  				ts = ts.Add(interval)
   100  				m := es.EventFactory(ctx)
   101  				e := Event[T]{
   102  					Timestamp: ts,
   103  					Event:     m,
   104  				}
   105  				select {
   106  				case <-ctx.Done():
   107  					return
   108  				case output <- e:
   109  				}
   110  			}
   111  
   112  			produced++
   113  			if es.Count > 0 && produced >= es.Count {
   114  				return
   115  			}
   116  
   117  			if realtime {
   118  				if es.StopAfter > 0 && time.Since(started) >= es.StopAfter {
   119  					return
   120  				}
   121  			} else {
   122  				if es.StopAfter > 0 && ts.Sub(started) >= es.StopAfter {
   123  					return
   124  				}
   125  			}
   126  		}
   127  	}()
   128  	return output
   129  }
   130  
   131  // Wait waits for the event source to finish producing.
   132  func (es EventSource[T]) Wait() {
   133  	<-es.stopped
   134  }