use sampler lock instead

Signed-off-by: Alex Couture-Beil <alex@earthly.dev>

Author: Alex Couture-Beil

executor/resources/sampler.go

@@ -10,7 +10,7 @@ type WithTimestamp interface {
 }
 
 type Sampler[T WithTimestamp] struct {
-	mu          sync.RWMutex
+	mu          sync.Mutex
 	minInterval time.Duration
 	maxSamples  int
 	callback    func(ts time.Time) (T, error)
@@ -26,19 +26,15 @@ type Sub[T WithTimestamp] struct {
 	first    time.Time
 	last     time.Time
 	samples  []T
-	mu       sync.RWMutex
 	err      error
 }
 
 func (s *Sub[T]) Close(captureLast bool) ([]T, error) {
 	s.sampler.mu.Lock()
 	delete(s.sampler.subs, s)
-	s.sampler.mu.Unlock()
-
-	s.mu.Lock()
-	defer s.mu.Unlock()
 
 	if s.err != nil {
+		s.sampler.mu.Unlock()
 		return nil, s.err
 	}
 	current := s.first
@@ -50,6 +46,7 @@ func (s *Sub[T]) Close(captureLast bool) ([]T, error) {
 			current = ts
 		}
 	}
+	s.sampler.mu.Unlock()
 
 	if captureLast {
 		v, err := s.sampler.callback(time.Now())
@@ -98,26 +95,26 @@ func (s *Sampler[T]) run() {
 			return
 		case <-ticker.C:
 			tm := time.Now()
-			s.mu.RLock()
+			s.mu.Lock()
 			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()
+			s.mu.Unlock()
 			ticker = time.NewTimer(s.minInterval)
 			if len(active) == 0 {
 				continue
 			}
 			value, err := s.callback(tm)
+			s.mu.Lock()
 			for _, ss := range active {
-				ss.mu.Lock()
+				if _, found := s.subs[ss]; !found {
+					continue // skip if Close() was called while the lock was released
+				}
 				if err != nil {
 					ss.err = err
 				} else {
@@ -128,8 +125,8 @@ func (s *Sampler[T]) run() {
 				if time.Duration(ss.interval)*time.Duration(s.maxSamples) <= dur {
 					ss.interval *= 2
 				}
-				ss.mu.Unlock()
 			}
+			s.mu.Unlock()
 		}
 	}
 }