github.com/emcfarlane/larking@v0.0.0-20220605172417-1704b45ee6c3/starlib/starlarkthread/thread.go (about)

     1  // Copyright 2021 Edward McFarlane. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package starlarkthread
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"io"
    11  	"runtime"
    12  	"sync"
    13  	"testing"
    14  
    15  	"github.com/emcfarlane/larking/starlib/starlarkstruct"
    16  	"go.starlark.net/starlark"
    17  )
    18  
    19  func NewModule() *starlarkstruct.Module {
    20  	return &starlarkstruct.Module{
    21  		Name: "thread",
    22  		Members: starlark.StringDict{
    23  			"os":   starlark.String(runtime.GOOS),
    24  			"arch": starlark.String(runtime.GOARCH),
    25  		},
    26  	}
    27  }
    28  
    29  const ctxkey = "context"
    30  
    31  // SetContext sets the thread context.
    32  func SetContext(thread *starlark.Thread, ctx context.Context) {
    33  	thread.SetLocal(ctxkey, ctx)
    34  }
    35  
    36  // GetContext gets the thread context or returns a TODO context.
    37  // ctx := starlarkthread.Context(thread)
    38  func GetContext(thread *starlark.Thread) context.Context {
    39  	if ctx, ok := thread.Local(ctxkey).(context.Context); ok {
    40  		return ctx
    41  	}
    42  	return context.TODO()
    43  }
    44  
    45  // Resource is a starlark.Value that requires close handling.
    46  type Resource interface {
    47  	starlark.Value
    48  	io.Closer
    49  }
    50  
    51  const rsckey = "resources"
    52  
    53  // ResourceStore is a thread local storage map for adding resources.
    54  // It is thread safe.
    55  type ResourceStore struct {
    56  	mu sync.Mutex
    57  	m  map[Resource]bool // resource whether it's living
    58  }
    59  
    60  // WithResourceStore returns a cleanup function. It is required for
    61  // packages that add resources.
    62  func WithResourceStore(thread *starlark.Thread) func() error {
    63  	store := &ResourceStore{
    64  		m: make(map[Resource]bool),
    65  	}
    66  	SetResourceStore(thread, store)
    67  	// TODO: runtime.SetFinalizer?
    68  	return func() error { return CloseResources(thread) }
    69  }
    70  
    71  func SetResourceStore(thread *starlark.Thread, store *ResourceStore) {
    72  	thread.SetLocal(rsckey, store)
    73  }
    74  func GetResourceStore(thread *starlark.Thread) (*ResourceStore, error) {
    75  	store, ok := thread.Local(rsckey).(*ResourceStore)
    76  	if !ok {
    77  		return nil, fmt.Errorf("thread missing resource store")
    78  	}
    79  	return store, nil
    80  }
    81  
    82  func AddResource(thread *starlark.Thread, rsc Resource) error {
    83  	store, err := GetResourceStore(thread)
    84  	if err != nil {
    85  		return err
    86  	}
    87  	store.mu.Lock()
    88  	store.m[rsc] = true
    89  	store.mu.Unlock()
    90  	return nil
    91  }
    92  
    93  func (s *ResourceStore) Close() (firstErr error) {
    94  	s.mu.Lock()
    95  	defer s.mu.Unlock()
    96  
    97  	for rsc, open := range s.m {
    98  		if !open {
    99  			continue
   100  		}
   101  		if err := rsc.Close(); err != nil && firstErr == nil {
   102  			// TODO: chain errors?
   103  			firstErr = err
   104  		}
   105  		delete(s.m, rsc)
   106  	}
   107  	return
   108  }
   109  
   110  func CloseResources(thread *starlark.Thread) (firstErr error) {
   111  	store, err := GetResourceStore(thread)
   112  	if err != nil {
   113  		return err
   114  	}
   115  	return store.Close()
   116  }
   117  
   118  // AssertOption implements starlarkassert.TestOption
   119  // Add like so:
   120  //
   121  // 	starlarkassert.RunTests(t, "*.star", globals, starlarkthread.AssertOption)
   122  //
   123  func AssertOption(t testing.TB, thread *starlark.Thread) func() {
   124  	close := WithResourceStore(thread)
   125  	return func() {
   126  		if err := close(); err != nil {
   127  			t.Error(err, "failed to close resources")
   128  		}
   129  	}
   130  }