talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package cli
6

7
import (
8
	"context"
9
	"fmt"
10
	"os"
11
	"os/signal"
12
	"syscall"
13
)
14

15
// WithContext wraps function call to provide a context cancellable with ^C.
16
func WithContext(ctx context.Context, f func(context.Context) error) error {
17
	wrappedCtx, wrappedCtxCancel := context.WithCancel(ctx)
18
	defer wrappedCtxCancel()
19

20
	// listen for ^C and SIGTERM and abort context
21
	sigCh := make(chan os.Signal, 1)
22
	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
23

24
	exited := make(chan struct{})
25
	defer close(exited)
26

27
	go func() {
28
		select {
29
		case <-sigCh:
30
			wrappedCtxCancel()
31

32
			signal.Stop(sigCh)
33
			fmt.Fprintln(os.Stderr, "Signal received, aborting, press Ctrl+C once again to abort immediately...")
34
		case <-wrappedCtx.Done():
35
			return
36
		case <-exited:
37
		}
38
	}()
39

40
	return f(wrappedCtx)
41
}
42