internal/engine/env.go - shac-project/shac - Git at Google

 // Copyright 2023 The Shac Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package engine

 import (
 	"context"
 	"errors"
 	"fmt"
 	"io"
 	"io/fs"
 	"path"
 	"strings"
 	"sync"

 	"go.starlark.net/resolve"
 	"go.starlark.net/starlark"
 	"go.starlark.net/syntax"
 )

 type printImpl func(th *starlark.Thread, msg string)

 // sourceKey is a reference as parsed by load().
 type sourceKey struct {
 	orig    string
 	pkg     string
 	relpath string
 }

 func (s *sourceKey) String() string {
 	if s.pkg == "__main__" {
 		return "//" + s.relpath
 	}
 	return "@" + s.pkg + "//" + s.relpath
 }

 func parseSourceKey(parent sourceKey, s string) (sourceKey, error) {
 	sk := sourceKey{orig: s}
 	if s == "" {
 		return sk, errors.New("empty reference")
 	}
 	// It is an external reference.
 	if strings.HasPrefix(s, "@") {
 		parts := strings.SplitN(s[1:], "//", 2)
 		sk.pkg = parts[0]
 		if len(parts) == 1 {
 			// The root of a package is api.star, to not conflict with self-tests that
 			// would live in shac.star.
 			if strings.HasSuffix(sk.pkg, "/") {
 				return sk, fmt.Errorf("illegal external reference trailing \"/\": %s", s)
 			}
 			sk.relpath = "api.star"
 		} else {
 			// A file within the package is referenced.
 			if len(parts[1]) == 0 {
 				return sk, fmt.Errorf("illegal external reference path empty: %s", s)
 			}
 			for _, p := range strings.SplitN(parts[1], "/", -1) {
 				if p == "internal" {
 					return sk, fmt.Errorf("illegal external reference path containing \"internal\": %s", s)
 				}
 				if p == ".." {
 					return sk, fmt.Errorf("illegal external reference path containing \"..\": %s", s)
 				}
 				if p == "" {
 					return sk, fmt.Errorf("illegal external reference path containing \"//\": %s", s)
 				}
 			}
 			sk.relpath = parts[1]
 		}
 		return sk, nil
 	}

 	// It is an internal reference.
 	sk.pkg = parent.pkg
 	if strings.HasPrefix(s, "//") {
 		// It is root relative.
 		sk.relpath = s[2:]
 	} else {
 		// It is path relative.
 		sk.relpath = path.Clean(path.Join(path.Dir(parent.relpath), s))
 	}
 	return sk, nil
 }

 // loadedSource is the outcome of a load() statement.
 type loadedSource struct {
 	// th is the starlark thread that loaded this source file.
 	th *starlark.Thread

 	mu      sync.Mutex
 	globals starlark.StringDict
 	err     error
 }

 // starlarkEnv is the running environment enabling to run multiple starlark
 // files in parallel.
 type starlarkEnv struct {
 	// Immutable.
 	// globals is available to all load() statements. They must be frozen via
 	// Freeze().
 	globals starlark.StringDict
 	// packages are all the available packages. It must include __main__.
 	packages map[string]fs.FS
 	// Options for parsing Starlark.
 	opts *syntax.FileOptions

 	// Mutable.
 	mu sync.Mutex
 	// sources are the processed sources. Augments as more sources are added.
 	sources map[string]*loadedSource
 }

 // thread returns a new starlark thread.
 //
 // load() statement is not allowed.
 func (e *starlarkEnv) thread(ctx context.Context, name string, pi printImpl) *starlark.Thread {
 	t := &starlark.Thread{Name: name, Print: pi}
 	t.SetLocal("shac.context", ctx)
 	return t
 }

 // getContext returns the context.Context given a starlark thread.
 func getContext(t *starlark.Thread) context.Context {
 	return t.Local("shac.context").(context.Context)
 }

 // load loads a starlark source file. It is safe to call it concurrently.
 //
 // A thread will be implicitly created.
 func (e *starlarkEnv) load(ctx context.Context, sk sourceKey, pi printImpl) (starlark.StringDict, error) {
 	// We are the root thread. Start a thread implicitly.
 	t := e.thread(ctx, sk.String(), pi)
 	t.Load = func(th *starlark.Thread, str string) (starlark.StringDict, error) {
 		skn, err := parseSourceKey(th.Local("shac.pkg").(sourceKey), str)
 		if err != nil {
 			return nil, err
 		}
 		return e.loadInner(th, skn)
 	}
 	t.SetLocal("shac.top", sk)
 	t.SetLocal("shac.pkg", sk)
 	return e.loadInner(t, sk)
 }

 func (e *starlarkEnv) loadInner(th *starlark.Thread, sk sourceKey) (starlark.StringDict, error) {
 	key := sk.String()
 	e.mu.Lock()
 	if source, ok := e.sources[key]; ok {
 		// source has already been loaded or is in the process of being loaded
 		// by another thread.

 		e.mu.Unlock()
 		if source.th == th {
 			// `loadInner` will not be called concurrently with the same
 			// Starlark thread, so if the current thread already has a lock on
 			// this source it means there's a dependency cycle that would cause
 			// a deadlock.
 			if !source.mu.TryLock() {
 				return nil, fmt.Errorf("%s was loaded in a cycle dependency graph", sk.String())
 			}
 		} else {
 			// The source may be concurrently processed by another starlark
 			// thread that's traversing the dependency graph of a separate
 			// shac.star file, wait for the processing to complete by taking the
 			// lock.
 			//
 			// This block is hard to cover since it's a race condition. We'd
 			// have to inject a builtin that would hang the starlark execution
 			// to force concurrency.
 			source.mu.Lock()
 		}
 		defer source.mu.Unlock()
 		return source.globals, source.err
 	}

 	source := &loadedSource{th: th, err: fmt.Errorf("load(%q) failed: panic while loading", sk)}
 	// We are the "master" of this file, since we will be loading it. Make sure
 	// others won't load it.
 	source.mu.Lock()
 	e.sources[key] = source
 	e.mu.Unlock()
 	defer source.mu.Unlock()

 	// It's a new file that wasn't load'ed yet. ss.mu is still held which will
 	// block concurrent loading.
 	if pkg := e.packages[sk.pkg]; pkg != nil {
 		if f, err := pkg.Open(sk.relpath); err == nil {
 			var d []byte
 			if d, err = io.ReadAll(f); err == nil {
 				oldsk := th.Local("shac.pkg").(sourceKey)
 				th.SetLocal("shac.pkg", sk)
 				fp := syntax.FilePortion{Content: d, FirstLine: 1, FirstCol: 1}
 				source.globals, source.err = starlark.ExecFileOptions(e.opts, th, sk.String(), fp, e.globals)
 				th.SetLocal("shac.pkg", oldsk)
 				var errl resolve.ErrorList
 				if errors.As(source.err, &errl) {
 					// Unwrap the error, only keep the first one.
 					source.err = errl[0]
 				}
 				var errre resolve.Error
 				if errors.As(source.err, &errre) {
 					// Synthesize a BacktraceableError since it's nicer for the
 					// user. Sadly we can't get the function context even if the
 					// error is within a function implementation, so hardcode
 					// "<toplevel>".
 					source.err = &failure{
 						Message: errre.Msg,
 						Stack: starlark.CallStack{
 							starlark.CallFrame{
 								Name: "<toplevel>",
 								Pos:  errre.Pos,
 							},
 						},
 					}
 				}
 			} else {
 				source.err = err
 			}
 			if err = f.Close(); source.err == nil {
 				source.err = err
 			}
 		} else if errors.Is(err, fs.ErrNotExist) {
 			// Hide the underlying error for determinism.
 			source.err = fmt.Errorf("%s not found", sk.relpath)
 		} else {
 			source.err = err
 		}
 	} else {
 		source.err = errors.New("package not found")
 	}
 	return source.globals, source.err
 }
	// Copyright 2023 The Shac Authors
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package engine

	import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"path"
	"strings"
	"sync"

	"go.starlark.net/resolve"
	"go.starlark.net/starlark"
	"go.starlark.net/syntax"
	)

	type printImpl func(th *starlark.Thread, msg string)

	// sourceKey is a reference as parsed by load().
	type sourceKey struct {
	orig string
	pkg string
	relpath string
	}

	func (s *sourceKey) String() string {
	if s.pkg == "__main__" {
	return "//" + s.relpath
	}
	return "@" + s.pkg + "//" + s.relpath
	}

	func parseSourceKey(parent sourceKey, s string) (sourceKey, error) {
	sk := sourceKey{orig: s}
	if s == "" {
	return sk, errors.New("empty reference")
	}
	// It is an external reference.
	if strings.HasPrefix(s, "@") {
	parts := strings.SplitN(s[1:], "//", 2)
	sk.pkg = parts[0]
	if len(parts) == 1 {
	// The root of a package is api.star, to not conflict with self-tests that
	// would live in shac.star.
	if strings.HasSuffix(sk.pkg, "/") {
	return sk, fmt.Errorf("illegal external reference trailing \"/\": %s", s)
	}
	sk.relpath = "api.star"
	} else {
	// A file within the package is referenced.
	if len(parts[1]) == 0 {
	return sk, fmt.Errorf("illegal external reference path empty: %s", s)
	}
	for _, p := range strings.SplitN(parts[1], "/", -1) {
	if p == "internal" {
	return sk, fmt.Errorf("illegal external reference path containing \"internal\": %s", s)
	}
	if p == ".." {
	return sk, fmt.Errorf("illegal external reference path containing \"..\": %s", s)
	}
	if p == "" {
	return sk, fmt.Errorf("illegal external reference path containing \"//\": %s", s)
	}
	}
	sk.relpath = parts[1]
	}
	return sk, nil
	}

	// It is an internal reference.
	sk.pkg = parent.pkg
	if strings.HasPrefix(s, "//") {
	// It is root relative.
	sk.relpath = s[2:]
	} else {
	// It is path relative.
	sk.relpath = path.Clean(path.Join(path.Dir(parent.relpath), s))
	}
	return sk, nil
	}

	// loadedSource is the outcome of a load() statement.
	type loadedSource struct {
	// th is the starlark thread that loaded this source file.
	th *starlark.Thread

	mu sync.Mutex
	globals starlark.StringDict
	err error
	}

	// starlarkEnv is the running environment enabling to run multiple starlark
	// files in parallel.
	type starlarkEnv struct {
	// Immutable.
	// globals is available to all load() statements. They must be frozen via
	// Freeze().
	globals starlark.StringDict
	// packages are all the available packages. It must include __main__.
	packages map[string]fs.FS
	// Options for parsing Starlark.
	opts *syntax.FileOptions

	// Mutable.
	mu sync.Mutex
	// sources are the processed sources. Augments as more sources are added.
	sources map[string]*loadedSource
	}

	// thread returns a new starlark thread.
	//
	// load() statement is not allowed.
	func (e starlarkEnv) thread(ctx context.Context, name string, pi printImpl) starlark.Thread {
	t := &starlark.Thread{Name: name, Print: pi}
	t.SetLocal("shac.context", ctx)
	return t
	}

	// getContext returns the context.Context given a starlark thread.
	func getContext(t *starlark.Thread) context.Context {
	return t.Local("shac.context").(context.Context)
	}

	// load loads a starlark source file. It is safe to call it concurrently.
	//
	// A thread will be implicitly created.
	func (e *starlarkEnv) load(ctx context.Context, sk sourceKey, pi printImpl) (starlark.StringDict, error) {
	// We are the root thread. Start a thread implicitly.
	t := e.thread(ctx, sk.String(), pi)
	t.Load = func(th *starlark.Thread, str string) (starlark.StringDict, error) {
	skn, err := parseSourceKey(th.Local("shac.pkg").(sourceKey), str)
	if err != nil {
	return nil, err
	}
	return e.loadInner(th, skn)
	}
	t.SetLocal("shac.top", sk)
	t.SetLocal("shac.pkg", sk)
	return e.loadInner(t, sk)
	}

	func (e starlarkEnv) loadInner(th starlark.Thread, sk sourceKey) (starlark.StringDict, error) {
	key := sk.String()
	e.mu.Lock()
	if source, ok := e.sources[key]; ok {
	// source has already been loaded or is in the process of being loaded
	// by another thread.

	e.mu.Unlock()
	if source.th == th {
	// `loadInner` will not be called concurrently with the same
	// Starlark thread, so if the current thread already has a lock on
	// this source it means there's a dependency cycle that would cause
	// a deadlock.
	if !source.mu.TryLock() {
	return nil, fmt.Errorf("%s was loaded in a cycle dependency graph", sk.String())
	}
	} else {
	// The source may be concurrently processed by another starlark
	// thread that's traversing the dependency graph of a separate
	// shac.star file, wait for the processing to complete by taking the
	// lock.
	//
	// This block is hard to cover since it's a race condition. We'd
	// have to inject a builtin that would hang the starlark execution
	// to force concurrency.
	source.mu.Lock()
	}
	defer source.mu.Unlock()
	return source.globals, source.err
	}

	source := &loadedSource{th: th, err: fmt.Errorf("load(%q) failed: panic while loading", sk)}
	// We are the "master" of this file, since we will be loading it. Make sure
	// others won't load it.
	source.mu.Lock()
	e.sources[key] = source
	e.mu.Unlock()
	defer source.mu.Unlock()

	// It's a new file that wasn't load'ed yet. ss.mu is still held which will
	// block concurrent loading.
	if pkg := e.packages[sk.pkg]; pkg != nil {
	if f, err := pkg.Open(sk.relpath); err == nil {
	var d []byte
	if d, err = io.ReadAll(f); err == nil {
	oldsk := th.Local("shac.pkg").(sourceKey)
	th.SetLocal("shac.pkg", sk)
	fp := syntax.FilePortion{Content: d, FirstLine: 1, FirstCol: 1}
	source.globals, source.err = starlark.ExecFileOptions(e.opts, th, sk.String(), fp, e.globals)
	th.SetLocal("shac.pkg", oldsk)
	var errl resolve.ErrorList
	if errors.As(source.err, &errl) {
	// Unwrap the error, only keep the first one.
	source.err = errl[0]
	}
	var errre resolve.Error
	if errors.As(source.err, &errre) {
	// Synthesize a BacktraceableError since it's nicer for the
	// user. Sadly we can't get the function context even if the
	// error is within a function implementation, so hardcode
	// "<toplevel>".
	source.err = &failure{
	Message: errre.Msg,
	Stack: starlark.CallStack{
	starlark.CallFrame{
	Name: "<toplevel>",
	Pos: errre.Pos,
	},
	},
	}
	}
	} else {
	source.err = err
	}
	if err = f.Close(); source.err == nil {
	source.err = err
	}
	} else if errors.Is(err, fs.ErrNotExist) {
	// Hide the underlying error for determinism.
	source.err = fmt.Errorf("%s not found", sk.relpath)
	} else {
	source.err = err
	}
	} else {
	source.err = errors.New("package not found")
	}
	return source.globals, source.err
	}