Skip to content

Commit 5d1eb81

Browse files
alt-dimaalt-dima
committed
context propagation
1 parent 35b2735 commit 5d1eb81

6 files changed

Lines changed: 52 additions & 44 deletions

File tree

‎cmd/agent.go‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,15 @@ package cmd
22

33
import (
44
"bufio"
5+
"context"
56
"encoding/json"
67
"fmt"
78
"log"
89
"math/rand"
910
"net/http"
1011
"os"
11-
"os/signal"
1212
"runtime"
1313
"strings"
14-
"syscall"
1514
"time"
1615

1716
"github.com/alt-dima/iacconsole-cli/utils"
@@ -44,25 +43,23 @@ var agentCmd = &cobra.Command{
4443
log.Printf("Agent ID: %s", agentID)
4544
log.Printf("Connecting to: %s", wsURL)
4645

47-
runAgent(wsURL, authHeader, agentID)
46+
runAgent(cmd.Context(), wsURL, authHeader, agentID)
4847
},
4948
}
5049

51-
func runAgent(wsURL, authHeader, agentID string) {
52-
interrupt := make(chan os.Signal, 1)
53-
signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)
54-
50+
func runAgent(ctx context.Context, wsURL, authHeader, agentID string) {
5551
for {
5652
header := http.Header{}
5753
header.Add("Authorization", authHeader)
5854

59-
c, _, err := websocket.DefaultDialer.Dial(wsURL, header)
55+
dialer := websocket.DefaultDialer
56+
c, _, err := dialer.DialContext(ctx, wsURL, header)
6057
if err != nil {
6158
log.Printf("Dial error: %v. Retrying in 5s...", err)
6259
select {
6360
case <-time.After(5 * time.Second):
6461
continue
65-
case <-interrupt:
62+
case <-ctx.Done():
6663
return
6764
}
6865
}
@@ -140,7 +137,7 @@ func runAgent(wsURL, authHeader, agentID string) {
140137
continue
141138
}
142139
// log.Printf("Received command: %+v", cmd)
143-
go executeCommand(c, cmd, autoExecute)
140+
go executeCommand(ctx, c, cmd, autoExecute)
144141
case "ping":
145142
pong := utils.AgentPong{
146143
AgentMessage: utils.AgentMessage{Type: "pong"},
@@ -158,8 +155,12 @@ func runAgent(wsURL, authHeader, agentID string) {
158155
case <-done:
159156
c.Close()
160157
log.Printf("Connection lost. Retrying in 5s...")
161-
time.Sleep(5 * time.Second)
162-
case <-interrupt:
158+
select {
159+
case <-time.After(5 * time.Second):
160+
case <-ctx.Done():
161+
return
162+
}
163+
case <-ctx.Done():
163164
log.Println("Interrupt received, closing connection...")
164165
// Cleanly close the connection by sending a close message and then
165166
// waiting (with timeout) for the server to close the connection.
@@ -177,7 +178,7 @@ func runAgent(wsURL, authHeader, agentID string) {
177178
}
178179
}
179180

180-
func executeCommand(c *websocket.Conn, cmd utils.AgentCommand, autoExecute bool) {
181+
func executeCommand(ctx context.Context, c *websocket.Conn, cmd utils.AgentCommand, autoExecute bool) {
181182
// Format command for display
182183
cmdStr := formatCommandString(cmd)
183184

@@ -231,7 +232,7 @@ func executeCommand(c *websocket.Conn, cmd utils.AgentCommand, autoExecute bool)
231232
state.IacconsoleApiUrl = os.Getenv("IACCONSOLE_API_URL")
232233
state.StateS3Path = "./state"
233234

234-
utils.ExecuteAgentCommand(c, cmd, state)
235+
utils.ExecuteAgentCommand(ctx, c, cmd, state)
235236
}
236237

237238
// formatCommandString creates a human-readable string representation of the command

‎cmd/exec.go‎

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"log"
55
"os"
66
"os/exec"
7-
"os/signal"
87
"path/filepath"
98
"strings"
109
"syscall"
@@ -23,9 +22,6 @@ var execCmd = &cobra.Command{
2322
initConfig()
2423
},
2524
Run: func(cmd *cobra.Command, args []string) {
26-
//Creating signal to be handled and send to the child tofu/terraform
27-
sigs := make(chan os.Signal, 2)
28-
signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
2925
var err error
3026

3127
// Creating Session State and filling with values
@@ -78,10 +74,10 @@ var execCmd = &cobra.Command{
7874
log.Fatalf("Failed to prepare temp directory: %v", err)
7975
}
8076

81-
if err := s.GenerateVarsByDims(); err != nil {
77+
if err := s.GenerateVarsByDims(cmd.Context()); err != nil {
8278
log.Fatalf("Failed to generate vars by dimensions: %v", err)
8379
}
84-
if err := s.GenerateVarsByDimOptional("defaults"); err != nil {
80+
if err := s.GenerateVarsByDimOptional(cmd.Context(), "defaults"); err != nil {
8581
log.Fatalf("Failed to generate optional vars: %v", err)
8682
}
8783
if err := s.GenerateVarsByEnvVars(); err != nil {
@@ -103,11 +99,14 @@ var execCmd = &cobra.Command{
10399
}
104100
cmdToExec := s.GetStringFromViperByOrgOrDefault("cmd_to_exec")
105101

106-
// Starting child and Waiting for it to finish, passing signals to it
102+
// Starting child and Waiting for it to finish
107103
log.Println("excuting: " + cmdToExec + " " + strings.Join(cmdArgs, " "))
108-
execChildCommand := exec.Command(cmdToExec, cmdArgs...)
104+
execChildCommand := exec.CommandContext(cmd.Context(), cmdToExec, cmdArgs...)
109105
execChildCommand.Dir = s.CmdWorkTempDir
110106
execChildCommand.Env = os.Environ()
107+
execChildCommand.Cancel = func() error {
108+
return execChildCommand.Process.Signal(syscall.SIGINT)
109+
}
111110
execChildCommand.Stdin = os.Stdin
112111
execChildCommand.Stdout = os.Stdout
113112
execChildCommand.Stderr = os.Stderr
@@ -116,13 +115,7 @@ var execCmd = &cobra.Command{
116115
log.Fatalf("cmd.Start() failed with %s\n", err)
117116
}
118117

119-
go func() {
120-
sig := <-sigs
121-
log.Println("Got singnal +" + sig.String())
122-
if err := execChildCommand.Process.Signal(sig); err != nil {
123-
log.Printf("Failed to send signal to child process: %v", err)
124-
}
125-
}()
118+
126119

127120
err = execChildCommand.Wait()
128121
exitCodeFinal := 0
@@ -135,7 +128,7 @@ var execCmd = &cobra.Command{
135128
exitCodeFinal = execChildCommand.ProcessState.ExitCode()
136129
}
137130

138-
s.ReportHistory(cmdToExec, cmdArgs, args[0], exitCodeFinal)
131+
s.ReportHistory(cmd.Context(), cmdToExec, cmdArgs, args[0], exitCodeFinal)
139132

140133
if (exitCodeFinal == 0 && (args[0] == "apply" || args[0] == "destroy")) || forceCleanTempDir {
141134
os.RemoveAll(s.CmdWorkTempDir)

‎cmd/root.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cmd
22

33
import (
4+
"context"
45
"fmt"
56
"log"
67
"os"
@@ -20,6 +21,9 @@ var rootCmd = &cobra.Command{
2021
// Uncomment the following line if your bare application
2122
// has an action associated with it:
2223
// Run: func(cmd *cobra.Command, args []string) { },
24+
CompletionOptions: cobra.CompletionOptions{
25+
DisableDefaultCmd: true,
26+
},
2327
}
2428

2529
func SetVersionInfo(version string) {
@@ -29,7 +33,7 @@ func SetVersionInfo(version string) {
2933
// Execute adds all child commands to the root command and sets flags appropriately.
3034
// This is called by main.main(). It only needs to happen once to the rootCmd.
3135
func Execute() {
32-
if err := rootCmd.Execute(); err != nil {
36+
if err := rootCmd.ExecuteContext(context.Background()); err != nil {
3337
fmt.Fprintln(os.Stderr, err)
3438
os.Exit(1)
3539
}

‎utils/agent_executor.go‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,20 +2,22 @@ package utils
22

33
import (
44
"bufio"
5+
"context"
56
"io"
67
"log"
78
"os"
89
"os/exec"
910
"path/filepath"
1011
"strings"
1112
"sync"
13+
"syscall"
1214
"time"
1315

1416
"github.com/gorilla/websocket"
1517
)
1618

1719
// ExecuteAgentCommand runs a command and streams output to WebSocket
18-
func ExecuteAgentCommand(conn *websocket.Conn, cmd AgentCommand, state *State) {
20+
func ExecuteAgentCommand(ctx context.Context, conn *websocket.Conn, cmd AgentCommand, state *State) {
1921
// 1. Prepare environment
2022
state.OrgName = cmd.Org
2123
state.UnitName = cmd.Unit
@@ -78,13 +80,13 @@ func ExecuteAgentCommand(conn *websocket.Conn, cmd AgentCommand, state *State) {
7880
}
7981

8082
// 6. Generate variables - handle errors gracefully
81-
if err := state.GenerateVarsByDims(); err != nil {
83+
if err := state.GenerateVarsByDims(ctx); err != nil {
8284
log.Printf("Error generating vars by dimensions: %v", err)
8385
sendComplete(conn, cmd.ID, 1, err.Error())
8486
return
8587
}
8688

87-
if err := state.GenerateVarsByDimOptional("defaults"); err != nil {
89+
if err := state.GenerateVarsByDimOptional(ctx, "defaults"); err != nil {
8890
log.Printf("Error generating optional vars: %v", err)
8991
sendComplete(conn, cmd.ID, 1, err.Error())
9092
return
@@ -118,9 +120,12 @@ func ExecuteAgentCommand(conn *websocket.Conn, cmd AgentCommand, state *State) {
118120

119121
// 8. Spawn process
120122
log.Printf("Agent executing: %s %s", cmdToExec, strings.Join(args, " "))
121-
child := exec.Command(cmdToExec, args...)
123+
child := exec.CommandContext(ctx, cmdToExec, args...)
122124
child.Dir = state.CmdWorkTempDir
123125
child.Env = os.Environ()
126+
child.Cancel = func() error {
127+
return child.Process.Signal(syscall.SIGINT)
128+
}
124129

125130
stdout, _ := child.StdoutPipe()
126131
stderr, _ := child.StderrPipe()
@@ -150,7 +155,7 @@ func ExecuteAgentCommand(conn *websocket.Conn, cmd AgentCommand, state *State) {
150155
}
151156
}
152157

153-
state.ReportHistory(cmdToExec, args, cmd.Action, exitCode)
158+
state.ReportHistory(ctx, cmdToExec, args, cmd.Action, exitCode)
154159

155160
// 10. Cleanup
156161
if exitCode == 0 && (cmd.Action == "apply" || cmd.Action == "destroy") {

‎utils/externals.go‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ func (s *State) SetupBackendConfig() map[string]interface{} {
6262
return backendConfigMap
6363
}
6464

65-
func (s *State) GetDimData(dimensionKey string, dimensionValue string, skipOnNotFound bool) (map[string]interface{}, error) {
65+
func (s *State) GetDimData(ctx context.Context, dimensionKey string, dimensionValue string, skipOnNotFound bool) (map[string]interface{}, error) {
6666
var dimensionJsonMap map[string]interface{}
6767

6868
if s.IacconsoleApiUrl == "" {
@@ -80,7 +80,11 @@ func (s *State) GetDimData(dimensionKey string, dimensionValue string, skipOnNot
8080
return nil, err
8181
}
8282
} else {
83-
resp, err := http.Get(s.IacconsoleApiUrl + "/v1/dimension/" + s.OrgName + "/" + dimensionKey + "/" + dimensionValue + "?workspace=" + s.Workspace + "&fallbacktomaster=true")
83+
req, err := http.NewRequestWithContext(ctx, "GET", s.IacconsoleApiUrl + "/v1/dimension/" + s.OrgName + "/" + dimensionKey + "/" + dimensionValue + "?workspace=" + s.Workspace + "&fallbacktomaster=true", nil)
84+
if err != nil {
85+
return nil, err
86+
}
87+
resp, err := http.DefaultClient.Do(req)
8488
if err != nil {
8589
return nil, err
8690
}
@@ -125,7 +129,7 @@ func (s *State) GetDimData(dimensionKey string, dimensionValue string, skipOnNot
125129
return dimensionJsonMap, nil
126130
}
127131

128-
func (s *State) ReportHistory(cmdToExec string, cmdArgs []string, cmdMainArg string, exitCode int) {
132+
func (s *State) ReportHistory(ctx context.Context, cmdToExec string, cmdArgs []string, cmdMainArg string, exitCode int) {
129133
if s.IacconsoleApiUrl == "" {
130134
// Only report to API if URL is configured
131135
return
@@ -176,7 +180,7 @@ func (s *State) ReportHistory(cmdToExec string, cmdArgs []string, cmdMainArg str
176180
}
177181

178182
url := fmt.Sprintf("%s/v1/history/%s/%s/%s", s.IacconsoleApiUrl, s.OrgName, workspace, s.UnitName)
179-
req, err := http.NewRequestWithContext(context.TODO(), "POST", url, strings.NewReader(string(payloadBytes)))
183+
req, err := http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(payloadBytes)))
180184
if err != nil {
181185
log.Printf("Failed to create request for history: %v", err)
182186
return

‎utils/generatevars.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
package utils
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"log"
78
"os"
89
"strings"
910
)
1011

11-
func (s *State) GenerateVarsByDims() error {
12+
func (s *State) GenerateVarsByDims(ctx context.Context) error {
1213
for dimKey, dimValue := range s.ParsedDimensions {
13-
dimensionJsonMap, err := s.GetDimData(dimKey, dimValue, false)
14+
dimensionJsonMap, err := s.GetDimData(ctx, dimKey, dimValue, false)
1415
if err != nil {
1516
return err
1617
}
@@ -29,9 +30,9 @@ func (s *State) GenerateVarsByDims() error {
2930
return nil
3031
}
3132

32-
func (s *State) GenerateVarsByDimOptional(optionType string) error {
33+
func (s *State) GenerateVarsByDimOptional(ctx context.Context, optionType string) error {
3334
for dimKey := range s.ParsedDimensions {
34-
dimensionJsonMap, err := s.GetDimData(dimKey, "dim_"+optionType, true)
35+
dimensionJsonMap, err := s.GetDimData(ctx, dimKey, "dim_"+optionType, true)
3536
if err != nil {
3637
return err
3738
}

0 commit comments

Comments
 (0)