Merge pull request #9 from suker200/support_watch_preemptible_events

Support watch preemptible events
suker200 · Dec 6, 2021 · c0e32d2 · c0e32d2
2 parents 662bbec + 48daf75
commit c0e32d2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/main.go b/main.go
@@ -116,6 +116,10 @@ func main() {
 		resources.eventVMWatching()
 	}()
 
+	go func() {
+		resources.eventPrempVMWatching()
+	}()
+
 	resources.process(pusherCfg)
 
 	for {

diff --git a/utils.go b/utils.go
@@ -133,8 +133,29 @@ func (r *resources) eventVMWatching() {
 				if err := myProcess.Signal(syscall.SIGTERM); err != nil {
 					fmt.Println(err)
 				}
+				break
 			}
 			time.Sleep(time.Duration(2) * time.Second)
 		}
 	}
 }
+
+func (r *resources) eventPrempVMWatching() {
+	httpClient, _ := NewRetryHttpClient("", "")
+	pid := os.Getpid()
+
+	myProcess, _ := os.FindProcess(pid)
+
+	for {
+		if resp, err := httpClient.Do(context.Background(), nil, "http://metadata.google.internal/computeMetadata/v1/instance/preempted", map[string]string{"Metadata-Flavor": "Google"}, "GET"); err == nil {
+			if string(resp) != "FALSE" {
+				if err := myProcess.Signal(syscall.SIGTERM); err != nil {
+					fmt.Println(err)
+				}
+				break
+			}
+			time.Sleep(time.Duration(2) * time.Second)
+		}
+	}
+}
+