Skip to content

Commit 4837cb3

Browse files
tonyredondotonyredondo
committed
fix(ci-visibility): stop signal handler on shutdown
1 parent 550ec3e commit 4837cb3

5 files changed

Lines changed: 372 additions & 18 deletions

File tree

‎internal/civisibility/integrations/civisibility.go‎

Lines changed: 109 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"regexp"
1212
"strings"
1313
"sync"
14+
"sync/atomic"
1415
"syscall"
1516

1617
"github.com/DataDog/dd-trace-go/v2/ddtrace/mocktracer"
@@ -37,10 +38,29 @@ type ciVisibilityIdleConnectionCloser interface {
3738
CloseIdleConnections()
3839
}
3940

41+
// ciVisibilitySignalHandler owns the SIGINT/SIGTERM goroutine registered by CI
42+
// Visibility and lets normal test shutdown stop it before goleak runs.
43+
type ciVisibilitySignalHandler struct {
44+
signals chan os.Signal // receives process interrupt and termination signals
45+
stop chan struct{} // asks the handler goroutine to exit during normal shutdown
46+
done chan struct{} // closes when the handler goroutine exits
47+
stopOnce sync.Once // guarantees stop is signaled at most once
48+
stopping atomic.Bool // suppresses os.Exit when normal shutdown already started
49+
}
50+
4051
var (
4152
// ciVisibilityInitializationOnce ensures we initialize the CI visibility tracer only once.
4253
ciVisibilityInitializationOnce sync.Once
4354

55+
// ciVisibilitySignalHandlerMu synchronizes access to activeCIVisibilitySignalHandler.
56+
ciVisibilitySignalHandlerMu sync.Mutex
57+
58+
// activeCIVisibilitySignalHandler contains the active process signal handler, if CI Visibility started one.
59+
activeCIVisibilitySignalHandler *ciVisibilitySignalHandler
60+
61+
// ciVisibilitySignalExitFunc terminates the process after signal-triggered shutdown; tests replace it.
62+
ciVisibilitySignalExitFunc = os.Exit
63+
4464
// closeActions holds CI visibility close actions.
4565
closeActions []ciVisibilityCloseAction
4666

@@ -129,15 +149,83 @@ func internalCiVisibilityInitialization(tracerInitializer func([]tracer.StartOpt
129149

130150
initializeCiVisibilityLogs(serviceName)
131151

132-
// Handle SIGINT and SIGTERM signals to ensure we close all open spans and flush the tracer before exiting
133-
signals := make(chan os.Signal, 1)
134-
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
135-
go func() {
136-
<-signals
137-
ExitCiVisibility()
138-
os.Exit(1)
139-
}()
152+
startCIVisibilitySignalHandler()
153+
})
154+
}
155+
156+
// run waits for either a process signal or a normal shutdown request.
157+
func (handler *ciVisibilitySignalHandler) run() {
158+
defer close(handler.done)
159+
160+
select {
161+
case <-handler.signals:
162+
if handler.stopping.Load() {
163+
return
164+
}
165+
exitCiVisibility(false)
166+
ciVisibilitySignalExitFunc(1)
167+
case <-handler.stop:
168+
return
169+
}
170+
}
171+
172+
// startCIVisibilitySignalHandler registers the process signal handler once for
173+
// the current CI Visibility initialization.
174+
func startCIVisibilitySignalHandler() {
175+
ciVisibilitySignalHandlerMu.Lock()
176+
defer ciVisibilitySignalHandlerMu.Unlock()
177+
178+
if activeCIVisibilitySignalHandler != nil {
179+
return
180+
}
181+
182+
handler := &ciVisibilitySignalHandler{
183+
signals: make(chan os.Signal, 1),
184+
stop: make(chan struct{}),
185+
done: make(chan struct{}),
186+
}
187+
188+
activeCIVisibilitySignalHandler = handler
189+
signal.Notify(handler.signals, syscall.SIGINT, syscall.SIGTERM)
190+
go handler.run()
191+
}
192+
193+
// markCIVisibilitySignalHandlerStopping prevents a late buffered signal from
194+
// converting an already-started normal shutdown into os.Exit(1).
195+
func markCIVisibilitySignalHandlerStopping() {
196+
ciVisibilitySignalHandlerMu.Lock()
197+
handler := activeCIVisibilitySignalHandler
198+
ciVisibilitySignalHandlerMu.Unlock()
199+
200+
if handler != nil {
201+
handler.stopping.Store(true)
202+
}
203+
}
204+
205+
// stopCIVisibilitySignalHandler stops the active signal handler and waits for
206+
// its goroutine to exit. It is safe to call repeatedly.
207+
func stopCIVisibilitySignalHandler() {
208+
ciVisibilitySignalHandlerMu.Lock()
209+
handler := activeCIVisibilitySignalHandler
210+
ciVisibilitySignalHandlerMu.Unlock()
211+
212+
if handler == nil {
213+
return
214+
}
215+
216+
handler.stopOnce.Do(func() {
217+
handler.stopping.Store(true)
218+
signal.Stop(handler.signals)
219+
close(handler.stop)
140220
})
221+
222+
<-handler.done
223+
224+
ciVisibilitySignalHandlerMu.Lock()
225+
if activeCIVisibilitySignalHandler == handler {
226+
activeCIVisibilitySignalHandler = nil
227+
}
228+
ciVisibilitySignalHandlerMu.Unlock()
141229
}
142230

