Add support for stable Func keys in App Engine second gen (#184)

Func keys include the filename where the Func is created. The filename is parsed according to these rules:

* Paths in package main are shortened to just the file name (github.com/foo/foo.go -> foo.go)
* Paths are stripped to just package paths (/go/src/github.com/foo/bar.go -> github.com/foo/bar.go)
* Module versions are stripped (/go/pkg/mod/github.com/foo/[email protected]/baz.go -> github.com/foo/bar/baz.go)

Author: sbuss

delay/delay.go

@@ -34,8 +34,16 @@ be associated with the request that invoked the Call method.
 The state of a function invocation that has not yet successfully
 executed is preserved by combining the file name in which it is declared
 with the string key that was passed to the Func function. Updating an app
-with pending function invocations is safe as long as the relevant
-functions have the (filename, key) combination preserved.
+with pending function invocations should safe as long as the relevant
+functions have the (filename, key) combination preserved. The filename is
+parsed according to these rules:
+  * Paths in package main are shortened to just the file name (github.com/foo/foo.go -> foo.go)
+  * Paths are stripped to just package paths (/go/src/github.com/foo/bar.go -> github.com/foo/bar.go)
+  * Module versions are stripped (/go/pkg/mod/github.com/foo/[email protected]/baz.go -> github.com/foo/bar/baz.go)
+
+There is some inherent risk of pending function invocations being lost during
+an update that contains large changes. For example, switching from using GOPATH
+to go.mod is a large change that may inadvertently cause file paths to change.
 
 The delay package uses the Task Queue API to create tasks that call the
 reserved application path "/_ah/queue/go/delay".
@@ -50,13 +58,19 @@ import (
 	"encoding/gob"
 	"errors"
 	"fmt"
+	"go/build"
+	stdlog "log"
 	"net/http"
+	"path/filepath"
 	"reflect"
+	"regexp"
 	"runtime"
+	"strings"
 
 	"golang.org/x/net/context"
 
 	"google.golang.org/appengine"
+	"google.golang.org/appengine/internal"
 	"google.golang.org/appengine/log"
 	"google.golang.org/appengine/taskqueue"
 )
@@ -98,6 +112,45 @@ func isContext(t reflect.Type) bool {
 	return t == stdContextType || t == netContextType
 }
 
+var modVersionPat = regexp.MustCompile("@v[^/]+")
+
+// fileKey finds a stable representation of the caller's file path.
+// For calls from package main: strip all leading path entries, leaving just the filename.
+// For calls from anywhere else, strip $GOPATH/src, leaving just the package path and file path.
+func fileKey(file string) (string, error) {
+	if !internal.IsSecondGen() || internal.MainPath == "" {
+		return file, nil
+	}
+	// If the caller is in the same Dir as mainPath, then strip everything but the file name.
+	if filepath.Dir(file) == internal.MainPath {
+		return filepath.Base(file), nil
+	}
+	// If the path contains "_gopath/src/", which is what the builder uses for
+	// apps which don't use go modules, strip everything up to and including src.
+	// Or, if the path starts with /tmp/staging, then we're importing a package
+	// from the app's module (and we must be using go modules), and we have a
+	// path like /tmp/staging1234/srv/... so strip everything up to and
+	// including the first /srv/.
+	// And be sure to look at the GOPATH, for local development.
+	s := string(filepath.Separator)
+	for _, s := range []string{filepath.Join("_gopath", "src") + s, s + "srv" + s, filepath.Join(build.Default.GOPATH, "src") + s} {
+		if idx := strings.Index(file, s); idx > 0 {
+			return file[idx+len(s):], nil
+		}
+	}
+
+	// Finally, if that all fails then we must be using go modules, and the file is a module,
+	// so the path looks like /go/pkg/mod/github.com/foo/[email protected]/baz.go
+	// So... remove everything up to and including mod, plus the @.... version string.
+	m := "/mod/"
+	if idx := strings.Index(file, m); idx > 0 {
+		file = file[idx+len(m):]
+	} else {
+		return file, fmt.Errorf("fileKey: unknown file path format for %q", file)
+	}
+	return modVersionPat.ReplaceAllString(file, ""), nil
+}
+
 // Func declares a new Function. The second argument must be a function with a
 // first argument of type context.Context.
 // This function must be called at program initialization time. That means it
@@ -111,7 +164,12 @@ func Func(key string, i interface{}) *Function {
 
 	// Derive unique, somewhat stable key for this func.
 	_, file, _, _ := runtime.Caller(1)
-	f.key = file + ":" + key
+	fk, err := fileKey(file)
+	if err != nil {
+		// Not fatal, but log the error
+		stdlog.Printf("delay: %v", err)
+	}
+	f.key = fk + ":" + key
 
 	t := f.fv.Type()
 	if t.Kind() != reflect.Func {

delay/delay_test.go

@@ -12,6 +12,8 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"path/filepath"
 	"reflect"
 	"testing"
 
@@ -462,3 +464,78 @@ func TestStandardContext(t *testing.T) {
 		t.Errorf("stdCtxRuns: got %d, want 1", stdCtxRuns)
 	}
 }
+
+func TestFileKey(t *testing.T) {
+	os.Setenv("GAE_ENV", "standard")
+	tests := []struct {
+		mainPath string
+		file     string
+		want     string
+	}{
+		// first-gen
+		{
+			"",
+			filepath.FromSlash("srv/foo.go"),
+			filepath.FromSlash("srv/foo.go"),
+		},
+		// gopath
+		{
+			filepath.FromSlash("/tmp/staging1234/srv/"),
+			filepath.FromSlash("/tmp/staging1234/srv/foo.go"),
+			"foo.go",
+		},
+		{
+			filepath.FromSlash("/tmp/staging1234/srv/_gopath/src/example.com/foo"),
+			filepath.FromSlash("/tmp/staging1234/srv/_gopath/src/example.com/foo/foo.go"),
+			"foo.go",
+		},
+		{
+			filepath.FromSlash("/tmp/staging2234/srv/_gopath/src/example.com/foo"),
+			filepath.FromSlash("/tmp/staging2234/srv/_gopath/src/example.com/foo/bar/bar.go"),
+			filepath.FromSlash("example.com/foo/bar/bar.go"),
+		},
+		{
+			filepath.FromSlash("/tmp/staging3234/srv/_gopath/src/example.com/foo"),
+			filepath.FromSlash("/tmp/staging3234/srv/_gopath/src/example.com/bar/main.go"),
+			filepath.FromSlash("example.com/bar/main.go"),
+		},
+		// go mod, same package
+		{
+			filepath.FromSlash("/tmp/staging3234/srv"),
+			filepath.FromSlash("/tmp/staging3234/srv/main.go"),
+			"main.go",
+		},
+		{
+			filepath.FromSlash("/tmp/staging3234/srv"),
+			filepath.FromSlash("/tmp/staging3234/srv/bar/main.go"),
+			filepath.FromSlash("bar/main.go"),
+		},
+		{
+			filepath.FromSlash("/tmp/staging3234/srv/cmd"),
+			filepath.FromSlash("/tmp/staging3234/srv/cmd/main.go"),
+			"main.go",
+		},
+		{
+			filepath.FromSlash("/tmp/staging3234/srv/cmd"),
+			filepath.FromSlash("/tmp/staging3234/srv/bar/main.go"),
+			filepath.FromSlash("bar/main.go"),
+		},
+		// go mod, other package
+		{
+			filepath.FromSlash("/tmp/staging3234/srv"),
+			filepath.FromSlash("/go/pkg/mod/github.com/foo/[email protected]/baz.go"),
+			filepath.FromSlash("github.com/foo/bar/baz.go"),
+		},
+	}
+	for i, tc := range tests {
+		internal.MainPath = tc.mainPath
+		got, err := fileKey(tc.file)
+		if err != nil {
+			t.Errorf("Unexpected error, call %v, file %q: %v", i, tc.file, err)
+			continue
+		}
+		if got != tc.want {
+			t.Errorf("Call %v, file %q: got %q, want %q", i, tc.file, got, tc.want)
+		}
+	}
+}

internal/main.go

@@ -11,5 +11,6 @@ import (
 )
 
 func Main() {
+	MainPath = ""
 	appengine_internal.Main()
 }

internal/main_common.go

@@ -0,0 +1,7 @@
+package internal
+
+// MainPath stores the file path of the main package. On App Engine Standard
+// using Go version 1.9 and below, this will be unset. On App Engine Flex and
+// App Engine Standard second-gen (Go 1.11 and above), this will be the
+// filepath to package main.
+var MainPath string

internal/main_test.go

@@ -0,0 +1,18 @@
+// +build !appengine
+
+package internal
+
+import (
+	"go/build"
+	"path/filepath"
+	"testing"
+)
+
+func TestFindMainPath(t *testing.T) {
+	// Tests won't have package main, instead they have testing.tRunner
+	want := filepath.Join(build.Default.GOROOT, "src", "testing", "testing.go")
+	got := findMainPath()
+	if want != got {
+		t.Errorf("findMainPath: want %s, got %s", want, got)
+	}
+}

internal/main_vm.go

@@ -12,9 +12,12 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"path/filepath"
+	"runtime"
 )
 
 func Main() {
+	MainPath = filepath.Dir(findMainPath())
 	installHealthChecker(http.DefaultServeMux)
 
 	port := "8080"
@@ -31,6 +34,24 @@ func Main() {
 	}
 }
 
+// Find the path to package main by looking at the root Caller.
+func findMainPath() string {
+	pc := make([]uintptr, 100)
+	n := runtime.Callers(2, pc)
+	frames := runtime.CallersFrames(pc[:n])
+	for {
+		frame, more := frames.Next()
+		// Tests won't have package main, instead they have testing.tRunner
+		if frame.Function == "main.main" || frame.Function == "testing.tRunner" {
+			return frame.File
+		}
+		if !more {
+			break
+		}
+	}
+	return ""
+}
+
 func installHealthChecker(mux *http.ServeMux) {
 	// If no health check handler has been installed by this point, add a trivial one.
 	const healthPath = "/_ah/health"