fix probemanager close serialization for task 55

2 files changed, 184 insertions, 31 deletions

diff --git a/internal/probemanager/manager.go b/internal/probemanager/manager.go
index bb6d44d..288af41 100644
--- a/internal/probemanager/manager.go
+++ b/internal/probemanager/manager.go
@@ -136,14 +136,6 @@ func (m *Manager) Attach(syscall string) error {
 		m.mu.Unlock()
 		return err
 	}
-	if entry.active {
-		m.mu.Unlock()
-		return nil
-	}
-
-	enterTP := entry.enterTP
-	exitTP := entry.exitTP
-	attacher := m.attacher
 	m.mu.Unlock()
 	entry.attachMu.Lock()
 	defer entry.attachMu.Unlock()
@@ -158,6 +150,9 @@ func (m *Manager) Attach(syscall string) error {
 		m.mu.Unlock()
 		return nil
 	}
+	enterTP := entry.enterTP
+	exitTP := entry.exitTP
+	attacher := m.attacher
 	m.mu.Unlock()
 
 	enterLink, exitLink, attachErr := attachPair(attacher, enterTP, exitTP)
@@ -230,10 +225,6 @@ func (m *Manager) Detach(syscall string) error {
 
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	entry, err = m.entryLocked(syscall)
-	if err != nil {
-		return err
-	}
 	if enterErr == nil {
 		entry.enterLink = nil
 	}
@@ -319,39 +310,47 @@ func (m *Manager) Close() error {
 		m.mu.Unlock()
 		return nil
 	}
-	type pairLinks struct {
-		syscall   string
-		enterLink Link
-		exitLink  Link
+	type pairEntry struct {
+		syscall  string
+		entry    *probeEntry
+		hasLinks bool
 	}
-	links := make([]pairLinks, 0, len(m.probes))
+	entries := make([]pairEntry, 0, len(m.probes))
 	for syscall, entry := range m.probes {
-		links = append(links, pairLinks{
-			syscall:   syscall,
-			enterLink: entry.enterLink,
-			exitLink:  entry.exitLink,
+		entries = append(entries, pairEntry{
+			syscall:  syscall,
+			entry:    entry,
+			hasLinks: entry.enterLink != nil || entry.exitLink != nil,
 		})
-		entry.enterLink = nil
-		entry.exitLink = nil
-		entry.active = false
-		entry.lastErr = nil
 	}
 	m.closed = true
 	m.mu.Unlock()
 
 	var firstErr error
-	for _, l := range links {
+	for _, item := range entries {
+		if item.hasLinks {
+			item.entry.attachMu.Lock()
+		}
 		var errForSyscall error
-		if l.enterLink != nil {
-			if err := l.enterLink.Destroy(); err != nil {
+		m.mu.Lock()
+		enterLink := item.entry.enterLink
+		exitLink := item.entry.exitLink
+		item.entry.enterLink = nil
+		item.entry.exitLink = nil
+		item.entry.active = false
+		item.entry.lastErr = nil
+		m.mu.Unlock()
+
+		if enterLink != nil {
+			if err := enterLink.Destroy(); err != nil {
 				errForSyscall = err
 				if firstErr == nil {
 					firstErr = err
 				}
 			}
 		}
-		if l.exitLink != nil {
-			if err := l.exitLink.Destroy(); err != nil {
+		if exitLink != nil {
+			if err := exitLink.Destroy(); err != nil {
 				if errForSyscall == nil {
 					errForSyscall = err
 				}
@@ -360,7 +359,10 @@ func (m *Manager) Close() error {
 				}
 			}
 		}
-		m.setLastError(l.syscall, errForSyscall)
+		m.setLastError(item.syscall, errForSyscall)
+		if item.hasLinks {
+			item.entry.attachMu.Unlock()
+		}
 	}
 	return firstErr
 }
diff --git a/internal/probemanager/manager_test.go b/internal/probemanager/manager_test.go
index 4626af0..dc0c474 100644
--- a/internal/probemanager/manager_test.go
+++ b/internal/probemanager/manager_test.go
@@ -9,15 +9,30 @@ import (
 )
 
 type fakeLink struct {
+	mu        sync.Mutex
 	destroyed int
 	err       error
+	onDestroy func()
 }
 
 func (l *fakeLink) Destroy() error {
+	l.mu.Lock()
 	l.destroyed++
+	onDestroy := l.onDestroy
+	l.mu.Unlock()
+
+	if onDestroy != nil {
+		onDestroy()
+	}
 	return l.err
 }
 
+func (l *fakeLink) destroyCalls() int {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.destroyed
+}
+
 type fakeProgram struct {
 	mu         sync.Mutex
 	tracepoint string
@@ -175,6 +190,142 @@ func TestManagerAttachSerializesConcurrentCalls(t *testing.T) {
 	}
 }
 
+func TestManagerAttachWaitsForDetachBeforeReturning(t *testing.T) {
+	enterDestroyStarted := make(chan struct{})
+	releaseDestroy := make(chan struct{})
+	var enterDestroyOnce sync.Once
+	enter := &fakeLink{}
+	enter.onDestroy = func() {
+		enterDestroyOnce.Do(func() { close(enterDestroyStarted) })
+		<-releaseDestroy
+	}
+	exit := &fakeLink{}
+	enterProg := &fakeProgram{link: enter}
+	exitProg := &fakeProgram{link: exit}
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{
+			"handle_sys_enter_close": enterProg,
+			"handle_sys_exit_close":  exitProg,
+		},
+		errs: map[string]error{},
+	}
+	mgr := NewManager(attacher)
+	if err := mgr.AttachAll(nil, []string{"sys_enter_close", "sys_exit_close"}); err != nil {
+		t.Fatalf("AttachAll returned error: %v", err)
+	}
+
+	detachErrCh := make(chan error, 1)
+	go func() {
+		detachErrCh <- mgr.Detach("close")
+	}()
+
+	select {
+	case <-enterDestroyStarted:
+	case <-time.After(time.Second):
+		t.Fatal("detach did not start destroying the enter link")
+	}
+
+	attachErrCh := make(chan error, 1)
+	go func() {
+		attachErrCh <- mgr.Attach("close")
+	}()
+
+	select {
+	case err := <-attachErrCh:
+		t.Fatalf("Attach returned before Detach completed: %v", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(releaseDestroy)
+
+	if err := <-detachErrCh; err != nil {
+		t.Fatalf("Detach returned error: %v", err)
+	}
+	if err := <-attachErrCh; err != nil {
+		t.Fatalf("Attach returned error: %v", err)
+	}
+	if got := enterProg.attachCalls(); got != 2 {
+		t.Fatalf("expected enter attach to run twice, got %d", got)
+	}
+	if got := exitProg.attachCalls(); got != 2 {
+		t.Fatalf("expected exit attach to run twice, got %d", got)
+	}
+	if got := enter.destroyCalls(); got != 1 {
+		t.Fatalf("expected enter link to be destroyed once, got %d", got)
+	}
+	if got := exit.destroyCalls(); got != 1 {
+		t.Fatalf("expected exit link to be destroyed once, got %d", got)
+	}
+	if !mgr.IsActive("close") {
+		t.Fatalf("expected probe to be active after detach followed by attach")
+	}
+}
+
+func TestManagerCloseWaitsForDetachAndDoesNotDoubleDestroy(t *testing.T) {
+	enterDestroyStarted := make(chan struct{})
+	releaseDestroy := make(chan struct{})
+	var enterDestroyOnce sync.Once
+	enter := &fakeLink{}
+	enter.onDestroy = func() {
+		enterDestroyOnce.Do(func() { close(enterDestroyStarted) })
+		<-releaseDestroy
+	}
+	exit := &fakeLink{}
+	enterProg := &fakeProgram{link: enter}
+	exitProg := &fakeProgram{link: exit}
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{
+			"handle_sys_enter_close": enterProg,
+			"handle_sys_exit_close":  exitProg,
+		},
+		errs: map[string]error{},
+	}
+	mgr := NewManager(attacher)
+	if err := mgr.AttachAll(nil, []string{"sys_enter_close", "sys_exit_close"}); err != nil {
+		t.Fatalf("AttachAll returned error: %v", err)
+	}
+
+	detachErrCh := make(chan error, 1)
+	go func() {
+		detachErrCh <- mgr.Detach("close")
+	}()
+
+	select {
+	case <-enterDestroyStarted:
+	case <-time.After(time.Second):
+		t.Fatal("detach did not start destroying the enter link")
+	}
+
+	closeErrCh := make(chan error, 1)
+	go func() {
+		closeErrCh <- mgr.Close()
+	}()
+
+	select {
+	case err := <-closeErrCh:
+		t.Fatalf("Close returned before Detach completed: %v", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(releaseDestroy)
+
+	if err := <-detachErrCh; err != nil {
+		t.Fatalf("Detach returned error: %v", err)
+	}
+	if err := <-closeErrCh; err != nil {
+		t.Fatalf("Close returned error: %v", err)
+	}
+	if got := enter.destroyCalls(); got != 1 {
+		t.Fatalf("expected enter link to be destroyed once, got %d", got)
+	}
+	if got := exit.destroyCalls(); got != 1 {
+		t.Fatalf("expected exit link to be destroyed once, got %d", got)
+	}
+	if mgr.IsActive("close") {
+		t.Fatalf("expected probe to be inactive after Close")
+	}
+}
+
 func TestManagerDetachDestroysLinks(t *testing.T) {
 	enter := &fakeLink{}
 	exit := &fakeLink{}