internal/kobject: range over objects

For matching devices based on some criteria.

Signed-off-by: Ophestra <cat@gensokyo.uk>

internal/kobject/kobject.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"maps"
+	"slices"
 	"strconv"
 	"sync"
 
@@ -96,6 +97,13 @@ func (o *Object) update(env map[string]string, strip bool) {
 	}
 }
 
+// A pendingIterator is a callback currently iterating through objects targeted
+// by ongoing events.
+type pendingIterator struct {
+	f    func(o *Object) bool
+	done chan<- struct{}
+}
+
 // State processes a stream of [Event] populated from [uevent.Message] received
 // from a NETLINK_KOBJECT_UEVENT socket and presents an efficient representation
 // of kernel state.
@@ -106,6 +114,10 @@ type State struct {
 	uevent map[string]*Object
 	// Synchronises access to uevent and its objects.
 	ueventMu sync.RWMutex
+	// Alive iterators.
+	iter []*pendingIterator
+	// Synchronises access to iter.
+	iterMu sync.Mutex
 	// UUID for synthetic [uevent.Coldboot] events.
 	coldboot uevent.UUID
 	// Called on [uevent.KOBJ_CHANGE] with stripped environment variables.
@@ -129,6 +141,59 @@ func New(
 	}
 }
 
+// deleteIter removes an iterator from s. Must be called after acquiring iterMu.
+func (s *State) deleteIter(p *pendingIterator) {
+	s.iter = slices.DeleteFunc(s.iter, func(v *pendingIterator) bool {
+		return p == v
+	})
+}
+
+// dispatchIter broadcasts an [Object] to all alive iterators.
+func (s *State) dispatchIter(o *Object) {
+	s.iterMu.Lock()
+	defer s.iterMu.Unlock()
+
+	for _, p := range s.iter {
+		if !p.f(o) {
+			s.deleteIter(p)
+			close(p.done)
+		}
+	}
+}
+
+// Range calls f on all current and upcoming [Object] values tracked by s until
+// f returns false or the context is cancelled. f must not retain o or modify
+// the value it points to.
+func (s *State) Range(ctx context.Context, f func(o *Object) bool) {
+	done := make(chan struct{})
+	p := pendingIterator{f, done}
+
+	s.iterMu.Lock()
+	s.ueventMu.RLock()
+	for _, o := range s.uevent {
+		if !f(o) {
+			s.ueventMu.RUnlock()
+			s.iterMu.Unlock()
+			return
+		}
+	}
+	s.ueventMu.RUnlock()
+	s.iter = append(s.iter, &p)
+	s.iterMu.Unlock()
+
+	select {
+	case <-ctx.Done():
+		s.iterMu.Lock()
+		s.deleteIter(&p)
+		s.iterMu.Unlock()
+		return
+
+	case <-done:
+		// deregistered by dispatchIter
+		return
+	}
+}
+
 // UnexpectedColdbootError is reported by [State.Consume] for a coldboot event
 // with action other than the expected [uevent.KOBJ_ADD].
 type UnexpectedColdbootError Event
@@ -219,6 +284,7 @@ func (s *State) processEvent(e *Event) {
 			if o, ok := s.uevent[e.DevPath]; ok {
 				s.reportErr(DuplicateAddError(e.Clone()))
 				o.merge(e.Env)
+				s.dispatchIter(o)
 				return
 			}
 		}
@@ -228,6 +294,7 @@ func (s *State) processEvent(e *Event) {
 		}
 		o.merge(e.Env)
 		s.uevent[e.DevPath] = o
+		s.dispatchIter(o)
 		return
 
 	case uevent.KOBJ_REMOVE:
@@ -253,12 +320,14 @@ func (s *State) processEvent(e *Event) {
 			o = e.makeColdboot()
 			o.merge(e.Env)
 			s.uevent[e.DevPath] = o
+			s.dispatchIter(o)
 			return
 		}
 		o.update(e.Env, true)
 		if s.handleChange != nil {
 			s.handleChange(o, e.Env)
 		}
+		s.dispatchIter(o)
 		return
 
 	case uevent.KOBJ_MOVE:
@@ -281,6 +350,7 @@ func (s *State) processEvent(e *Event) {
 		o.merge(e.Env)
 		s.uevent[e.DevPath] = o
 		o.DevPath = e.DevPath
+		s.dispatchIter(o)
 		return
 
 	case uevent.KOBJ_ONLINE:
@@ -296,6 +366,7 @@ func (s *State) processEvent(e *Event) {
 			s.reportErr(UnexpectedOfflineError(o.Clone()))
 		}
 		o.Offline = false
+		s.dispatchIter(o)
 		return
 
 	case uevent.KOBJ_OFFLINE:
@@ -311,6 +382,7 @@ func (s *State) processEvent(e *Event) {
 			s.reportErr(UnexpectedOfflineError(o.Clone()))
 		}
 		o.Offline = true
+		s.dispatchIter(o)
 		return
 
 	case uevent.KOBJ_BIND:
@@ -326,6 +398,7 @@ func (s *State) processEvent(e *Event) {
 		}
 		o.State = StateBound
 		o.merge(e.Env)
+		s.dispatchIter(o)
 		return
 
 	case uevent.KOBJ_UNBIND:
@@ -341,6 +414,7 @@ func (s *State) processEvent(e *Event) {
 		}
 		o.State = StateNew
 		o.Driver = ""
+		s.dispatchIter(o)
 		return
 
 	default: // not reached