goでファイルを一時的に書き換えた処理をSIGINT時にもrollbackしたい場合のメモ

はじめに

goでhttptrace付きでgo runを実行するものを作っていた。詳しい話はアドベントカレンダーに記事を書いた。PoC的なものだったのでまじめには作っていなかったのだけれど意外と便利なので常用してみようとした。そのときに１つだけ困ったことがあったのでそれへの対応のメモ。

問題

現状の処理では以下の様にファイルを一時的に書き換えている。

main.goをAST越しにいじった結果に変更
変更後のmain.goでgo runを実行
main.goを復元(rollback)

ここで、内部のgo runで実行した長めの処理の途中で中断したいことがある。一方で内部で利用しているgo runは単にexec.CommandのRunを呼んでいるだけなのでCtrl + cなどでsignal(SIGINT)が送られたら即座に終了してしまう(挙動としては正しい)。

rollbackは単にdeferで付けているだけなので当然ファイルは書き換わった状態のまま残る。

defer rollback()
doSomething() // <- この処理中でSIGINTならrollbackが呼ばれない

signalを受け取る方法

os/signalを使う。signal.Notifyでdefaultのhandlerから挙動を変えている(defaultはpanic)。あと、signalの制御を完全には奪い取りたくないのでしっかり終了しておけるようにしたい(kill -9 <pid>が必要になる自体は避けたい)。これにはsignal.Stopを使えば良さそう。

そんなわけで以下のようなサンプルコードを作ってみる。

defaultの状態は0
実行前に1に変更 (ファイルの変更に対応)
rollbackすると-1される

正常な遷移なら 0 -> 1 -> 0 となり0で終了されるはず。加えて、SIGINTが送られてもrollbackされるようにする。

実行結果だけを先に転記しておく。上手く動いていそうに見える。

# signalが送られていない場合
$ go run main.go
2018/12/14 07:50:51 start 0
2018/12/14 07:50:51 inc 1
2018/12/14 07:50:53 rollback 1
2018/12/14 07:50:53 end 0

# signalを２度送った場合（途中で中断したい場合)
$ go run main.go
2018/12/14 07:51:12 start 0
2018/12/14 07:51:12 inc 1
  C-c C-c2018/12/14 07:51:12 rollback 1
  C-c C-csignal: interrupt

コード

package main

import (
    "log"
    "os"
    "os/signal"
    "syscall"
    "time"
)

func main() {
    var i int

    log.Println("start", i)
    i++
    log.Println("inc", i)

    defer func() {
        log.Println("end", i)
    }()
    rollback := func() {
        log.Println("rollback", i)
        i--
    }

    c := make(chan os.Signal)
    signal.Notify(c, syscall.SIGINT)
    go func() {
        <-c
        rollback()
        signal.Stop(c)
    }()

    time.Sleep(2 * time.Second)
    rollback()
}

rollbackが複数回呼ばれてしまうことがある

ちょっとした問題がある。rollbackが複数回呼ばれることがある模様。Ctrl + cを一度だけ入力させrollbackさせたのちに何もせずに正常終了という形になった場合にはrollbackが2度呼ばれてしまう。

$ go run main.go
2018/12/14 07:50:56 start 0
2018/12/14 07:50:56 inc 1
  C-c C-c2018/12/14 07:50:57 rollback 1
2018/12/14 07:50:58 rollback 0
2018/12/14 07:50:58 end -1

これはsync.Onceで雑に対応してしまえば良さそう。

--- 01tough-sig/main.go  2018-12-14 06:43:30.756478993 +0900
+++ 02once/main.go    2018-12-14 06:43:29.833204660 +0900
@@ -4,6 +4,7 @@
    "log"
    "os"
    "os/signal"
+   "sync"
    "syscall"
    "time"
 )
@@ -18,9 +19,12 @@
    defer func() {
        log.Println("end", i)
    }()
+   var once sync.Once
    rollback := func() {
-      log.Println("rollback", i)
-      i--
+       once.Do(func() {
+           log.Println("rollback", i)
+           i--
+       })
    }
 
    c := make(chan os.Signal)

今度は大丈夫。

$ go run main.go
2018/12/14 07:58:44 start 0
2018/12/14 07:58:44 inc 1
  C-c C-c2018/12/14 07:58:45 rollback 1
2018/12/14 07:58:46 end 0

おわりに

無事に本体自体も修正された。

$ go-run-httptrace 03sleep/main.go
2018/12/14 08:00:02 parse 03sleep/main.go
2018/12/14 08:00:02 transform AST
2018/12/14 08:00:02 format
running via github.com/podhmo/go-traceable/cmd/go-run-httptrace/

# テキトウな実行コード
sleep 2
# 途中でSIGINT
  C-c C-c2018/12/14 08:00:03 rollback 03sleep/main.go

signal: interrupt
2018/12/14 08:00:03 exit status 1

ちなみに実行中は以下のコードが

package main

import (
    "fmt"
    "time"
)

func main() {
    fmt.Println("sleep", 2)
    time.Sleep(2 * time.Second)
    fmt.Println("end")
}

以下に変更されている(今まではこれが残っていた)。

package main

import (
    "fmt"
    "time"

    httptrace "github.com/podhmo/go-traceable/httptrace"
)

func mainInner() {
    fmt.Println("sleep", 2)
    time.Sleep(2 * time.Second)
    fmt.Println("end")
}
func main() {
    fmt.Println("running via github.com/podhmo/go-traceable/cmd/go-run-httptrace/")
    teardown := httptrace.Patch()
    defer teardown()
    mainInner()
}