lab.nexedi.com/kirr/go123@v0.0.0-20240207185015-8299741fa871/tracing/cmd/gotrace/gotrace.go (about)

     1  // Copyright (C) 2018-2021  Nexedi SA and Contributors.
     2  //                          Kirill Smelkov <kirr@nexedi.com>
     3  //
     4  // This program is free software: you can Use, Study, Modify and Redistribute
     5  // it under the terms of the GNU General Public License version 3, or (at your
     6  // option) any later version, as published by the Free Software Foundation.
     7  //
     8  // You can also Link and Combine this program with other software covered by
     9  // the terms of any of the Free Software licenses or any of the Open Source
    10  // Initiative approved licenses and Convey the resulting work. Corresponding
    11  // source of such a combination shall include the source code for all other
    12  // software used.
    13  //
    14  // This program is distributed WITHOUT ANY WARRANTY; without even the implied
    15  // warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
    16  //
    17  // See COPYING file for full licensing terms.
    18  // See https://www.nexedi.com/licensing for rationale and options.
    19  
    20  /*
    21  Gotrace is a program to support and interact with go tracing subsystem.
    22  
    23  Gotrace is a common entry to tracing and provides several subcommands:
    24  
    25  	gen	generate code according to tracing annotations and imports
    26  	list	lists tracepoints defined by a package
    27  
    28  See package lab.nexedi.com/kirr/go123/tracing documentation on how to define
    29  and use trace events in programs.
    30  
    31  TODO automatically turn every trace:event into an USDT probe so that they can
    32  be traced from outside of the process too.
    33  See e.g. https://github.com/iovisor/bcc/issues/327 for context.
    34  
    35  FIXME build tags not taken into account
    36  */
    37  package main
    38  
    39  import (
    40  	"bufio"
    41  	"crypto/sha1"
    42  	"flag"
    43  	"fmt"
    44  	"go/ast"
    45  	"go/build"
    46  	"go/parser"
    47  	"go/token"
    48  	"go/types"
    49  	"io"
    50  	"io/ioutil"
    51  	"log"
    52  	"os"
    53  	"path/filepath"
    54  	"sort"
    55  	"strconv"
    56  	"strings"
    57  	"text/template"
    58  
    59  	"golang.org/x/tools/go/loader"
    60  
    61  	"lab.nexedi.com/kirr/go123/prog"
    62  	"lab.nexedi.com/kirr/go123/xerr"
    63  )
    64  
    65  // traceEvent represents 1 trace:event declaration.
    66  type traceEvent struct {
    67  	Pos  token.Position
    68  	Pkgt *Package // package this trace event is part of
    69  
    70  	// declaration of function to signal the event
    71  	// the declaration is constructed on the fly via converting e.g.
    72  	//
    73  	//	//trace:event traceConnRecv(c *Conn, msg Msg)
    74  	//
    75  	// into
    76  	//
    77  	//	func traceConnRecv(c *Conn, msg Msg)
    78  	//
    79  	// when trace:event is parsed the func declaration is not added
    80  	// anywhere in the sources - just its AST + package is virtually
    81  	// constructed.
    82  	//
    83  	// See parseTraceEvent for details.
    84  	*ast.FuncDecl
    85  }
    86  
    87  // traceImport represents 1 trace:import directive.
    88  type traceImport struct {
    89  	Pos     token.Position
    90  	PkgName string // "" if import name was not explicitly specified
    91  	PkgPath string
    92  }
    93  
    94  // traceImported represents 1 imported trace:event.
    95  type traceImported struct {
    96  	*traceEvent                   // imported event
    97  	ImportSpec  *traceImport      // imported via this spec
    98  	ImporterPkg *types.Package    // from this package
    99  	ImportedAs  map[string]string // in context where some packages are imported as named (pkgpath -> pkgname)
   100  }
   101  
   102  // Package represents tracing-related information about a package.
   103  type Package struct {
   104  	Pkgi *loader.PackageInfo // original non-augmented package
   105  
   106  	Eventv  []*traceEvent  // trace events this package defines
   107  	Importv []*traceImport // trace imports of other packages
   108  
   109  	// original package is augmented with tracing code
   110  	// information about tracing code is below:
   111  
   112  	traceFilev []*ast.File    // files for added trace:event funcs
   113  	traceFset  *token.FileSet // fset for ^^^
   114  
   115  	traceChecker  *types.Checker // to typecheck ^^^
   116  	tracePkg      *types.Package // original package augmented with ^^^
   117  	traceTypeInfo *types.Info    // typeinfo for ^^^
   118  }
   119  
   120  // parseTraceEvent parses trace event definition into traceEvent.
   121  //
   122  // text is text argument after "//trace:event ".
   123  func (p *Package) parseTraceEvent(srcfile *ast.File, pos token.Position, text string) (*traceEvent, error) {
   124  	posErr := func(format string, argv ...interface{}) error {
   125  		return fmt.Errorf("%v: "+format, append([]interface{}{pos}, argv...)...)
   126  	}
   127  
   128  	if !strings.HasPrefix(text, "trace") {
   129  		return nil, posErr("trace event must start with \"trace\"")
   130  	}
   131  
   132  	// prepare artificial package with trace event definition as func declaration
   133  	buf := &Buffer{}
   134  	buf.emit("package %s", p.Pkgi.Pkg.Name())
   135  
   136  	// add all imports from original source file
   137  	// so that inside it all looks like as if it was in original source context
   138  	buf.emit("\nimport (")
   139  
   140  	for _, imp := range srcfile.Imports {
   141  		impline := ""
   142  		if imp.Name != nil {
   143  			impline += imp.Name.Name + " "
   144  		}
   145  		impline += imp.Path.Value
   146  		buf.emit("\t%s", impline)
   147  	}
   148  
   149  	buf.emit(")")
   150  
   151  	// func itself
   152  	buf.emit("\nfunc " + text)
   153  
   154  	// now parse/typecheck
   155  	filename := fmt.Sprintf("%v:%v+trace:event %v", pos.Filename, pos.Line, text)
   156  	//println("---- 8< ----", filename)
   157  	//println(buf.String())
   158  	//println("---- 8< ----")
   159  	tf, err := parser.ParseFile(p.traceFset, filename, buf.String(), 0)
   160  	if err != nil {
   161  		return nil, err // already has pos' as prefix
   162  	}
   163  
   164  	p.traceFilev = append(p.traceFilev, tf)
   165  
   166  	// must be:
   167  	// GenDecl{IMPORT}
   168  	// FuncDecl
   169  	if len(tf.Decls) != 2 {
   170  		return nil, posErr("trace event must be func-like")
   171  	}
   172  
   173  	declf, ok := tf.Decls[1].(*ast.FuncDecl)
   174  	if !ok {
   175  		return nil, posErr("trace event must be func-like, not %v", tf.Decls[0])
   176  	}
   177  	// XXX ok to allow methods (declf.Recv != nil) ?
   178  	if declf.Type.Results != nil {
   179  		return nil, posErr("trace event must not return results")
   180  	}
   181  
   182  	// typecheck prepared file to get trace func argument types
   183  	// (type information lands into p.traceTypeInfo)
   184  	err = p.traceChecker.Files([]*ast.File{tf})
   185  	if err != nil {
   186  		return nil, err // already has pos' as prefix
   187  	}
   188  
   189  	return &traceEvent{Pos: pos, Pkgt: p, FuncDecl: declf}, nil
   190  }
   191  
   192  // parseTraceImport parses trace import directive into traceImport.
   193  //
   194  // text is text argument after "//trace:import ".
   195  func (p *Package) parseTraceImport(pos token.Position, text string) (*traceImport, error) {
   196  	// //trace:import "path/to/pkg"
   197  	// //trace:import name "path/to/pkg"
   198  
   199  	if len(text) == 0 {
   200  		return nil, fmt.Errorf("%v: empty trace-import spec", pos)
   201  	}
   202  
   203  	pkgname := ""
   204  	pkgqpath := text
   205  
   206  	if !(text[0] == '"' || text[0] == '\'') {
   207  		textv := strings.SplitN(text, " ", 2)
   208  		if len(textv) != 2 {
   209  			return nil, fmt.Errorf("%v: invalid trace-import spec %v", pos, text)
   210  		}
   211  		pkgname = textv[0]
   212  		pkgqpath = textv[1]
   213  	}
   214  
   215  	// Unquote pkgqpath as regular import does
   216  	pkgpath, err := strconv.Unquote(pkgqpath)
   217  	if err != nil || pkgpath == "" || pkgpath[0] == '\'' {
   218  		return nil, fmt.Errorf("%v: invalid trace-import path %v", pos, pkgqpath)
   219  	}
   220  
   221  	// reject duplicate imports
   222  	for _, imported := range p.Importv {
   223  		if pkgpath == imported.PkgPath {
   224  			return nil, fmt.Errorf("%v: duplicate trace import of %v (previous at %v)", pos, pkgpath, imported.Pos)
   225  		}
   226  	}
   227  
   228  	return &traceImport{Pos: pos, PkgName: pkgname, PkgPath: pkgpath}, nil
   229  }
   230  
   231  // progImporter is types.Importer that imports packages from loaded loader.Program .
   232  type progImporter struct {
   233  	prog *loader.Program
   234  }
   235  
   236  func (pi *progImporter) Import(path string) (*types.Package, error) {
   237  	pkgi := pi.prog.Package(path)
   238  	if pkgi == nil {
   239  		return nil, fmt.Errorf("package %q not found", path)
   240  	}
   241  
   242  	return pkgi.Pkg, nil
   243  }
   244  
   245  // packageTrace returns tracing information about a package.
   246  func packageTrace(prog *loader.Program, pkgi *loader.PackageInfo) (*Package, error) {
   247  	// prepare Package with typechecker ready to typecheck trace files
   248  	// (to get trace func argument types)
   249  	tconf := &types.Config{
   250  		Importer: &progImporter{prog},
   251  
   252  		// to ignore traceXXX() calls from original package code
   253  		IgnoreFuncBodies: true,
   254  
   255  		// we take imports from original source file verbatim,
   256  		// but most of them probably won't be used.
   257  		DisableUnusedImportCheck: true,
   258  	}
   259  
   260  	// tfset := token.NewFileSet() // XXX ok to separate or use original package fset?
   261  	tfset := prog.Fset
   262  	tpkg := types.NewPackage(pkgi.Pkg.Path(), pkgi.Pkg.Name())
   263  	tinfo := &types.Info{Types: make(map[ast.Expr]types.TypeAndValue)}
   264  
   265  	p := &Package{
   266  		Pkgi: pkgi,
   267  
   268  		// XXX vvv do we need separate field for traceFset if it is = prog.Fset?
   269  		traceFset:     tfset,
   270  		traceChecker:  types.NewChecker(tconf, tfset, tpkg, tinfo),
   271  		tracePkg:      tpkg,
   272  		traceTypeInfo: tinfo,
   273  	}
   274  
   275  	// preload original package files into tracing package
   276  	err := p.traceChecker.Files(p.Pkgi.Files)
   277  	if err != nil {
   278  		// must not happen
   279  		panic(fmt.Errorf("%v: error rechecking original package: %v", pkgi.Pkg.Path(), err))
   280  	}
   281  
   282  	// go through files of the original package and process //trace: directives
   283  	//
   284  	// NOTE before go1.10 we don't process cgo files as there go/loader passes to us
   285  	// already preprocessed results with comments stripped, not original source.
   286  	// This problem was fixed for go1.10 in https://github.com/golang/go/commit/85c3ebf4.
   287  	for _, file := range pkgi.Files {                        // ast.File
   288  		for _, commgroup := range file.Comments {        // ast.CommentGroup
   289  			for _, comment := range commgroup.List { // ast.Comment
   290  				pos := prog.Fset.Position(comment.Slash)
   291  				//fmt.Printf("%v %q\n", pos, comment.Text)
   292  
   293  				// only directives starting from beginning of line
   294  				if pos.Column != 1 {
   295  					continue
   296  				}
   297  
   298  				if !strings.HasPrefix(comment.Text, "//trace:") {
   299  					continue
   300  				}
   301  
   302  				textv := strings.SplitN(comment.Text, " ", 2)
   303  				if len(textv) != 2 {
   304  					return nil, fmt.Errorf("%v: invalid directive format", pos)
   305  				}
   306  
   307  				directive, arg := textv[0], textv[1]
   308  				switch directive {
   309  				case "//trace:event":
   310  					//fmt.Println("*", textv)
   311  					event, err := p.parseTraceEvent(file, pos, arg)
   312  					if err != nil {
   313  						return nil, err
   314  					}
   315  
   316  					// XXX needed here? - better append in parseTraceEvent
   317  					p.Eventv = append(p.Eventv, event)
   318  
   319  				case "//trace:import":
   320  					imported, err := p.parseTraceImport(pos, arg)
   321  					if err != nil {
   322  						return nil, err
   323  					}
   324  
   325  					// XXX needed here? - better append in parseTraceImport
   326  					p.Importv = append(p.Importv, imported)
   327  
   328  				default:
   329  					return nil, fmt.Errorf("%v: unknown tracing directive %q", pos, directive)
   330  				}
   331  			}
   332  		}
   333  	}
   334  
   335  	// events and imports go in canonical order
   336  	sort.Sort(byEventName(p.Eventv))
   337  	sort.Sort(byPkgPath(p.Importv))
   338  
   339  	return p, nil
   340  }
   341  
   342  // byEventName provides []*traceEvent ordering by event name.
   343  type byEventName []*traceEvent
   344  
   345  func (v byEventName) Less(i, j int) bool { return v[i].Name.Name < v[j].Name.Name }
   346  func (v byEventName) Swap(i, j int)      { v[i], v[j] = v[j], v[i] }
   347  func (v byEventName) Len() int           { return len(v) }
   348  
   349  // byPkgPath provides []*traceImport ordering by package path.
   350  type byPkgPath []*traceImport
   351  
   352  func (v byPkgPath) Less(i, j int) bool { return v[i].PkgPath < v[j].PkgPath }
   353  func (v byPkgPath) Swap(i, j int)      { v[i], v[j] = v[j], v[i] }
   354  func (v byPkgPath) Len() int           { return len(v) }
   355  
   356  // SplitTests splits package into main and test parts, each covering trace-related things accordingly.
   357  func (p *Package) SplitTests() (testPkg *Package) {
   358  	__ := *p
   359  	testPkg = &__
   360  
   361  	// relevant for tracing are only: .Eventv & .Importv
   362  	eventv := p.Eventv
   363  	importv := p.Importv
   364  	p.Eventv = nil
   365  	p.Importv = nil
   366  	testPkg.Eventv = nil
   367  	testPkg.Importv = nil
   368  
   369  	for _, e := range eventv {
   370  		if strings.HasSuffix(e.Pos.Filename, "_test.go") {
   371  			testPkg.Eventv = append(testPkg.Eventv, e)
   372  		} else {
   373  			p.Eventv = append(p.Eventv, e)
   374  		}
   375  	}
   376  
   377  	for _, i := range importv {
   378  		if strings.HasSuffix(i.Pos.Filename, "_test.go") {
   379  			testPkg.Importv = append(testPkg.Importv, i)
   380  		} else {
   381  			p.Importv = append(p.Importv, i)
   382  		}
   383  	}
   384  
   385  	return testPkg
   386  }
   387  
   388  // ----------------------------------------
   389  
   390  // Argv returns comma-separated argument-list.
   391  func (te *traceEvent) Argv() string {
   392  	argv := []string{}
   393  
   394  	for _, field := range te.FuncDecl.Type.Params.List {
   395  		for _, name := range field.Names {
   396  			argv = append(argv, name.Name)
   397  		}
   398  	}
   399  
   400  	return strings.Join(argv, ", ")
   401  }
   402  
   403  // ArgvTyped returns argument list with types.
   404  //
   405  // types are qualified relative to original package
   406  func (te *traceEvent) ArgvTyped() string {
   407  	return te.ArgvTypedRelativeTo(te.Pkgt.tracePkg, nil)
   408  }
   409  
   410  // ArgvTypedRelativeTo returns argument list with types qualified relative to specified package.
   411  //
   412  // importedAs specifies under which name a package was imported, if name was explicitly set
   413  func (te *traceEvent) ArgvTypedRelativeTo(pkg *types.Package, importedAs map[string]string /*pkgpath -> pkgname*/) string {
   414  	argv := []string{}
   415  
   416  	// default qualifier - relative to original package
   417  	qf := func(p *types.Package) string {
   418  		// specified package - unqualified
   419  		if p == pkg {
   420  			return ""
   421  		}
   422  
   423  		// qualify as explicitly named
   424  		pkgname := importedAs[p.Path()]
   425  		if pkgname != "" {
   426  			return pkgname
   427  		}
   428  
   429  		// default qualification
   430  		return p.Name()
   431  	}
   432  
   433  	for _, field := range te.FuncDecl.Type.Params.List {
   434  		namev := []string{}
   435  		for _, name := range field.Names {
   436  			namev = append(namev, name.Name)
   437  		}
   438  
   439  		arg := strings.Join(namev, ", ")
   440  		typ := te.Pkgt.traceTypeInfo.Types[field.Type].Type
   441  		arg += " " + types.TypeString(typ, qf)
   442  
   443  		argv = append(argv, arg)
   444  	}
   445  
   446  	return strings.Join(argv, ", ")
   447  }
   448  
   449  // NeedPkgv returns packages that are needed for argument types.
   450  func (te *traceEvent) NeedPkgv() []string {
   451  	pkgset := StrSet{ /*pkgpath*/ }
   452  	qf := func(pkg *types.Package) string {
   453  		// if we are called - pkg is used
   454  		pkgset.Add(pkg.Path())
   455  		return "" // don't care
   456  	}
   457  
   458  	for _, field := range te.FuncDecl.Type.Params.List {
   459  		typ := te.Pkgt.traceTypeInfo.Types[field.Type].Type
   460  		_ = types.TypeString(typ, qf)
   461  	}
   462  
   463  	return pkgset.Itemv()
   464  }
   465  
   466  // ImportSpec returns string representation of import spec.
   467  func (ti *traceImport) ImportSpec() string {
   468  	t := ti.PkgName
   469  	if t != "" {
   470  		t += " "
   471  	}
   472  	t += fmt.Sprintf("%q", ti.PkgPath)
   473  	return t
   474  }
   475  
   476  // traceEventCodeTmpl is code template generated for one trace event.
   477  var traceEventCodeTmpl = template.Must(template.New("traceevent").Parse(`
   478  // traceevent: {{.Name}}({{.ArgvTyped}})
   479  
   480  {{/* probe type for this trace event */ -}}
   481  type _t_{{.Name}} struct {
   482  	tracing.Probe
   483  	probefunc     func({{.ArgvTyped}})
   484  }
   485  
   486  {{/* list of probes attached (nil if nothing) */ -}}
   487  var _{{.Name}} *_t_{{.Name}}
   488  
   489  {{/* function which event producer calls to notify about the event
   490     *
   491     * after https://github.com/golang/go/issues/19348 is done this separate
   492     * checking function will be inlined and tracepoint won't cost a function
   493     * call when it is disabled */ -}}
   494  func {{.Name}}({{.ArgvTyped}}) {
   495  	if _{{.Name}} != nil {
   496  		_{{.Name}}_run({{.Argv}})
   497  	}
   498  }
   499  
   500  {{/* function to notify attached probes */ -}}
   501  func _{{.Name}}_run({{.ArgvTyped}}) {
   502  	for p := _{{.Name}}; p != nil; p = (*_t_{{.Name}})(unsafe.Pointer(p.Next())) {
   503  		p.probefunc({{.Argv}})
   504  	}
   505  }
   506  
   507  {{/* function to attach a probe to tracepoint */ -}}
   508  func {{.Name}}_Attach(pg *tracing.ProbeGroup, probe func({{.ArgvTyped}})) *tracing.Probe {
   509  	p := _t_{{.Name}}{probefunc: probe}
   510  	tracing.AttachProbe(pg, (**tracing.Probe)(unsafe.Pointer(&_{{.Name}})), &p.Probe)
   511  	return &p.Probe
   512  }
   513  `))
   514  
   515  // traceEventImportTmpl is code template generated for importing one trace event.
   516  var traceEventImportTmpl = template.Must(template.New("traceimport").Parse(`
   517  {{/* function to attach a probe to tracepoint imported via go:linkname */ -}}
   518  //go:linkname {{.ImportSpec.PkgName}}_{{.Name}}_Attach {{.ImportSpec.PkgPath}}.{{.Name}}_Attach
   519  func {{.ImportSpec.PkgName}}_{{.Name}}_Attach(*tracing.ProbeGroup, func({{.ArgvTypedRelativeTo .ImporterPkg .ImportedAs}})) *tracing.Probe
   520  `))
   521  
   522  // traceEventImportCheckTmpl is code template generated to check consistency with one imported package.
   523  var traceEventImportCheckTmpl = template.Must(template.New("traceimportcheck").Parse(`
   524  {{/* linking will fail if trace import code becomes out of sync wrt imported package */ -}}
   525  // rerun "gotrace gen" if you see link failure ↓↓↓
   526  //go:linkname {{.ImportSpec.PkgName}}_trace_exporthash {{.ImportSpec.PkgPath}}._trace_exporthash_{{.ExportHash}}
   527  func {{.ImportSpec.PkgName}}_trace_exporthash()
   528  func init() { {{.ImportSpec.PkgName}}_trace_exporthash() }
   529  `))
   530  
   531  // magic begins all files generated by gotrace.
   532  const magic = "// Code generated by lab.nexedi.com/kirr/go123/tracing/cmd/gotrace; DO NOT EDIT.\n"
   533  
   534  // checkCanWrite checks whether it is safe to write to file at path.
   535  //
   536  // it is safe to write when either
   537  // - the file does not exist, or
   538  // - it exits but was previously generated by us
   539  func checkCanWrite(path string) error {
   540  	f, err := os.Open(path)
   541  	if e, ok := err.(*os.PathError); ok && os.IsNotExist(e.Err) {
   542  		return nil
   543  	}
   544  
   545  	defer f.Close()
   546  	bf := bufio.NewReader(f)
   547  
   548  	headline, err := bf.ReadString('\n')
   549  	if err != nil || headline != magic {
   550  		return fmt.Errorf("refusing to make output: %v exists but was not generated by gotrace", path)
   551  	}
   552  
   553  	return nil
   554  }
   555  
   556  // writeFile writes data to a file at path after checking it is safe to write there.
   557  func writeFile(path string, data []byte) error {
   558  	err := checkCanWrite(path)
   559  	if err != nil {
   560  		return err
   561  	}
   562  
   563  	return ioutil.WriteFile(path, data, 0666)
   564  }
   565  
   566  // removeFile make sure there is no file at path after checking it is safe to write to that file.
   567  func removeFile(path string) error {
   568  	err := checkCanWrite(path)
   569  	if err != nil {
   570  		return err
   571  	}
   572  
   573  	err = os.Remove(path)
   574  	if e, ok := err.(*os.PathError); ok && os.IsNotExist(e.Err) {
   575  		err = nil
   576  	}
   577  	return err
   578  }
   579  
   580  // Program represents loaded program for tracepoint analysis.
   581  //
   582  // It is generalization of loader.Program due to loader not allowing to
   583  // construct programs incrementally.
   584  type Program struct {
   585  	// list of loader.Programs in use
   586  	//
   587  	// We generally need to have several programs because a package can
   588  	// trace:import another package which is not otherwise imported by
   589  	// original program.
   590  	//
   591  	// Since go/loader does not support incrementally augmenting loaded
   592  	// program with more packages, we work-around it with having several
   593  	// progs.
   594  	progv []*loader.Program
   595  
   596  	// config for loading programs
   597  	loaderConf *loader.Config
   598  }
   599  
   600  // NewProgram constructs new empty Program ready to load packages according to specified build context.
   601  func NewProgram(ctxt *build.Context, cwd string) *Program {
   602  	// adjust build context to filter-out ztrace* files when discovering packages
   603  	//
   604  	// we don't load what should be generated by us for 2 reasons:
   605  	// - code generated could be wrong with older version of the
   606  	//   tool - it should not prevent from regenerating.
   607  	// - generated code imports packages which might be not there
   608  	//   yet in gopath (lab.nexedi.com/kirr/go123/tracing)
   609  	ctxtReadDir := ctxt.ReadDir
   610  	if ctxtReadDir == nil {
   611  		ctxtReadDir = ioutil.ReadDir
   612  	}
   613  	ctxtNoZTrace := *ctxt
   614  	ctxtNoZTrace.ReadDir = func(dir string) ([]os.FileInfo, error) {
   615  		fv, err := ctxtReadDir(dir)
   616  		okv := fv[:0]
   617  		for _, f := range fv {
   618  			if !strings.HasPrefix(f.Name(), "ztrace") {
   619  				okv = append(okv, f)
   620  			}
   621  		}
   622  		return okv, err
   623  	}
   624  
   625  	p := &Program{}
   626  	p.loaderConf = &loader.Config{
   627  		ParserMode:          parser.ParseComments,
   628  		TypeCheckFuncBodies: func(path string) bool { return false },
   629  		Build:               &ctxtNoZTrace,
   630  		Cwd:                 cwd,
   631  	}
   632  
   633  	return p
   634  }
   635  
   636  // Import imports a package and returns associated package info and program
   637  // under which it was loaded.
   638  func (p *Program) Import(pkgpath string) (prog *loader.Program, pkgi *loader.PackageInfo, err error) {
   639  	// let's see - maybe it is already there
   640  	for _, prog := range p.progv {
   641  		pkgi := prog.Package(pkgpath)
   642  		if pkgi != nil {
   643  			return prog, pkgi, nil
   644  		}
   645  	}
   646  
   647  	// not found - we have to load new program rooted at pkgpath
   648  	p.loaderConf.ImportPkgs = nil
   649  	p.loaderConf.Import(pkgpath)
   650  
   651  	prog, err = p.loaderConf.Load()
   652  	if err != nil {
   653  		return nil, nil, err
   654  	}
   655  
   656  	if !(len(prog.Created) == 0 && len(prog.Imported) == 1) {
   657  		panic("import")
   658  	}
   659  
   660  	p.progv = append(p.progv, prog)
   661  	pkgi = prog.InitialPackages()[0]
   662  	return prog, pkgi, nil
   663  }
   664  
   665  // ImportWithTests imports a package augmented with code from _test.go files +
   666  // imports external test package (if present).
   667  func (p *Program) ImportWithTests(pkgpath string) (prog *loader.Program, pkgi *loader.PackageInfo, xtestPkgi *loader.PackageInfo, err error) {
   668  	// NOTE always reimporting not to interfere with regular imports
   669  	p.loaderConf.ImportPkgs = nil
   670  	p.loaderConf.ImportWithTests(pkgpath)
   671  
   672  	prog, err = p.loaderConf.Load()
   673  	if err != nil {
   674  		return nil, nil, nil, err
   675  	}
   676  
   677  	if len(prog.Imported) != 1 {
   678  		panic("import with tests")
   679  	}
   680  
   681  	if len(prog.Created) > 0 {
   682  		xtestPkgi = prog.Created[0]
   683  	}
   684  	for _, pkgi = range prog.Imported {
   685  	}
   686  
   687  	return prog, pkgi, xtestPkgi, nil
   688  }
   689  
   690  // ---- `gotrace gen` ----
   691  
   692  // tracegen generates code according to tracing directives in a package @ pkgpath.
   693  //
   694  // ctxt is build context for discovering packages
   695  // cwd is "current" directory for resolving local imports (e.g. packages like "./some/package")
   696  func tracegen(pkgpath string, ctxt *build.Context, cwd string) error {
   697  	P := NewProgram(ctxt, cwd)
   698  
   699  	lprog, pkgi, xtestPkgi, err := P.ImportWithTests(pkgpath)
   700  	if err != nil {
   701  		return err
   702  	}
   703  
   704  	// determine package directory
   705  	if len(pkgi.Files) == 0 {
   706  		return fmt.Errorf("package %s is empty", pkgi.Pkg.Path())
   707  	}
   708  
   709  	pkgdir := filepath.Dir(lprog.Fset.File(pkgi.Files[0].Pos()).Name())
   710  
   711  	// tracing info for this specified package
   712  	tpkg, err := packageTrace(lprog, pkgi)
   713  	if err != nil {
   714  		return err // XXX err ctx
   715  	}
   716  
   717  	// split everything related to tracing into plain and test (not xtest) packages
   718  	testTpkg := tpkg.SplitTests()
   719  
   720  	err1 := tracegen1(P, tpkg, pkgdir, "")
   721  	err2 := tracegen1(P, testTpkg, pkgdir, "_test")
   722  
   723  	// also handle xtest package
   724  	xtestTpkg := &Package{} // dummy package with empty .Eventv & .Importv
   725  	if xtestPkgi != nil {
   726  		xtestTpkg, err = packageTrace(lprog, xtestPkgi)
   727  		if err != nil {
   728  			return err // XXX err ctx
   729  		}
   730  	}
   731  
   732  	err3 := tracegen1(P, xtestTpkg, pkgdir, "_x_test")
   733  
   734  	return xerr.Merge(err1, err2, err3)
   735  }
   736  
   737  // tracegen1 generates code according to tracing directives for a (sub)package @pkgpath.
   738  //
   739  // subpackage is either original package, testing code, or external test package
   740  func tracegen1(P *Program, tpkg *Package, pkgdir string, kind string) error {
   741  	var err error
   742  
   743  	// write ztrace.go with code generated for trace events and imports
   744  	ztrace_go := filepath.Join(pkgdir, "ztrace"+kind+".go")
   745  	if len(tpkg.Eventv) == 0 && len(tpkg.Importv) == 0 {
   746  		err = removeFile(ztrace_go)
   747  		if err != nil {
   748  			return err
   749  		}
   750  	} else {
   751  		// prologue
   752  		prologue := &Buffer{}
   753  		prologue.WriteString(magic)
   754  		prologue.emit("\npackage %v", tpkg.Pkgi.Pkg.Name())
   755  		prologue.emit("// code generated for tracepoints")
   756  		prologue.emit("\nimport (")
   757  		prologue.emit("\t%q", "lab.nexedi.com/kirr/go123/tracing")
   758  
   759  		// pkgpaths of all packages needed for used types
   760  		needPkg := StrSet{}
   761  
   762  		// some packages are imported with explicit name
   763  		importedAs := map[string]string{} // pkgpath -> pkgname
   764  
   765  		text := &Buffer{}
   766  
   767  		// code for trace:event definitions
   768  		for _, event := range tpkg.Eventv {
   769  			needPkg.Add("unsafe") // used in tr
   770  			needPkg.Add(event.NeedPkgv()...)
   771  			err = traceEventCodeTmpl.Execute(text, event)
   772  			if err != nil {
   773  				panic(err)
   774  			}
   775  		}
   776  
   777  		// export hash symbol so that if importing package is out of
   778  		// sync - it will have it different and linking will fail.
   779  		if len(tpkg.Eventv) > 0 {
   780  			text.emit("\n// trace export signature")
   781  			//text.emit("---- 8< ----")
   782  			//fmt.Fprintf(text, "%s", traceExport(tpkg, kind))
   783  			//text.emit("---- 8< ----")
   784  			text.emit("func _trace_exporthash_%s() {}", traceExportHash(tpkg, kind))
   785  		}
   786  
   787  		// code for trace:import imports
   788  		for _, timport := range tpkg.Importv {
   789  			text.emit("\n// traceimport: %s", timport.ImportSpec())
   790  
   791  			impProg, impPkgi, err := P.Import(timport.PkgPath)
   792  			if err != nil {
   793  				return fmt.Errorf("%v: error trace-importing %s: %v", timport.Pos, timport.PkgPath, err)
   794  			}
   795  
   796  			// set name of the package if it was not explicitly specified
   797  			if timport.PkgName == "" {
   798  				timport.PkgName = impPkgi.Pkg.Name()
   799  			} else {
   800  				importedAs[timport.PkgPath] = timport.PkgName
   801  			}
   802  
   803  			impPkg, err := packageTrace(impProg, impPkgi)
   804  			if err != nil {
   805  				return err // XXX err ctx
   806  			}
   807  
   808  			if len(impPkg.Eventv) == 0 {
   809  				return fmt.Errorf("%v: package %v does not export anything trace-related", timport.Pos, timport.PkgPath)
   810  			}
   811  
   812  			// verify export hash so link fails if it gets out of sync with imported package
   813  			err = traceEventImportCheckTmpl.Execute(text, struct {
   814  				ImportSpec *traceImport
   815  				ExportHash string
   816  			}{
   817  				timport,
   818  				traceExportHash(impPkg, "" /*regular package*/)})
   819  			if err != nil {
   820  				return err // XXX err ctx
   821  			}
   822  
   823  			text.emit("")
   824  
   825  			// import individual events
   826  			for _, event := range impPkg.Eventv {
   827  				needPkg.Add(event.NeedPkgv()...)
   828  				importedEvent := traceImported{
   829  					traceEvent:  event,
   830  					ImportSpec:  timport,
   831  					ImporterPkg: tpkg.Pkgi.Pkg,
   832  					ImportedAs:  importedAs,
   833  				}
   834  				err = traceEventImportTmpl.Execute(text, importedEvent)
   835  				if err != nil {
   836  					panic(err)
   837  				}
   838  			}
   839  		}
   840  
   841  		// finish prologue with needed imports
   842  		if !needPkg.Has("unsafe") {
   843  			// we need it anyway because go:linkname is not allowed without unsafe
   844  			prologue.emit("\t_ %q", "unsafe")
   845  		} else {
   846  			prologue.emit("\t%q", "unsafe")
   847  			needPkg.Delete("unsafe")
   848  		}
   849  
   850  		needPkg.Delete(tpkg.Pkgi.Pkg.Path()) // our pkg - no need to import
   851  		needPkgv := needPkg.Itemv()
   852  		if len(needPkgv) > 0 {
   853  			prologue.emit("")
   854  		}
   855  
   856  		for _, needpkg := range needPkgv {
   857  			pkgname := importedAs[needpkg]
   858  			if pkgname != "" {
   859  				pkgname += " "
   860  			}
   861  			prologue.emit("\t%s%q", pkgname, needpkg)
   862  		}
   863  		prologue.emit(")")
   864  
   865  		// write output to ztrace.go
   866  		fulltext := append(prologue.Bytes(), text.Bytes()...)
   867  		err = writeFile(ztrace_go, fulltext)
   868  		if err != nil {
   869  			return err
   870  		}
   871  	}
   872  
   873  	// write empty ztrace.s so go:linkname works, if there are trace imports
   874  	ztrace_s := filepath.Join(pkgdir, "ztrace"+kind+".s")
   875  	if len(tpkg.Importv) == 0 {
   876  		err = removeFile(ztrace_s)
   877  	} else {
   878  		text := &Buffer{}
   879  		text.WriteString(magic)
   880  		text.emit("// empty .s so `go build` does not use -complete for go:linkname to work")
   881  
   882  		err = writeFile(ztrace_s, text.Bytes())
   883  	}
   884  
   885  	if err != nil {
   886  		return err
   887  	}
   888  
   889  	return nil
   890  }
   891  
   892  // traceExport returns signatures of all tracing-related exports of a package
   893  // in canonical order as would be seen from universe scope.
   894  func traceExport(tpkg *Package, kind string) []byte {
   895  	pkgpath := tpkg.Pkgi.Pkg.Path()
   896  	pkgname := tpkg.Pkgi.Pkg.Name()
   897  
   898  	exported := &Buffer{}
   899  	exported.emit("%q %q", pkgpath, kind)
   900  
   901  	for _, event := range tpkg.Eventv {
   902  		importedEvent := traceImported{
   903  			traceEvent:  event,
   904  			ImportSpec:  &traceImport{PkgName: pkgname, PkgPath: pkgpath},
   905  			ImporterPkg: nil, // from nowhere
   906  			ImportedAs:  nil, // no naming for imports
   907  		}
   908  		err := traceEventImportTmpl.Execute(exported, importedEvent)
   909  		if err != nil {
   910  			panic(err)
   911  		}
   912  	}
   913  
   914  	return exported.Bytes()
   915  }
   916  
   917  // traceExportHash computes signature of tracing-related exports of a package
   918  // implementation note: it is sha1 of associated header + importing code as
   919  // if it was executed from universe scope.
   920  func traceExportHash(tpkg *Package, kind string) string {
   921  	return fmt.Sprintf("%x", sha1.Sum(traceExport(tpkg, kind)))
   922  }
   923  
   924  const genSummary = "generate code according to tracing annotations and imports"
   925  
   926  func genUsage(w io.Writer) {
   927  	fmt.Fprintf(w,
   928  `Usage: gotrace gen <package>
   929  Generate code according to tracing annotations and imports
   930  
   931    options:
   932  
   933          -h --help       this help text.
   934  `)
   935  }
   936  
   937  func genMain(argv []string) {
   938  	flags := flag.FlagSet{Usage: func() { genUsage(os.Stderr) }}
   939  	flags.Init("", flag.ExitOnError)
   940  	flags.Parse(argv[1:])
   941  
   942  	argv = flags.Args()
   943  	if len(argv) < 1 {
   944  		flags.Usage()
   945  		prog.Exit(2)
   946  	}
   947  	pkgpath := argv[0]
   948  
   949  	cwd, err := os.Getwd()
   950  	if err != nil {
   951  		prog.Fatal(err)
   952  	}
   953  
   954  	err = tracegen(pkgpath, &build.Default, cwd)
   955  	if err != nil {
   956  		prog.Fatal(err)
   957  	}
   958  }
   959  
   960  
   961  // ---- `gotrace list` ----
   962  
   963  // tracelist lists trace-events defined by a package @ pkgpath.
   964  //
   965  // ctxt and cwd are tunables for discovering packages. See tracegen for details.
   966  //
   967  // TODO support listing by pkgspec (e.g. "./...")
   968  func tracelist(w io.Writer, pkgpath string, ctxt *build.Context, cwd string) error {
   969  	P := NewProgram(ctxt, cwd)
   970  
   971  	// NOTE only listing trace-events provided by main package, not tests or xtest
   972  	lprog, pkgi, err := P.Import(pkgpath)
   973  	if err != nil {
   974  		return err
   975  	}
   976  
   977  	tpkg, err := packageTrace(lprog, pkgi)
   978  	if err != nil {
   979  		return err // XXX err ctx
   980  	}
   981  
   982  	for _, event := range tpkg.Eventv {
   983  		_, err = fmt.Fprintf(w, "%s:%s\n", event.Pkgt.Pkgi.Pkg.Path(), event.Name)
   984  		if err != nil {
   985  			return err
   986  		}
   987  	}
   988  
   989  	return nil
   990  }
   991  
   992  const listSummary = "lists tracepoints defined by a package"
   993  
   994  func listUsage(w io.Writer) {
   995  	fmt.Fprintf(w,
   996  `Usage: gotrace list <package>
   997  List tracepoints defined by a package
   998  
   999    options:
  1000  
  1001          -h --help       this help text.
  1002  `)
  1003  }
  1004  
  1005  func listMain(argv []string) {
  1006  	flags := flag.FlagSet{Usage: func() { listUsage(os.Stderr) }}
  1007  	flags.Init("", flag.ExitOnError)
  1008  	flags.Parse(argv[1:])
  1009  
  1010  	argv = flags.Args()
  1011  	if len(argv) < 1 {
  1012  		flags.Usage()
  1013  		prog.Exit(2)
  1014  	}
  1015  	pkgpath := argv[0]
  1016  
  1017  	cwd, err := os.Getwd()
  1018  	if err != nil {
  1019  		prog.Fatal(err)
  1020  	}
  1021  
  1022  	err = tracelist(os.Stdout, pkgpath, &build.Default, cwd)
  1023  	if err != nil {
  1024  		prog.Fatal(err)
  1025  	}
  1026  }
  1027  
  1028  // ---- main driver ----
  1029  
  1030  var commands = prog.CommandRegistry{
  1031  	{"gen",	 genSummary,  genUsage,  genMain},
  1032  	{"list", listSummary, listUsage, listMain},
  1033  }
  1034  
  1035  var helpTopics = prog.HelpRegistry{
  1036  	// XXX for now empty
  1037  }
  1038  
  1039  var gotrace = prog.MainProg{
  1040  	Name:       "gotrace",
  1041  	Summary:    "Gotrace is a program to support and interact with go tracing subsystem",
  1042  	Commands:   commands,
  1043  	HelpTopics: helpTopics,
  1044  }
  1045  
  1046  func main() {
  1047  	log.SetFlags(0)
  1048  	log.SetPrefix("gotrace: ")
  1049  	gotrace.Main()
  1050  }