prometheus_race_test.go

Documentation: github.com/cybertec-postgresql/pgwatch/v6/internal/sinks

     1  package sinks
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
     9  	"github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
    10  	"github.com/prometheus/client_golang/prometheus"
    11  )
    12  
    13  func TestCollect_RaceCondition_Real(_ *testing.T) {
    14  	// 1. Initialize the real PrometheusWriter
    15  	promw, _ := NewPrometheusWriter(testutil.TestContext, "127.0.0.1:0/")
    16  	promw.Println("init done")
    17  
    18  	// 2. Register a metric so Write() actually puts data into the map
    19  	_ = promw.SyncMetric("race_db", "test_metric", AddOp)
    20  
    21  	var wg sync.WaitGroup
    22  	done := make(chan struct{})
    23  
    24  	// --- The Writer (Simulating Database Updates) ---
    25  	wg.Go(func() {
    26  		for {
    27  			select {
    28  			case <-done:
    29  				return
    30  			default:
    31  				// Call the REAL Write method
    32  				_ = promw.Write(metrics.MeasurementEnvelope{
    33  					DBName:     "race_db",
    34  					MetricName: "test_metric",
    35  					Data: metrics.Measurements{
    36  						{
    37  							metrics.EpochColumnName: time.Now().UnixNano(),
    38  							"value":                 int64(100),
    39  						},
    40  					},
    41  				})
    42  				// No sleep here -> hammer the map as fast as possible
    43  			}
    44  		}
    45  	})
    46  
    47  	// --- The Collector (Simulating Prometheus Scrapes) ---
    48  	wg.Go(func() {
    49  		// Prometheus provides a channel to receive metrics
    50  		ch := make(chan prometheus.Metric, 10000)
    51  
    52  		// Scrape 50 times (more than enough to trigger a race in a tight loop)
    53  		for range 50 {
    54  			// Call the REAL Collect method
    55  			promw.Collect(ch)
    56  
    57  			// Drain the channel so it doesn't block
    58  		drainLoop:
    59  			for {
    60  				select {
    61  				case <-ch:
    62  				default:
    63  					break drainLoop
    64  				}
    65  			}
    66  		}
    67  		close(done) // Tell the writer to stop
    68  	})
    69  
    70  	wg.Wait()
    71  	_ = promw.SyncMetric("race_db", "test_metric", DeleteOp)
    72  }
    73  
    74  func TestGaugesMap_RaceCondition(_ *testing.T) {
    75  	// 1. Initialize PrometheusWriter
    76  	promw, _ := NewPrometheusWriter(testutil.TestContext, "127.0.0.1:0/pgwatch")
    77  
    78  	// 2. Register a metric so Write() actually puts data into the map
    79  	_ = promw.SyncMetric("race_db", "test_metric", AddOp)
    80  
    81  	// 3. Pre-fill cache so Collect has something to do
    82  	_ = promw.Write(metrics.MeasurementEnvelope{
    83  		DBName:     "race_db",
    84  		MetricName: "test_metric",
    85  		Data: metrics.Measurements{
    86  			{
    87  				metrics.EpochColumnName: time.Now().UnixNano(),
    88  				"value":                 int64(100),
    89  			},
    90  		},
    91  	})
    92  
    93  	var wg sync.WaitGroup
    94  	done := make(chan struct{})
    95  
    96  	// --- The Config Reloader (Simulating configuration updates) ---
    97  	wg.Go(func() {
    98  		for {
    99  			select {
   100  			case <-done:
   101  				return
   102  			default:
   103  				// Call the REAL DefineMetrics method (Writes to gauges map)
   104  				_ = promw.DefineMetrics(&metrics.Metrics{
   105  					MetricDefs: metrics.MetricDefs{
   106  						"test_metric": {Gauges: []string{"value"}},
   107  					},
   108  				})
   109  			}
   110  		}
   111  	})
   112  
   113  	// --- The Collector (Simulating Prometheus Scrapes) ---
   114  	wg.Go(func() {
   115  		// Prometheus provides a channel to receive metrics
   116  		ch := make(chan prometheus.Metric, 10000)
   117  
   118  		// Scrape 50 times (more than enough to trigger a race in a tight loop)
   119  		for range 50 {
   120  			// Call the REAL Collect method (Reads from gauges map)
   121  			promw.Collect(ch)
   122  
   123  			// Drain the channel so it doesn't block
   124  		drainLoop:
   125  			for {
   126  				select {
   127  				case <-ch:
   128  				default:
   129  					break drainLoop
   130  				}
   131  			}
   132  		}
   133  		close(done) // Tell the reloader to stop
   134  	})
   135  
   136  	wg.Wait()
   137  	_ = promw.SyncMetric("race_db", "", DeleteOp)
   138  }
   139