-
-
Notifications
You must be signed in to change notification settings - Fork 3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #5872 from ipfs/misc/move-intr-handler
main: move InterruptHandler to util
- Loading branch information
Showing
3 changed files
with
97 additions
and
69 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
// +build !wasm | ||
|
||
package util | ||
|
||
import ( | ||
"context" | ||
"fmt" | ||
"io" | ||
"os" | ||
"os/signal" | ||
"sync" | ||
"syscall" | ||
) | ||
|
||
// IntrHandler helps set up an interrupt handler that can | ||
// be cleanly shut down through the io.Closer interface. | ||
type IntrHandler struct { | ||
sig chan os.Signal | ||
wg sync.WaitGroup | ||
} | ||
|
||
func NewIntrHandler() *IntrHandler { | ||
ih := &IntrHandler{} | ||
ih.sig = make(chan os.Signal, 1) | ||
return ih | ||
} | ||
|
||
func (ih *IntrHandler) Close() error { | ||
close(ih.sig) | ||
ih.wg.Wait() | ||
return nil | ||
} | ||
|
||
// Handle starts handling the given signals, and will call the handler | ||
// callback function each time a signal is catched. The function is passed | ||
// the number of times the handler has been triggered in total, as | ||
// well as the handler itself, so that the handling logic can use the | ||
// handler's wait group to ensure clean shutdown when Close() is called. | ||
func (ih *IntrHandler) Handle(handler func(count int, ih *IntrHandler), sigs ...os.Signal) { | ||
signal.Notify(ih.sig, sigs...) | ||
ih.wg.Add(1) | ||
go func() { | ||
defer ih.wg.Done() | ||
count := 0 | ||
for range ih.sig { | ||
count++ | ||
handler(count, ih) | ||
} | ||
signal.Stop(ih.sig) | ||
}() | ||
} | ||
|
||
func SetupInterruptHandler(ctx context.Context) (io.Closer, context.Context) { | ||
intrh := NewIntrHandler() | ||
ctx, cancelFunc := context.WithCancel(ctx) | ||
|
||
handlerFunc := func(count int, ih *IntrHandler) { | ||
switch count { | ||
case 1: | ||
fmt.Println() // Prevent un-terminated ^C character in terminal | ||
|
||
ih.wg.Add(1) | ||
go func() { | ||
defer ih.wg.Done() | ||
cancelFunc() | ||
}() | ||
|
||
default: | ||
fmt.Println("Received another interrupt before graceful shutdown, terminating...") | ||
os.Exit(-1) | ||
} | ||
} | ||
|
||
intrh.Handle(handlerFunc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM) | ||
|
||
return intrh, ctx | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
package util | ||
|
||
import ( | ||
"context" | ||
"io" | ||
) | ||
|
||
type ctxCloser context.CancelFunc | ||
|
||
func (c ctxCloser) Close() error { | ||
c() | ||
return nil | ||
} | ||
|
||
func SetupInterruptHandler(ctx context.Context) (io.Closer, context.Context) { | ||
ctx, cancel := context.WithCancel(ctx) | ||
return ctxCloser(cancel), ctx | ||
} |