sampler data race fix

Protect against data race occuring while both sampler.Run and sub.Close
accesses the same data from multiple threads.

Signed-off-by: Alex Couture-Beil <[email protected]>

Author: Alex Couture-Beil

executor/resources/sampler.go

@@ -26,6 +26,7 @@ type Sub[T WithTimestamp] struct {
 	first    time.Time
 	last     time.Time
 	samples  []T
+	mu       sync.RWMutex
 	err      error
 }
 
@@ -34,6 +35,9 @@ func (s *Sub[T]) Close(captureLast bool) ([]T, error) {
 	delete(s.sampler.subs, s)
 	s.sampler.mu.Unlock()
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	if s.err != nil {
 		return nil, s.err
 	}
@@ -94,13 +98,16 @@ func (s *Sampler[T]) run() {
 			return
 		case <-ticker.C:
 			tm := time.Now()
-			active := make([]*Sub[T], 0, len(s.subs))
 			s.mu.RLock()
+			active := make([]*Sub[T], 0, len(s.subs))
 			for ss := range s.subs {
+				ss.mu.Lock()
 				if tm.Sub(ss.last) < ss.interval {
+					ss.mu.Unlock()
 					continue
 				}
 				ss.last = tm
+				ss.mu.Unlock()
 				active = append(active, ss)
 			}
 			s.mu.RUnlock()
@@ -110,6 +117,7 @@ func (s *Sampler[T]) run() {
 			}
 			value, err := s.callback(tm)
 			for _, ss := range active {
+				ss.mu.Lock()
 				if err != nil {
 					ss.err = err
 				} else {
@@ -120,6 +128,7 @@ func (s *Sampler[T]) run() {
 				if time.Duration(ss.interval)*time.Duration(s.maxSamples) <= dur {
 					ss.interval *= 2
 				}
+				ss.mu.Unlock()
 			}
 		}
 	}