From b702f127ed340010a8730fae8e2ddfc872d33c77 Mon Sep 17 00:00:00 2001
From: Tim Voronov <ziflex@users.noreply.github.com>
Date: Tue, 26 Feb 2019 15:32:50 -0500
Subject: [PATCH] Fixed graceful process termination (#240)

---
 cli/exec.go            | 13 ++-----------
 cli/options.go         | 26 +++++++++++++++-----------
 cli/repl.go            | 16 +++-------------
 pkg/runtime/program.go |  6 ++++--
 4 files changed, 24 insertions(+), 37 deletions(-)

diff --git a/cli/exec.go b/cli/exec.go
index dac8c0ec..14ffa970 100644
--- a/cli/exec.go
+++ b/cli/exec.go
@@ -9,7 +9,6 @@ import (
 	"io/ioutil"
 	"os"
 	"os/signal"
-	"syscall"
 )
 
 func ExecFile(pathToFile string, opts Options) {
@@ -38,18 +37,10 @@ func Exec(query string, opts Options) {
 
 	l := NewLogger()
 
-	ctx, err := opts.WithContext(context.Background())
+	ctx, cancel := opts.WithContext(context.Background())
 
-	if err != nil {
-		fmt.Println("Failed to register HTML drivers")
-		fmt.Println(err)
-		os.Exit(1)
-		return
-	}
-
-	ctx, cancel := context.WithCancel(ctx)
 	c := make(chan os.Signal, 1)
-	signal.Notify(c, syscall.SIGHUP)
+	signal.Notify(c, os.Interrupt)
 
 	go func() {
 		for {
diff --git a/cli/options.go b/cli/options.go
index 3b3a8f86..83f55307 100644
--- a/cli/options.go
+++ b/cli/options.go
@@ -16,24 +16,28 @@ type Options struct {
 	ShowTime  bool
 }
 
-func (opts Options) WithContext(ctx context.Context) (context.Context, error) {
+func (opts Options) WithContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	httpDriver := http.NewDriver(
+		http.WithProxy(opts.Proxy),
+		http.WithUserAgent(opts.UserAgent),
+	)
+
 	ctx = drivers.WithContext(
 		ctx,
-		http.NewDriver(
-			http.WithProxy(opts.Proxy),
-			http.WithUserAgent(opts.UserAgent),
-		),
+		httpDriver,
 		drivers.AsDefault(),
 	)
 
+	cdpDriver := cdp.NewDriver(
+		cdp.WithAddress(opts.Cdp),
+		cdp.WithProxy(opts.Proxy),
+		cdp.WithUserAgent(opts.UserAgent),
+	)
+
 	ctx = drivers.WithContext(
 		ctx,
-		cdp.NewDriver(
-			cdp.WithAddress(opts.Cdp),
-			cdp.WithProxy(opts.Proxy),
-			cdp.WithUserAgent(opts.UserAgent),
-		),
+		cdpDriver,
 	)
 
-	return ctx, nil
+	return context.WithCancel(ctx)
 }
diff --git a/cli/repl.go b/cli/repl.go
index da188b16..8a7e8a94 100644
--- a/cli/repl.go
+++ b/cli/repl.go
@@ -3,12 +3,10 @@ package cli
 import (
 	"context"
 	"fmt"
+	"github.com/MontFerret/ferret/pkg/parser/fql"
 	"os"
 	"os/signal"
 	"strings"
-	"syscall"
-
-	"github.com/MontFerret/ferret/pkg/parser/fql"
 
 	"github.com/MontFerret/ferret/pkg/compiler"
 	"github.com/MontFerret/ferret/pkg/runtime"
@@ -50,18 +48,10 @@ func Repl(version string, opts Options) {
 
 	l := NewLogger()
 
-	ctx, err := opts.WithContext(context.Background())
-
-	if err != nil {
-		fmt.Println("Failed to register HTML drivers")
-		fmt.Println(err)
-		os.Exit(1)
-		return
-	}
+	ctx, cancel := opts.WithContext(context.Background())
 
-	ctx, cancel := context.WithCancel(ctx)
 	c := make(chan os.Signal, 1)
-	signal.Notify(c, syscall.SIGHUP)
+	signal.Notify(c, os.Interrupt)
 
 	exit := func() {
 		cancel()
diff --git a/pkg/runtime/program.go b/pkg/runtime/program.go
index a406fc8e..d51a46c0 100644
--- a/pkg/runtime/program.go
+++ b/pkg/runtime/program.go
@@ -2,11 +2,12 @@ package runtime
 
 import (
 	"context"
+	"runtime"
+
 	"github.com/MontFerret/ferret/pkg/runtime/core"
 	"github.com/MontFerret/ferret/pkg/runtime/logging"
 	"github.com/MontFerret/ferret/pkg/runtime/values"
 	"github.com/pkg/errors"
-	"runtime"
 )
 
 type Program struct {
@@ -61,12 +62,13 @@ func (p *Program) Run(ctx context.Context, setters ...Option) (result []byte, er
 	}()
 
 	scope, closeFn := core.NewRootScope()
+
 	defer func() {
 		if err := closeFn(); err != nil {
 			logger.Error().
 				Timestamp().
 				Err(err).
-				Msg("Closing root scope")
+				Msg("closing root scope")
 		}
 	}()