github.com/grafana/pyroscope@v1.18.0/pkg/metastore/tracing/context_registry_test.go (about)

     1  package tracing
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  type contextKey string
    14  
    15  func TestContextRegistry_StoreAndRetrieve(t *testing.T) {
    16  	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
    17  	defer r.Shutdown()
    18  
    19  	ctx := context.WithValue(context.Background(), contextKey("key"), "value")
    20  	r.Store("id-1", ctx)
    21  
    22  	retrieved, found := r.Retrieve("id-1")
    23  	require.True(t, found)
    24  	assert.Equal(t, "value", retrieved.Value(contextKey("key")))
    25  }
    26  
    27  func TestContextRegistry_RetrieveNotFound(t *testing.T) {
    28  	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
    29  	defer r.Shutdown()
    30  
    31  	retrieved, found := r.Retrieve("nonexistent-id")
    32  	require.False(t, found)
    33  	assert.Equal(t, context.Background(), retrieved)
    34  }
    35  
    36  func TestContextRegistry_Delete(t *testing.T) {
    37  	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
    38  	defer r.Shutdown()
    39  
    40  	ctx := context.WithValue(context.Background(), contextKey("key"), "value")
    41  	r.Store("id-1", ctx)
    42  
    43  	_, found := r.Retrieve("id-1")
    44  	require.True(t, found)
    45  
    46  	r.Delete("id-1")
    47  
    48  	_, found = r.Retrieve("id-1")
    49  	require.False(t, found)
    50  }
    51  
    52  func TestContextRegistry_Cleanup(t *testing.T) {
    53  	// Use short TTL for a faster test
    54  	r := newContextRegistry(100*time.Millisecond, 200*time.Millisecond, nil)
    55  	defer r.Shutdown()
    56  
    57  	ctx := context.WithValue(context.Background(), contextKey("key"), "value")
    58  	r.Store("id-1", ctx)
    59  
    60  	_, found := r.Retrieve("id-1")
    61  	require.True(t, found)
    62  	assert.Equal(t, 1, r.Size())
    63  
    64  	time.Sleep(400 * time.Millisecond)
    65  
    66  	_, found = r.Retrieve("id-1")
    67  	require.False(t, found)
    68  	assert.Equal(t, 0, r.Size())
    69  }
    70  
    71  func TestContextRegistry_Size(t *testing.T) {
    72  	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
    73  	defer r.Shutdown()
    74  
    75  	assert.Equal(t, 0, r.Size())
    76  
    77  	ctx := context.Background()
    78  	r.Store("id-1", ctx)
    79  	r.Store("id-2", ctx)
    80  	r.Store("id-3", ctx)
    81  
    82  	assert.Equal(t, 3, r.Size())
    83  
    84  	r.Delete("id-2")
    85  	assert.Equal(t, 2, r.Size())
    86  }
    87  
    88  func TestContextRegistry_ConcurrentAccess(t *testing.T) {
    89  	r := newContextRegistry(1*time.Second, 5*time.Second, nil)
    90  	defer r.Shutdown()
    91  
    92  	done := make(chan bool)
    93  
    94  	for i := 0; i < 10; i++ {
    95  		go func(id int) {
    96  			ctx := context.WithValue(context.Background(), contextKey("id"), id)
    97  			for j := 0; j < 100; j++ {
    98  				r.Store(fmt.Sprintf("%d-%d", id, j), ctx)
    99  				time.Sleep(1 * time.Millisecond)
   100  			}
   101  			done <- true
   102  		}(i)
   103  	}
   104  
   105  	for i := 0; i < 10; i++ {
   106  		go func(id int) {
   107  			for j := 0; j < 100; j++ {
   108  				r.Retrieve(fmt.Sprintf("%d-%d", id, j))
   109  				time.Sleep(1 * time.Millisecond)
   110  			}
   111  			done <- true
   112  		}(i)
   113  	}
   114  
   115  	for i := 0; i < 20; i++ {
   116  		<-done
   117  	}
   118  
   119  	assert.True(t, r.Size() > 0)
   120  }