From a1fadab4b09a3287cb93f0305becfdf6ab2f0985 Mon Sep 17 00:00:00 2001
From: Shiming Zhang <wzshiming@foxmail.com>
Date: Fri, 11 Mar 2022 17:48:46 +0800
Subject: [PATCH] Atomic write status file

---
 pkg/kubelet/nodeshutdown/storage.go | 35 ++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/pkg/kubelet/nodeshutdown/storage.go b/pkg/kubelet/nodeshutdown/storage.go
index 07d7de09a67..fbf1ea90664 100644
--- a/pkg/kubelet/nodeshutdown/storage.go
+++ b/pkg/kubelet/nodeshutdown/storage.go
@@ -18,7 +18,10 @@ package nodeshutdown
 
 import (
 	"encoding/json"
+	"io"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 	"time"
 )
 
@@ -36,7 +39,7 @@ func (l localStorage) Store(data interface{}) (err error) {
 	if err != nil {
 		return err
 	}
-	return os.WriteFile(l.Path, b, 0644)
+	return atomicWrite(l.Path, b, 0644)
 }
 
 func (l localStorage) Load(data interface{}) (err error) {
@@ -61,3 +64,33 @@ type state struct {
 	StartTime time.Time `json:"startTime"`
 	EndTime   time.Time `json:"endTime"`
 }
+
+// atomicWrite atomically writes data to a file specified by filename.
+func atomicWrite(filename string, data []byte, perm os.FileMode) error {
+	f, err := ioutil.TempFile(filepath.Dir(filename), ".tmp-"+filepath.Base(filename))
+	if err != nil {
+		return err
+	}
+	err = os.Chmod(f.Name(), perm)
+	if err != nil {
+		f.Close()
+		return err
+	}
+	n, err := f.Write(data)
+	if err != nil {
+		f.Close()
+		return err
+	}
+	if n < len(data) {
+		f.Close()
+		return io.ErrShortWrite
+	}
+	if err := f.Sync(); err != nil {
+		f.Close()
+		return err
+	}
+	if err := f.Close(); err != nil {
+		return err
+	}
+	return os.Rename(f.Name(), filename)
+}