.

go/xcommon/tracing/cmd/gotracegen/gotracegen.go

@@ -26,14 +26,18 @@ XXX tracepoints this package imports
 package main
 
 import (
+	"bufio"
+	"bytes"
 	"flag"
 	"fmt"
 	"go/ast"
 	"go/parser"
 	"go/token"
 	"go/types"
+	"io/ioutil"
 	"log"
 	"os"
+	"path/filepath"
 	"sort"
 	"strings"
 	"text/template"
@@ -68,7 +72,7 @@ type traceImport struct {
 
 // Package represents tracing-related information about a package
 type Package struct {
-	PkgPath string
+	Pkgi	*loader.PackageInfo
 	Eventv  []*traceEvent  // trace events this package defines
 	Importv []*traceImport // packages this package trace imports
 }
@@ -173,7 +177,7 @@ func packageTrace(lprog *loader.Program, pkgi *loader.PackageInfo) *Package {
 	sort.Sort(byEventName(eventv))
 	sort.Sort(byPkgPath(importv))
 
-	return &Package{PkgPath: pkgi.Pkg.Path(), Eventv: eventv, Importv: importv}
+	return &Package{Pkgi: pkgi, Eventv: eventv, Importv: importv}
 }
 
 // ----------------------------------------
@@ -260,12 +264,43 @@ func {{.Pkgi.Pkg.Name}}_{{.Name}}_Attach(*tracing.ProbeGroup, func({{.TypedArgv}
 // XXX go123 in magic
 const magic = "// Code generated by lab.nexedi.com/kirr/go123/tracing/cmd/gotrace; DO NOT EDIT.\n"
 
-struct Buffer {
+// checkCanWrite check whether it is safe to write at path
+// it is safe to write when either
+// - the file does not exist, or
+// - it exits but was previously generated by us
+func checkCanWrite(path string) error {
+	f, err := os.Open(path)
+	if e, ok := err.(*os.PathError); ok && os.IsNotExist(e.Err) {
+		return nil
+	}
+
+	defer f.Close()
+	bf := bufio.NewReader(f)
+
+	headline, err := bf.ReadString('\n')
+	if err != nil || headline != magic {
+		return fmt.Errorf("refusing to make output: %v exists but was not generated by gotrace", path)
+	}
+
+	return nil
+}
+
+// writeFile writes data to a file at path after checking it is safe to write there
+func writeFile(path string, data []byte) error {
+	err := checkCanWrite(path)
+	if err != nil {
+		return err
+	}
+
+	return ioutil.WriteFile(path, data, 0666)
+}
+
+type Buffer struct {
 	bytes.Buffer
 }
 
 func (b *Buffer) emit(format string, argv ...interface{}) {
-	fmt.Fprintf(b, format + "\n", ...argv)
+	fmt.Fprintf(b, format + "\n", argv...)
 }
 
 // tracegen generates code according to tracing directives in a package @ pkgpath
@@ -285,22 +320,34 @@ func tracegen(pkgpath string) error {
 		log.Fatal(err)
 	}
 
-	// tracing info for this specified package
 	pkgi := lprog.InitialPackages()[0]
+
+	// determine package directory
+	if len(pkgi.Files) == 0 {
+		log.Fatalf("package %s is empty", pkgi.Pkg.Path)
+	}
+
+	pkgdir := filepath.Dir(lprog.Fset.File(pkgi.Files[0].Pos()).Name())
+	//println("pkgpath", pkgpath)
+	//println("pkgdir", pkgdir)
+	//return nil
+
+	// tracing info for this specified package
 	pkg := packageTrace(lprog, pkgi)
 
-	buf := &bytes.Buffer{}
+	buf := &Buffer{}
 
 	// prologue
-	buf.Write(magic)
-	buf.emit("\npackage %v", pkg.Pkg.Name)
+	buf.WriteString(magic)
+	buf.emit("\npackage %v", pkg.Pkgi.Pkg.Name)
 	buf.emit("\nimport (")
 	buf.emit("\t\tlab.nexedi.com/kirr/neo/go/xcommon/tracing")
+	// TODO import all packages for used types
+	buf.emit(")")
 
-
-	// generate code for trace:event definitions
+	// code for trace:event definitions
 	for _, event := range pkg.Eventv {
-		err = traceEventCodeTmpl.Execute(os.Stdout, event)
+		err = traceEventCodeTmpl.Execute(buf, event)
 		if err != nil {
 			panic(err)	// XXX
 		}
@@ -308,10 +355,10 @@ func tracegen(pkgpath string) error {
 
 	// TODO export hash
 
-	// generate code for trace:import imports
-	fmt.Println()
+	// code for trace:import imports
+	buf.emit("")
 	for _, timport := range pkg.Importv {
-		fmt.Printf("// traceimport: %v\n", timport.PkgPath)
+		buf.emit("// traceimport: %v", timport.PkgPath)
 
 		pkgi = lprog.Package(timport.PkgPath)
 		if pkgi == nil {
@@ -322,7 +369,7 @@ func tracegen(pkgpath string) error {
 		pkg = packageTrace(lprog, pkgi)
 
 		for _, event := range pkg.Eventv {
-			err = traceEventImportTmpl.Execute(os.Stdout, event)
+			err = traceEventImportTmpl.Execute(buf, event)
 			if err != nil {
 				panic(err)	// XXX
 			}
@@ -331,12 +378,30 @@ func tracegen(pkgpath string) error {
 
 	// TODO check export hash
 
-	// TODO trace.s with "// empty .s so `go build` does not use -complete for go:linkname to work"
+
+	// before flushing check we can write to trace.{go,s}
+	err = writeFile(filepath.Join(pkgdir, "trace.go"), buf.Bytes())
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// write trace.s so go:linkname works
+	buf.Reset()
+	buf.WriteString(magic)
+	buf.emit("// empty .s so `go build` does not use -complete for go:linkname to work")
+
+	err = writeFile(filepath.Join(pkgdir, "trace.s"), buf.Bytes())
+	if err != nil {
+		log.Fatal(err)
+	}
 
 	return nil	// XXX
 }
 
 func main() {
+	log.SetFlags(0)
+	log.SetPrefix("gotrace: ")
+
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr,
 `gotracegen [options] [package]