Skip to content

Commit

Permalink
Merge pull request #1 from shogo82148/handle-panics
Browse files Browse the repository at this point in the history
handle panics
  • Loading branch information
shogo82148 authored Aug 13, 2022
2 parents 4a642d5 + 7d9fd08 commit 54b9d53
Show file tree
Hide file tree
Showing 2 changed files with 131 additions and 19 deletions.
107 changes: 88 additions & 19 deletions memoize.go
Original file line number Diff line number Diff line change
@@ -1,14 +1,46 @@
package memoize

import (
"bytes"
"context"
"errors"
"fmt"
"runtime/debug"
"sync"
"time"
)

// for testing
var nowFunc = time.Now

// errGoexit indicates the runtime.Goexit was called in
// the user given function.
var errGoexit = errors.New("runtime.Goexit was called")

// A panicError is an arbitrary value recovered from a panic
// with the stack trace during the execution of given function.
type panicError struct {
value interface{}
stack []byte
}

// Error implements error interface.
func (p *panicError) Error() string {
return fmt.Sprintf("%v\n\n%s", p.value, p.stack)
}

func newPanicError(v interface{}) error {
stack := debug.Stack()

// The first line of the stack trace is of the form "goroutine N [status]:"
// but by the time the panic reaches Do the goroutine may no longer exist
// and its status will have changed. Trim out the misleading line.
if line := bytes.IndexByte(stack[:], '\n'); line >= 0 {
stack = stack[line+1:]
}
return &panicError{value: v, stack: stack}
}

// Group memoizes the calls of Func with expiration.
type Group[K comparable, V any] struct {
mu sync.Mutex // protects m
Expand Down Expand Up @@ -93,28 +125,65 @@ func (g *Group[K, V]) Do(ctx context.Context, key K, fn func(ctx context.Context
}

func do[K comparable, V any](g *Group[K, V], e *entry[V], c *call[V], key K, fn func(ctx context.Context, key K) (V, time.Time, error)) {
defer c.cancel()
var ret result[V]

v, expiresAt, err := fn(c.ctx, key)
ret := result[V]{
val: v,
expiresAt: expiresAt,
err: err,
}
normalReturn := false
recovered := false

// save to the cache
e.mu.Lock()
e.call = nil // to avoid adding new channels to c.chans
chans := c.chans
if err == nil {
e.val = v
e.expiresAt = expiresAt
}
e.mu.Unlock()
// use double-defer to distinguish panic from runtime.Goexit,
// more details see https://golang.org/cl/134395
defer func() {
// the given function invoked runtime.Goexit
if !normalReturn && !recovered {
ret.err = errGoexit
}

// save to the cache
e.mu.Lock()
e.call = nil // to avoid adding new channels to c.chans
chans := c.chans
if ret.err == nil {
e.val = ret.val
e.expiresAt = ret.expiresAt
}
e.mu.Unlock()

if e, ok := ret.err.(*panicError); ok {
panic(e)
} else if ret.err == errGoexit {
// Already in the process of goexit, no need to call again
} else {
// Normal return
// notify the result to the callers
for _, ch := range chans {
ch <- ret
}
}
}()

func() {
defer func() {
c.cancel()
if !normalReturn {
// Ideally, we would wait to take a stack trace until we've determined
// whether this is a panic or a runtime.Goexit.
//
// Unfortunately, the only way we can distinguish the two is to see
// whether the recover stopped the goroutine from terminating, and by
// the time we know that, the part of the stack trace relevant to the
// panic has been discarded.
if r := recover(); r != nil {
ret.err = newPanicError(r)
}
}
}()

ret.val, ret.expiresAt, ret.err = fn(c.ctx, key)
normalReturn = true
}()

// notify the result to the callers
for _, ch := range chans {
ch <- ret
if !normalReturn {
recovered = true
}
}

Expand Down
43 changes: 43 additions & 0 deletions memoize_test.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,13 @@
package memoize

import (
"bytes"
"context"
"errors"
"os"
"os/exec"
"runtime"
"strings"
"sync"
"sync/atomic"
"testing"
Expand Down Expand Up @@ -244,6 +249,44 @@ func TestDoContext(t *testing.T) {
}
}

func TestPanicDo(t *testing.T) {
if runtime.GOOS == "js" {
t.Skipf("js does not support exec")
}

if os.Getenv("TEST_PANIC_DO") != "" {
var g Group[string, int]
fn := func(ctx context.Context, _ string) (int, time.Time, error) {
panic("Panicking in Do")
}
g.Do(context.Background(), "key", fn)
t.Fatalf("Do unexpectedly returned")
}

t.Parallel()

cmd := exec.Command(os.Args[0], "-test.run="+t.Name(), "-test.v")
cmd.Env = append(os.Environ(), "TEST_PANIC_DO=1")
out := new(bytes.Buffer)
cmd.Stdout = out
cmd.Stderr = out
if err := cmd.Start(); err != nil {
t.Fatal(err)
}

err := cmd.Wait()
t.Logf("%s:\n%s", strings.Join(cmd.Args, " "), out)
if err == nil {
t.Errorf("Test subprocess passed; want a crash due to panic in DoChan")
}
if bytes.Contains(out.Bytes(), []byte("Do unexpectedly returned")) {
t.Errorf("Test subprocess failed with an unexpected failure mode.")
}
if !bytes.Contains(out.Bytes(), []byte("Panicking in Do")) {
t.Errorf("Test subprocess failed, but the crash isn't caused by panicking in Do")
}
}

func benchmarkDo(parallelism int) func(b *testing.B) {
return func(b *testing.B) {
b.SetParallelism(parallelism)
Expand Down

0 comments on commit 54b9d53

Please sign in to comment.