143231
// initializeCiVisibilityLogs starts CI Visibility log shipping only when logs are enabled and Bazel offline/file modes do not suppress it.
@@ -172,12 +260,25 @@ func PushCiVisibilityCloseAction(action ciVisibilityCloseAction) {
172260

173261
// ExitCiVisibility executes all registered close actions and stops the tracer.
174262
func ExitCiVisibility() {
263+
markCIVisibilitySignalHandlerStopping()
264+
exitCiVisibility(true)
265+
}
266+
267+
// exitCiVisibility executes CI Visibility shutdown and optionally stops the
268+
// signal handler. Signal-triggered shutdown skips that wait to avoid self-deadlock.
269+
func exitCiVisibility(stopSignalHandler bool) {
175270
if civisibility.GetState() != civisibility.StateInitialized {
176271
log.Debug("civisibility: already closed or not initialized")
272+
if stopSignalHandler {
273+
stopCIVisibilitySignalHandler()
274+
}
177275
return
178276
}
179277

180278
civisibility.SetState(civisibility.StateExiting)
279+
if stopSignalHandler {
280+
defer stopCIVisibilitySignalHandler()
281+
}
181282
defer civisibility.SetState(civisibility.StateExited)
182283
log.Debug("civisibility: exiting")
183284
closeActionsMutex.Lock()

‎internal/civisibility/integrations/civisibility_features.go‎

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,18 @@ type (
3939
RemoteCommits []string
4040
IsOk bool
4141
}
42+
43+
// repositoryUploadHooks captures repository upload test seams for a single upload attempt.
44+
repositoryUploadHooks struct {
45+
// uploadRepositoryChanges replaces the full upload path when tests need to bypass git operations.
46+
uploadRepositoryChanges func() (int64, error)
47+
// getSearchCommits retrieves the local/remote commit comparison for the current repository state.
48+
getSearchCommits func() (*searchCommitsResponse, error)
49+
// unshallowGitRepository expands shallow clones before computing the final missing commits.
50+
unshallowGitRepository func() (bool, error)
51+
// sendObjectsPackFile uploads a git packfile containing the missing commits.
52+
sendObjectsPackFile func(string, []string, []string) (int64, error)
53+
}
4254
)
4355

4456
var (
@@ -51,6 +63,9 @@ var (
5163
// additionalFeaturesInitializationMu serializes additional feature initialization with test-only state resets.
5264
additionalFeaturesInitializationMu sync.Mutex
5365

66+
// repositoryUploadHooksMu protects repository upload hooks that tests replace while settings upload work may still be running.
67+
repositoryUploadHooksMu sync.RWMutex
68+
5469
// ciVisibilityRapidClient contains the http rapid client to do CI Visibility queries and upload to the rapid backend
5570
ciVisibilityClient net.Client
5671

@@ -75,8 +90,8 @@ var (
7590
// ciVisibilityImpactedTestsAnalyzer contains the CI Visibility impacted tests analyzer
7691
ciVisibilityImpactedTestsAnalyzer *impactedtests.ImpactedTestAnalyzer
7792

78-
// uploadRepositoryChangesFunc is a must-not-call test seam used to prove offline/file modes suppress git upload.
79-
uploadRepositoryChangesFunc = uploadRepositoryChanges
93+
// uploadRepositoryChangesFunc is a must-not-call test seam used to prove offline/file modes suppress git upload. A nil value uses the default git upload path.
94+
uploadRepositoryChangesFunc func() (int64, error)
8095

8196
// getSearchCommitsFunc allows tests to exercise repository upload control flow without reading local git state.
8297
getSearchCommitsFunc = getSearchCommits
@@ -107,12 +122,14 @@ func ensureSettingsInitialization(serviceName string) {
107122
log.Debug("civisibility: settings initialization mode [manifest:%t payload_files:%t manifest_file:%s payload_root:%s repository_upload_enabled:%t]",
108123
testOptimizationMode.ManifestEnabled, testOptimizationMode.PayloadFilesEnabled, bazel.TestOptimizationPathForLog(testOptimizationMode.ManifestPath), testOptimizationMode.PayloadsRoot, uploadEnabled)
109124
if uploadEnabled {
125+
repositoryUpload := snapshotRepositoryUploadHooks()
126+
110127
// upload the repository changes
111128
go func() {
112129
defer func() {
113130
close(uploadChannel)
114131
}()
115-
bytes, err := uploadRepositoryChangesFunc()
132+
bytes, err := repositoryUpload.run()
116133
if err != nil {
117134
log.Error("civisibility: error uploading repository changes: %s", err.Error())
118135
} else {
@@ -404,10 +421,36 @@ func GetImpactedTestsAnalyzer() *impactedtests.ImpactedTestAnalyzer {
404421
return ciVisibilityImpactedTestsAnalyzer
405422
}
406423

424+
// snapshotRepositoryUploadHooks returns a stable hook set so asynchronous upload work is isolated from test resets.
425+
func snapshotRepositoryUploadHooks() repositoryUploadHooks {
426+
repositoryUploadHooksMu.RLock()
427+
defer repositoryUploadHooksMu.RUnlock()
428+
429+
return repositoryUploadHooks{
430+
uploadRepositoryChanges: uploadRepositoryChangesFunc,
431+
getSearchCommits: getSearchCommitsFunc,
432+
unshallowGitRepository: unshallowGitRepositoryFunc,
433+
sendObjectsPackFile: sendObjectsPackFileFunc,
434+
}
435+
}
436+
437+
// run executes either a full upload replacement hook or the default upload path with the captured hooks.
438+
func (hooks repositoryUploadHooks) run() (int64, error) {
439+
if hooks.uploadRepositoryChanges != nil {
440+
return hooks.uploadRepositoryChanges()
441+
}
442+
return uploadRepositoryChangesWithHooks(hooks)
443+
}
444+
407445
// uploadRepositoryChanges discovers the commits and packfiles that must be uploaded so backend features can reason about the current repo state.
408446
func uploadRepositoryChanges() (bytes int64, err error) {
447+
return uploadRepositoryChangesWithHooks(snapshotRepositoryUploadHooks())
448+
}
449+
450+
// uploadRepositoryChangesWithHooks runs the default git upload flow against a stable hook snapshot.
451+
func uploadRepositoryChangesWithHooks(hooks repositoryUploadHooks) (bytes int64, err error) {
409452
// get the search commits response
410-
initialCommitData, err := getSearchCommitsFunc()
453+
initialCommitData, err := hooks.getSearchCommits()
411454
if err != nil {
412455
return 0, fmt.Errorf("civisibility: error getting the search commits response: %s", err)
413456
}
@@ -436,7 +479,7 @@ func uploadRepositoryChanges() (bytes int64, err error) {
436479
}
437480

438481
// there's some missing commits on the backend, first we need to check if we need to unshallow before sending anything...
439-
hasBeenUnshallowed, err := unshallowGitRepositoryFunc()
482+
hasBeenUnshallowed, err := hooks.unshallowGitRepository()
440483
if err != nil || !hasBeenUnshallowed {
441484
if err != nil {
442485
log.Warn("%s", err.Error())
@@ -445,11 +488,11 @@ func uploadRepositoryChanges() (bytes int64, err error) {
445488
// the initial commit data
446489

447490
// send the pack file with the missing commits
448-
return sendObjectsPackFileFunc(initialCommitData.LocalCommits[0], initialMissingCommits, initialCommitData.RemoteCommits)
491+
return hooks.sendObjectsPackFile(initialCommitData.LocalCommits[0], initialMissingCommits, initialCommitData.RemoteCommits)
449492
}
450493

451494
// after unshallowing the repository we need to get the search commits to calculate the missing commits again
452-
commitsData, err := getSearchCommitsFunc()
495+
commitsData, err := hooks.getSearchCommits()
453496
if err != nil {
454497
return 0, fmt.Errorf("civisibility: error getting the search commits response: %s", err)
455498
}
@@ -460,7 +503,7 @@ func uploadRepositoryChanges() (bytes int64, err error) {
460503
}
461504

462505
// send the pack file with the missing commits
463-
return sendObjectsPackFileFunc(commitsData.LocalCommits[0], commitsData.missingCommits(), commitsData.RemoteCommits)
506+
return hooks.sendObjectsPackFile(commitsData.LocalCommits[0], commitsData.missingCommits(), commitsData.RemoteCommits)
464507
}
465508

466509
// getSearchCommits gets the search commits response with the local and remote commits

‎internal/civisibility/integrations/civisibility_features_test.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,48 @@ func TestUploadRepositoryChangesReusesInitialMissingCommitsWhenUnshallowIsUnavai
103103
assert.Equal(t, []string{"remote-1"}, uploadedExcludes)
104104
}
105105

106+
func TestUploadRepositoryChangesUsesHookSnapshot(t *testing.T) {
107+
resetCIVisibilityStateForTesting()
108+
t.Cleanup(resetCIVisibilityStateForTesting)
109+
110+
uploadStarted := make(chan struct{})
111+
uploadRelease := make(chan struct{})
112+
getSearchCalls := 0
113+
114+
getSearchCommitsFunc = func() (*searchCommitsResponse, error) {
115+
getSearchCalls++
116+
if getSearchCalls == 1 {
117+
close(uploadStarted)
118+
<-uploadRelease
119+
return newSearchCommitsResponse([]string{"local-1"}, nil, true), nil
120+
}
121+
return newSearchCommitsResponse([]string{"local-1"}, []string{"local-1"}, true), nil
122+
}
123+
unshallowGitRepositoryFunc = func() (bool, error) {
124+
return true, nil
125+
}
126+
sendObjectsPackFileFunc = func(_ string, _ []string, _ []string) (int64, error) {
127+
return 42, nil
128+
}
129+
130+
done := make(chan struct{})
131+
var bytes int64
132+
var err error
133+
go func() {
134+
defer close(done)
135+
bytes, err = uploadRepositoryChanges()
136+
}()
137+
138+
<-uploadStarted
139+
resetCIVisibilityStateForTesting()
140+
close(uploadRelease)
141+
<-done
142+
143+
require.NoError(t, err)
144+
assert.Equal(t, int64(42), bytes)
145+
assert.Equal(t, 2, getSearchCalls)
146+
}
147+
106148
func TestEnsureSettingsInitializationNilClientFactoryDoesNotStartUpload(t *testing.T) {
107149
resetCIVisibilityStateForTesting()
108150
t.Cleanup(resetCIVisibilityStateForTesting)

0 commit comments

Comments
 (0)