container: sync stubbed wait4 loop after notify

This ensures consistent state observed by wait4 loop when running against stub.

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

1 files changed, 16 insertions, 4 deletions

diff --git a/container/dispatcher_test.go b/container/dispatcher_test.go
index 834c700d..26c49323 100644
--- a/container/dispatcher_test.go
+++ b/container/dispatcher_test.go
@@ -162,7 +162,8 @@ func checkSimple(t *testing.T, fname string, testCases []simpleTestCase) {
 			t.Parallel()
 
 			wait4signal := make(chan struct{})
-			k := &kstub{wait4signal, stub.New(t, func(s *stub.Stub[syscallDispatcher]) syscallDispatcher { return &kstub{wait4signal, s} }, tc.want)}
+			lockNotify := make(chan struct{})
+			k := &kstub{wait4signal, lockNotify, stub.New(t, func(s *stub.Stub[syscallDispatcher]) syscallDispatcher { return &kstub{wait4signal, lockNotify, s} }, tc.want)}
 			defer stub.HandleExit(t)
 			if err := tc.f(k); !reflect.DeepEqual(err, tc.wantErr) {
 				t.Errorf("%s: error = %v, want %v", fname, err, tc.wantErr)
@@ -200,8 +201,8 @@ func checkOpBehaviour(t *testing.T, testCases []opBehaviourTestCase) {
 				t.Helper()
 				t.Parallel()
 
-				k := &kstub{nil, stub.New(t,
-					func(s *stub.Stub[syscallDispatcher]) syscallDispatcher { return &kstub{nil, s} },
+				k := &kstub{nil, nil, stub.New(t,
+					func(s *stub.Stub[syscallDispatcher]) syscallDispatcher { return &kstub{nil, nil, s} },
 					stub.Expect{Calls: slices.Concat(tc.early, []stub.Call{{Name: stub.CallSeparator}}, tc.apply)},
 				)}
 				state := &setupState{Params: tc.params, Msg: k}
@@ -322,12 +323,19 @@ const (
 
 type kstub struct {
 	wait4signal chan struct{}
+	lockNotify  chan struct{}
 	*stub.Stub[syscallDispatcher]
 }
 
 func (k *kstub) new(f func(k syscallDispatcher)) { k.Helper(); k.New(f) }
 
-func (k *kstub) lockOSThread() { k.Helper(); k.Expects("lockOSThread") }
+func (k *kstub) lockOSThread() {
+	k.Helper()
+	expect := k.Expects("lockOSThread")
+	if k.lockNotify != nil && expect.Ret == magicWait4Signal {
+		<-k.lockNotify
+	}
+}
 
 func (k *kstub) setPtracer(pid uintptr) error {
 	k.Helper()
@@ -472,6 +480,10 @@ func (k *kstub) notify(c chan<- os.Signal, sig ...os.Signal) {
 		k.FailNow()
 	}
 
+	if k.lockNotify != nil && expect.Ret == magicWait4Signal {
+		defer close(k.lockNotify)
+	}
+
 	// export channel for external instrumentation
 	if chanf, ok := expect.Args[0].(func(c chan<- os.Signal)); ok && chanf != nil {
 		chanf(c)