runtime: move exit hooks into internal/runtime/exithookHEADmaster

This removes a //go:linkname usage in the coverage implementation.

For #67401.

Change-Id: I0602172c7e372a84465160dbf46d9fa371582fff
Reviewed-on: https://go-review.googlesource.com/c/go/+/586259
Auto-Submit: Russ Cox <rsc@golang.org>
Reviewed-by: Than McIntosh <thanm@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Cherry Mui <cherryyz@google.com>

2 files changed, 67 insertions, 7 deletions

diff --git a/src/internal/coverage/cfile/hooks.go b/src/internal/coverage/cfile/hooks.go
index 003d6ca1e56..3821d1e91e0 100644
--- a/src/internal/coverage/cfile/hooks.go
+++ b/src/internal/coverage/cfile/hooks.go
@@ -4,7 +4,7 @@
 
 package cfile
 
-import _ "unsafe"
+import "internal/runtime/exithook"
 
 // InitHook is invoked from the main package "init" routine in
 // programs built with "-cover". This function is intended to be
@@ -29,14 +29,10 @@ func InitHook(istest bool) {
 	// Note: hooks are run in reverse registration order, so
 	// register the counter data hook before the meta-data hook
 	// (in the case where two hooks are needed).
-	runOnNonZeroExit := true
-	runtime_addExitHook(emitCounterData, runOnNonZeroExit)
+	exithook.Add(exithook.Hook{F: emitCounterData, RunOnFailure: true})
 	if istest {
-		runtime_addExitHook(emitMetaData, runOnNonZeroExit)
+		exithook.Add(exithook.Hook{F: emitMetaData, RunOnFailure: true})
 	} else {
 		emitMetaData()
 	}
 }
-
-//go:linkname runtime_addExitHook runtime.addExitHook
-func runtime_addExitHook(f func(), runOnNonZeroExit bool)
diff --git a/src/internal/runtime/exithook/hooks.go b/src/internal/runtime/exithook/hooks.go
new file mode 100644
index 00000000000..931154c45d0
--- /dev/null
+++ b/src/internal/runtime/exithook/hooks.go
@@ -0,0 +1,64 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package exithook provides limited support for on-exit cleanup.
+//
+// CAREFUL! The expectation is that Add should only be called
+// from a safe context (e.g. not an error/panic path or signal
+// handler, preemption enabled, allocation allowed, write barriers
+// allowed, etc), and that the exit function F will be invoked under
+// similar circumstances. That is the say, we are expecting that F
+// uses normal / high-level Go code as opposed to one of the more
+// restricted dialects used for the trickier parts of the runtime.
+package exithook
+
+// A Hook is a function to be run at program termination
+// (when someone invokes os.Exit, or when main.main returns).
+// Hooks are run in reverse order of registration:
+// the first hook added is the last one run.
+type Hook struct {
+	F            func() // func to run
+	RunOnFailure bool   // whether to run on non-zero exit code
+}
+
+var (
+	hooks   []Hook
+	running bool
+)
+
+// Add adds a new exit hook.
+func Add(h Hook) {
+	hooks = append(hooks, h)
+}
+
+// Run runs the exit hooks.
+// It returns an error if Run is already running or
+// if one of the hooks panics.
+func Run(code int) (err error) {
+	if running {
+		return exitError("exit hook invoked exit")
+	}
+	running = true
+
+	defer func() {
+		if x := recover(); x != nil {
+			err = exitError("exit hook invoked panic")
+		}
+	}()
+
+	local := hooks
+	hooks = nil
+	for i := len(local) - 1; i >= 0; i-- {
+		h := local[i]
+		if code == 0 || h.RunOnFailure {
+			h.F()
+		}
+	}
+	running = false
+	return nil
+}
+
+type exitError string
+
+func (e exitError) Error() string { return string(e) }