Merge pull request #15 from gyuho/gofail-go

*: add "gofail-go" for manual "Release" call

code/binding.go

@@ -40,10 +40,12 @@ func (b *Binding) Write(dst io.Writer) error {
 	for _, fp := range b.fps {
 		_, err := fmt.Fprintf(
 			dst,
-			"var %s *runtime.Failpoint = runtime.NewFailpoint(%q, %q)\n",
+			"var %s *runtime.Failpoint = runtime.NewFailpoint(%q, %q, %v)\n",
 			fp.Runtime(),
 			b.fppath,
-			fp.Name())
+			fp.Name(),
+			fp.goFailGo,
+		)
 		if err != nil {
 			return err
 		}

code/failpoint.go

@@ -21,6 +21,8 @@ import (
 )
 
 type Failpoint struct {
+	goFailGo bool
+
 	name    string
 	varType string
 	code    []string
@@ -32,16 +34,22 @@ type Failpoint struct {
 // newFailpoint makes a new failpoint based on the a line containing a
 // failpoint comment header.
 func newFailpoint(l string) (*Failpoint, error) {
-	if !strings.HasPrefix(strings.TrimSpace(l), "// gofail:") {
+	lt := strings.TrimSpace(l)
+	isGoFail, isGoFailGo := strings.HasPrefix(lt, pfxGofail), strings.HasPrefix(lt, pfxGofailGo)
+	if !isGoFail && !isGoFailGo {
 		// not a failpoint
 		return nil, nil
 	}
-	cmd := strings.SplitAfter(l, "// gofail:")[1]
+	pfx := pfxGofail
+	if isGoFailGo {
+		pfx = pfxGofailGo
+	}
+	cmd := strings.SplitAfter(l, pfx)[1]
 	fields := strings.Fields(cmd)
 	if len(fields) != 3 || fields[0] != "var" {
 		return nil, fmt.Errorf("failpoint: malformed comment header %q", l)
 	}
-	return &Failpoint{name: fields[1], varType: fields[2], ws: strings.Split(l, "//")[0]}, nil
+	return &Failpoint{goFailGo: isGoFailGo, name: fields[1], varType: fields[2], ws: strings.Split(l, "//")[0]}, nil
 }
 
 // flush writes the failpoint code to a buffer
@@ -53,8 +61,16 @@ func (fp *Failpoint) flush(dst io.Writer) error {
 }
 
 func (fp *Failpoint) hdr(varname string) string {
-	hdr := fp.ws + "if v" + fp.name + ", __fpErr := " + fp.Runtime() + ".Acquire(); __fpErr == nil { "
-	hdr = hdr + "defer " + fp.Runtime() + ".Release(); "
+	ev := errVarGoFail
+	if fp.goFailGo {
+		ev = errVarGoFailGo
+	}
+	hdr := fp.ws + "if v" + fp.name + fmt.Sprintf(", %s := ", ev) + fp.Runtime() + ".Acquire();" + fmt.Sprintf(" %s == nil { ", ev)
+	exec := "defer "
+	if fp.goFailGo {
+		exec = "go "
+	}
+	hdr = hdr + exec + fp.Runtime() + ".Release(); "
 	if fp.varType == "struct{}" {
 		// unused
 		varname = "_"

code/rewrite.go

@@ -16,11 +16,22 @@ package code
 
 import (
 	"bufio"
+	"fmt"
 	"io"
 	"strings"
 	"unicode"
 )
 
+const (
+	pfxGofail    = `// gofail:`
+	labelGofail  = `/* gofail-label */`
+	errVarGoFail = `__fpErr`
+
+	pfxGofailGo    = `// gofail-go:`
+	labelGofailGo  = `/* gofail-go-label */`
+	errVarGoFailGo = `__fpGoErr`
+)
+
 // ToFailpoints turns all gofail comments into failpoint code. Returns a list of
 // all failpoints it activated.
 func ToFailpoints(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error) {
@@ -51,7 +62,10 @@ func ToFailpoints(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error)
 				fps = append(fps, curfp)
 				curfp = nil
 			}
-		} else if label := gofailLabel(l); label != "" {
+		} else if label := gofailLabel(l, pfxGofail, labelGofail); label != "" {
+			// expose gofail label
+			l = label
+		} else if label := gofailLabel(l, pfxGofailGo, labelGofailGo); label != "" {
 			// expose gofail label
 			l = label
 		} else if curfp, err = newFailpoint(l); err != nil {
@@ -94,12 +108,18 @@ func ToComments(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error) {
 			continue
 		}
 
-		isHdr := strings.Contains(l, ", __fpErr := __fp_") && strings.HasPrefix(lTrim, "if")
+		isErrVarGoFail := strings.Contains(l, fmt.Sprintf(", %s := __fp_", errVarGoFail))
+		isErrVarGoFailGo := strings.Contains(l, fmt.Sprintf(", %s := __fp_", errVarGoFailGo))
+		isHdr := (isErrVarGoFail || isErrVarGoFailGo) && strings.HasPrefix(lTrim, "if")
 		if isHdr {
+			pfx := pfxGofail
+			if isErrVarGoFailGo {
+				pfx = pfxGofailGo
+			}
 			ws = strings.Split(l, "i")[0]
 			n := strings.Split(strings.Split(l, "__fp_")[1], ".")[0]
 			t := strings.Split(strings.Split(l, ".(")[1], ")")[0]
-			dst.WriteString(ws + "// gofail: var " + n + " " + t + "\n")
+			dst.WriteString(ws + pfx + " var " + n + " " + t + "\n")
 			if !strings.Contains(l, "; __badType") {
 				// not single liner
 				unmatchedBraces = 1
@@ -108,8 +128,11 @@ func ToComments(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error) {
 			continue
 		}
 
-		if isLabel := strings.Contains(l, "\t/* gofail-label */"); isLabel {
-			l = strings.Replace(l, "/* gofail-label */", "// gofail:", 1)
+		if isLabel := strings.Contains(l, "\t"+labelGofail); isLabel {
+			l = strings.Replace(l, labelGofail, pfxGofail, 1)
+		}
+		if isLabel := strings.Contains(l, "\t"+labelGofailGo); isLabel {
+			l = strings.Replace(l, labelGofailGo, pfxGofailGo, 1)
 		}
 
 		if _, werr := dst.WriteString(l); werr != nil {
@@ -123,15 +146,15 @@ func ToComments(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error) {
 	return
 }
 
-func gofailLabel(l string) string {
-	if !strings.HasPrefix(strings.TrimSpace(l), "// gofail:") {
+func gofailLabel(l string, pfx string, lb string) string {
+	if !strings.HasPrefix(strings.TrimSpace(l), pfx) {
 		return ""
 	}
-	label := strings.SplitAfter(l, "// gofail:")[1]
+	label := strings.SplitAfter(l, pfx)[1]
 	if len(label) == 0 || !strings.Contains(label, ":") {
 		return ""
 	}
-	return strings.Replace(l, "// gofail:", "/* gofail-label */", 1)
+	return strings.Replace(l, pfx, lb, 1)
 }
 
 func numBraces(l string) (opening int, closing int) {

code/rewrite_test.go

@@ -46,6 +46,18 @@ func f() {
 	}
 }
 `, 1},
+	{`
+	func f() {
+		// gofail: labelTest:
+		for {
+			if g() {
+				// gofail-go: var testLabel struct{}
+				// continue labelTest
+				return
+			}
+		}
+	}
+	`, 1},
 }
 
 func TestToFailpoint(t *testing.T) {

examples/cmd/cmd.go

@@ -21,9 +21,23 @@ import (
 	"github.com/etcd-io/gofail/examples"
 )
 
+/*
+GOFAIL_HTTP=:22381 go run cmd.go
+
+curl -L http://localhost:22381
+
+curl \
+  -L http://localhost:22381/github.com/coreos/gofail/examples/ExampleLabelsGo \
+  -X PUT -d'return'
+
+curl \
+  -L http://localhost:22381/github.com/coreos/gofail/examples/ExampleLabelsGo \
+  -X DELETE
+*/
+
 func main() {
 	for {
-		log.Println(examples.ExampleFunc())
+		log.Println(examples.ExampleLabelsGoFunc())
 		time.Sleep(time.Second)
 	}
 }

examples/examples.go

@@ -39,3 +39,18 @@ func ExampleLabelsFunc() (s string) {
 	}
 	return s
 }
+
+func ExampleLabelsGoFunc() (s string) {
+	i := 0
+	// gofail: myLabel:
+	for i < 5 {
+		s = s + "i"
+		i++
+		for j := 0; j < 5; j++ {
+			s = s + "j"
+			// gofail-go: var ExampleLabelsGo struct{}
+			// continue myLabel
+		}
+	}
+	return s
+}

runtime/failpoint.go

@@ -15,17 +15,28 @@
 package runtime
 
 import (
+	"context"
 	"fmt"
 	"sync"
 )
 
 type Failpoint struct {
+	cmu      sync.RWMutex
+	ctx      context.Context
+	cancel   context.CancelFunc
+	donec    chan struct{}
+	released bool
+
 	mu sync.RWMutex
 	t  *terms
 }
 
-func NewFailpoint(pkg, name string) *Failpoint {
+func NewFailpoint(pkg, name string, goFailGo bool) *Failpoint {
 	fp := &Failpoint{}
+	if goFailGo {
+		fp.ctx, fp.cancel = context.WithCancel(context.Background())
+		fp.donec = make(chan struct{})
+	}
 	register(pkg+"/"+name, fp)
 	return fp
 }
@@ -49,7 +60,22 @@ func (fp *Failpoint) Acquire() (interface{}, error) {
 }
 
 // Release is called when the failpoint exists.
-func (fp *Failpoint) Release() { fp.mu.RUnlock() }
+func (fp *Failpoint) Release() {
+	fp.cmu.RLock()
+	ctx := fp.ctx
+	donec := fp.donec
+	fp.cmu.RUnlock()
+	if ctx != nil && !fp.released {
+		<-ctx.Done()
+		select {
+		case <-donec:
+		default:
+			close(donec)
+		}
+	}
+
+	fp.mu.RUnlock()
+}
 
 // BadType is called when the failpoint evaluates to the wrong type.
 func (fp *Failpoint) BadType(v interface{}, t string) {

runtime/runtime.go

@@ -15,6 +15,7 @@
 package runtime
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -76,6 +77,7 @@ func enableAndLock(failpath, inTerms string) (func(), error) {
 	}
 	fp.mu.Lock()
 	fp.t = t
+	fp.released = false
 	return func() { fp.mu.Unlock() }, nil
 }
 
@@ -87,6 +89,22 @@ func Disable(failpath string) error {
 	if fp == nil {
 		return ErrNoExist
 	}
+
+	fp.cmu.RLock()
+	cancel := fp.cancel
+	donec := fp.donec
+	fp.cmu.RUnlock()
+	if cancel != nil && donec != nil {
+		cancel()
+		<-donec
+
+		fp.cmu.Lock()
+		fp.ctx, fp.cancel = context.WithCancel(context.Background())
+		fp.donec = make(chan struct{})
+		fp.released = true
+		fp.cmu.Unlock()
+	}
+
 	fp.mu.Lock()
 	defer fp.mu.Unlock()
 	if fp.t == nil {