fix probemanager attach race for task 55

2 files changed, 105 insertions, 2 deletions

diff --git a/internal/probemanager/manager.go b/internal/probemanager/manager.go
index cad755e..bb6d44d 100644
--- a/internal/probemanager/manager.go
+++ b/internal/probemanager/manager.go
@@ -38,6 +38,7 @@ type probeEntry struct {
 
 	enterLink Link
 	exitLink  Link
+	attachMu  sync.Mutex
 
 	active  bool
 	lastErr error
@@ -144,6 +145,20 @@ func (m *Manager) Attach(syscall string) error {
 	exitTP := entry.exitTP
 	attacher := m.attacher
 	m.mu.Unlock()
+	entry.attachMu.Lock()
+	defer entry.attachMu.Unlock()
+
+	m.mu.Lock()
+	entry, err = m.entryLocked(syscall)
+	if err != nil {
+		m.mu.Unlock()
+		return err
+	}
+	if entry.active {
+		m.mu.Unlock()
+		return nil
+	}
+	m.mu.Unlock()
 
 	enterLink, exitLink, attachErr := attachPair(attacher, enterTP, exitTP)
 
@@ -183,6 +198,16 @@ func (m *Manager) Detach(syscall string) error {
 		m.mu.Unlock()
 		return err
 	}
+	m.mu.Unlock()
+	entry.attachMu.Lock()
+	defer entry.attachMu.Unlock()
+
+	m.mu.Lock()
+	entry, err = m.entryLocked(syscall)
+	if err != nil {
+		m.mu.Unlock()
+		return err
+	}
 	enterLink := entry.enterLink
 	exitLink := entry.exitLink
 	m.mu.Unlock()
diff --git a/internal/probemanager/manager_test.go b/internal/probemanager/manager_test.go
index cc0233b..4626af0 100644
--- a/internal/probemanager/manager_test.go
+++ b/internal/probemanager/manager_test.go
@@ -3,7 +3,9 @@ package probemanager
 import (
 	"errors"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 )
 
 type fakeLink struct {
@@ -17,16 +19,23 @@ func (l *fakeLink) Destroy() error {
 }
 
 type fakeProgram struct {
+	mu         sync.Mutex
 	tracepoint string
+	attachs    int
 	link       *fakeLink
 	err        error
 	onAttach   func()
 }
 
 func (p *fakeProgram) AttachTracepoint(_, name string) (Link, error) {
+	p.mu.Lock()
 	p.tracepoint = name
-	if p.onAttach != nil {
-		p.onAttach()
+	p.attachs++
+	onAttach := p.onAttach
+	p.mu.Unlock()
+
+	if onAttach != nil {
+		onAttach()
 	}
 	if p.err != nil {
 		return nil, p.err
@@ -37,6 +46,12 @@ func (p *fakeProgram) AttachTracepoint(_, name string) (Link, error) {
 	return p.link, nil
 }
 
+func (p *fakeProgram) attachCalls() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.attachs
+}
+
 type fakeAttacher struct {
 	programs map[string]*fakeProgram
 	errs     map[string]error
@@ -97,6 +112,69 @@ func TestManagerAttachAllToggleAndCounts(t *testing.T) {
 	}
 }
 
+func TestManagerAttachSerializesConcurrentCalls(t *testing.T) {
+	enterBlocked := make(chan struct{})
+	releaseEnter := make(chan struct{})
+	enter := &fakeProgram{link: &fakeLink{}}
+	var enteredOnce sync.Once
+	enter.onAttach = func() {
+		enteredOnce.Do(func() { close(enterBlocked) })
+		<-releaseEnter
+	}
+	exit := &fakeProgram{link: &fakeLink{}}
+	attacher := &fakeAttacher{
+		programs: map[string]*fakeProgram{
+			"handle_sys_enter_close": enter,
+			"handle_sys_exit_close":  exit,
+		},
+		errs: map[string]error{},
+	}
+	mgr := NewManager(attacher)
+	mgr.Register("close", TracepointPair{Enter: "sys_enter_close", Exit: "sys_exit_close"})
+
+	errCh1 := make(chan error, 1)
+	go func() {
+		errCh1 <- mgr.Attach("close")
+	}()
+
+	select {
+	case <-enterBlocked:
+	case <-time.After(time.Second):
+		t.Fatal("first attach did not start")
+	}
+
+	errCh2 := make(chan error, 1)
+	go func() {
+		errCh2 <- mgr.Attach("close")
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+	if got := enter.attachCalls(); got != 1 {
+		t.Fatalf("expected only one enter attach while first call was in flight, got %d", got)
+	}
+	if got := exit.attachCalls(); got != 0 {
+		t.Fatalf("expected exit attach to wait for enter completion, got %d", got)
+	}
+
+	close(releaseEnter)
+
+	if err := <-errCh1; err != nil {
+		t.Fatalf("first Attach returned error: %v", err)
+	}
+	if err := <-errCh2; err != nil {
+		t.Fatalf("second Attach returned error: %v", err)
+	}
+	if got := enter.attachCalls(); got != 1 {
+		t.Fatalf("expected enter attach to run once, got %d", got)
+	}
+	if got := exit.attachCalls(); got != 1 {
+		t.Fatalf("expected exit attach to run once, got %d", got)
+	}
+	if !mgr.IsActive("close") {
+		t.Fatalf("expected probe to remain active after concurrent attach calls")
+	}
+}
+
 func TestManagerDetachDestroysLinks(t *testing.T) {
 	enter := &fakeLink{}
 	exit := &fakeLink{}