github.com/cockroachdb/tools@v0.0.0-20230222021103-a6d27438930d/internal/memoize/memoize_test.go (about)

     1  // Copyright 2019 The Go Authors. 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 memoize_test
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	"golang.org/x/tools/internal/memoize"
    14  )
    15  
    16  func TestGet(t *testing.T) {
    17  	var store memoize.Store
    18  
    19  	evaled := 0
    20  
    21  	h, release := store.Promise("key", func(context.Context, interface{}) interface{} {
    22  		evaled++
    23  		return "res"
    24  	})
    25  	defer release()
    26  	expectGet(t, h, "res")
    27  	expectGet(t, h, "res")
    28  	if evaled != 1 {
    29  		t.Errorf("got %v calls to function, wanted 1", evaled)
    30  	}
    31  }
    32  
    33  func expectGet(t *testing.T, h *memoize.Promise, wantV interface{}) {
    34  	t.Helper()
    35  	gotV, gotErr := h.Get(context.Background(), nil)
    36  	if gotV != wantV || gotErr != nil {
    37  		t.Fatalf("Get() = %v, %v, wanted %v, nil", gotV, gotErr, wantV)
    38  	}
    39  }
    40  
    41  func TestNewPromise(t *testing.T) {
    42  	calls := 0
    43  	f := func(context.Context, interface{}) interface{} {
    44  		calls++
    45  		return calls
    46  	}
    47  
    48  	// All calls to Get on the same promise return the same result.
    49  	p1 := memoize.NewPromise("debug", f)
    50  	expectGet(t, p1, 1)
    51  	expectGet(t, p1, 1)
    52  
    53  	// A new promise calls the function again.
    54  	p2 := memoize.NewPromise("debug", f)
    55  	expectGet(t, p2, 2)
    56  	expectGet(t, p2, 2)
    57  
    58  	// The original promise is unchanged.
    59  	expectGet(t, p1, 1)
    60  }
    61  
    62  func TestStoredPromiseRefCounting(t *testing.T) {
    63  	var store memoize.Store
    64  	v1 := false
    65  	v2 := false
    66  	p1, release1 := store.Promise("key1", func(context.Context, interface{}) interface{} {
    67  		return &v1
    68  	})
    69  	p2, release2 := store.Promise("key2", func(context.Context, interface{}) interface{} {
    70  		return &v2
    71  	})
    72  	expectGet(t, p1, &v1)
    73  	expectGet(t, p2, &v2)
    74  
    75  	expectGet(t, p1, &v1)
    76  	expectGet(t, p2, &v2)
    77  
    78  	p2Copy, release2Copy := store.Promise("key2", func(context.Context, interface{}) interface{} {
    79  		return &v1
    80  	})
    81  	if p2 != p2Copy {
    82  		t.Error("Promise returned a new value while old is not destroyed yet")
    83  	}
    84  	expectGet(t, p2Copy, &v2)
    85  
    86  	release2()
    87  	if got, want := v2, false; got != want {
    88  		t.Errorf("after destroying first v2 ref, got %v, want %v", got, want)
    89  	}
    90  	release2Copy()
    91  	if got, want := v1, false; got != want {
    92  		t.Errorf("after destroying v2, got %v, want %v", got, want)
    93  	}
    94  	release1()
    95  
    96  	p2Copy, release2Copy = store.Promise("key2", func(context.Context, interface{}) interface{} {
    97  		return &v2
    98  	})
    99  	if p2 == p2Copy {
   100  		t.Error("Promise returned previously destroyed value")
   101  	}
   102  	release2Copy()
   103  }
   104  
   105  func TestPromiseDestroyedWhileRunning(t *testing.T) {
   106  	// Test that calls to Promise.Get return even if the promise is destroyed while running.
   107  
   108  	var store memoize.Store
   109  	c := make(chan int)
   110  
   111  	var v int
   112  	h, release := store.Promise("key", func(ctx context.Context, _ interface{}) interface{} {
   113  		<-c
   114  		<-c
   115  		if err := ctx.Err(); err != nil {
   116  			t.Errorf("ctx.Err() = %v, want nil", err)
   117  		}
   118  		return &v
   119  	})
   120  
   121  	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) // arbitrary timeout; may be removed if it causes flakes
   122  	defer cancel()
   123  
   124  	var wg sync.WaitGroup
   125  	wg.Add(1)
   126  	var got interface{}
   127  	var err error
   128  	go func() {
   129  		got, err = h.Get(ctx, nil)
   130  		wg.Done()
   131  	}()
   132  
   133  	c <- 0    // send once to enter the promise function
   134  	release() // release before the promise function returns
   135  	c <- 0    // let the promise function proceed
   136  
   137  	wg.Wait()
   138  
   139  	if err != nil {
   140  		t.Errorf("Get() failed: %v", err)
   141  	}
   142  	if got != &v {
   143  		t.Errorf("Get() = %v, want %v", got, v)
   144  	}
   145  }
   146  
   147  func TestDoubleReleasePanics(t *testing.T) {
   148  	var store memoize.Store
   149  	_, release := store.Promise("key", func(ctx context.Context, _ interface{}) interface{} { return 0 })
   150  
   151  	panicked := false
   152  
   153  	func() {
   154  		defer func() {
   155  			if recover() != nil {
   156  				panicked = true
   157  			}
   158  		}()
   159  		release()
   160  		release()
   161  	}()
   162  
   163  	if !panicked {
   164  		t.Errorf("calling release() twice did not panic")
   165  	}
   166  }