register a callback that is called when watchdog forces GC (#15)

notification.go

@@ -0,0 +1,80 @@
+package watchdog
+
+import "sync"
+
+var (
+	gcNotifeeMutex sync.Mutex
+	gcNotifees     []notifeeEntry
+
+	forcedGCNotifeeMutex sync.Mutex
+	forcedGCNotifees     []notifeeEntry
+)
+
+// RegisterPostGCNotifee registers a function that is called every time a GC has happened,
+// both GC runs triggered by the Go runtime and by watchdog.
+// The unregister function returned can be used to unregister this notifee.
+func RegisterPostGCNotifee(f func()) (unregister func()) {
+	gcNotifeeMutex.Lock()
+	defer gcNotifeeMutex.Unlock()
+
+	var id int
+	if len(gcNotifees) > 0 {
+		id = gcNotifees[len(gcNotifees)-1].id + 1
+	}
+	gcNotifees = append(gcNotifees, notifeeEntry{id: id, f: f})
+
+	return func() {
+		gcNotifeeMutex.Lock()
+		defer gcNotifeeMutex.Unlock()
+
+		for i, entry := range gcNotifees {
+			if entry.id == id {
+				gcNotifees = append(gcNotifees[:i], gcNotifees[i+1:]...)
+			}
+		}
+	}
+}
+
+func notifyGC() {
+	if NotifyGC != nil {
+		NotifyGC()
+	}
+	gcNotifeeMutex.Lock()
+	defer gcNotifeeMutex.Unlock()
+	for _, entry := range gcNotifees {
+		entry.f()
+	}
+}
+
+// RegisterPreGCNotifee registers a function that is called before watchdog triggers a GC run.
+// It is ONLY called when watchdog triggers a GC run, not when the Go runtime triggers it.
+// The unregister function returned can be used to unregister this notifee.
+func RegisterPreGCNotifee(f func()) (unregister func()) {
+	forcedGCNotifeeMutex.Lock()
+	defer forcedGCNotifeeMutex.Unlock()
+
+	var id int
+	if len(forcedGCNotifees) > 0 {
+		id = forcedGCNotifees[len(forcedGCNotifees)-1].id + 1
+	}
+	forcedGCNotifees = append(forcedGCNotifees, notifeeEntry{id: id, f: f})
+
+	return func() {
+		forcedGCNotifeeMutex.Lock()
+		defer forcedGCNotifeeMutex.Unlock()
+
+		for i, entry := range forcedGCNotifees {
+			if entry.id == id {
+				forcedGCNotifees = append(forcedGCNotifees[:i], forcedGCNotifees[i+1:]...)
+			}
+		}
+	}
+}
+
+func notifyForcedGC() {
+	forcedGCNotifeeMutex.Lock()
+	defer forcedGCNotifeeMutex.Unlock()
+	for _, entry := range forcedGCNotifees {
+		entry.f()
+	}
+}

watchdog.go

@@ -37,7 +37,7 @@ var (
 	Clock = clock.New()
 
 	// NotifyGC, if non-nil, will be called when a GC has happened.
-	// Deprecated: use RegisterNotifee instead.
+	// Deprecated: use RegisterPostGCNotifee instead.
 	NotifyGC func() = func() {}
 
 	// HeapProfileThreshold sets the utilization threshold that will trigger a
@@ -99,9 +99,6 @@ var (
 	// See: https://github.com/prometheus/client_golang/issues/403
 	memstatsFn = runtime.ReadMemStats
 	sysmemFn   = (*gosigar.Mem).Get
-
-	notifeeMutex sync.Mutex
-	notifees     []notifeeEntry
 )
 
 type notifeeEntry struct {
@@ -365,16 +362,19 @@ func pollingWatchdog(policy Policy, frequency time.Duration, limit uint64, usage
 func forceGC(memstats *runtime.MemStats) {
 	Logger.Infof("watchdog is forcing GC")
 
+	startNotify := time.Now()
+	notifyForcedGC()
 	// it's safe to assume that the finalizer will attempt to run before
 	// runtime.GC() returns because runtime.GC() waits for the sweep phase to
 	// finish before returning.
 	// finalizers are run in the sweep phase.
 	start := time.Now()
+	notificationsTook := start.Sub(startNotify)
 	runtime.GC()
 	took := time.Since(start)
 
 	memstatsFn(memstats)
-	Logger.Infof("watchdog-triggered GC finished; took: %s; current heap allocated: %d bytes", took, memstats.HeapAlloc)
+	Logger.Infof("watchdog-triggered GC finished; notifications took: %s, took: %s; current heap allocated: %d bytes", notificationsTook, took, memstats.HeapAlloc)
 }
 
 func setupGCSentinel(gcTriggered chan struct{}) {
@@ -510,38 +510,3 @@ func wdrecover() {
 		}
 	}
 }
-
-// RegisterNotifee registers a function that will be called when a GC has happened.
-// The unregister function returned can be used to unregister this notifee.
-func RegisterNotifee(f func()) (unregister func()) {
-	notifeeMutex.Lock()
-	defer notifeeMutex.Unlock()
-
-	var id int
-	if len(notifees) > 0 {
-		id = notifees[len(notifees)-1].id + 1
-	}
-	notifees = append(notifees, notifeeEntry{id: id, f: f})
-
-	return func() {
-		notifeeMutex.Lock()
-		defer notifeeMutex.Unlock()
-
-		for i, entry := range notifees {
-			if entry.id == id {
-				notifees = append(notifees[:i], notifees[i+1:]...)
-			}
-		}
-	}
-}
-
-func notifyGC() {
-	if NotifyGC != nil {
-		NotifyGC()
-	}
-	notifeeMutex.Lock()
-	defer notifeeMutex.Unlock()
-	for _, entry := range notifees {
-		entry.f()
-	}
-}

watchdog_test.go

@@ -153,7 +153,7 @@ func TestSystemDriven_Isolated(t *testing.T) {
 	NotifyGC = func() {
 		notifyChDeprecated <- struct{}{}
 	}
-	unregister := RegisterNotifee(func() {
+	unregister := RegisterPostGCNotifee(func() {
 		notifyCh <- struct{}{}
 	})
 	defer unregister()