Move DialFunc to Dial

kylecarbs · kylecarbs · commit 0af778936b26 · 2021-07-21T00:04:13.000Z
diff --git a/wsnet/cache.go b/wsnet/cache.go
@@ -9,16 +9,15 @@ import (
 )
 
 // dialerFunc is used to reference a dialer returned for caching.
-type dialerFunc func(ctx context.Context, key string) (*Dialer, error)
+type dialerFunc func() (*Dialer, error)
 
 // DialCache constructs a new DialerCache.
 // The cache clears connections that:
 // 1. Are older than the TTL and have no active user-created connections.
 // 2. Have been closed.
-func DialCache(ttl time.Duration, dialer dialerFunc) *DialerCache {
+func DialCache(ttl time.Duration) *DialerCache {
 	dc := &DialerCache{
 		ttl:         ttl,
-		dialerFunc:  dialer,
 		closed:      make(chan struct{}),
 		flightGroup: &singleflight.Group{},
 		mut:         sync.RWMutex{},
@@ -30,7 +29,6 @@ func DialCache(ttl time.Duration, dialer dialerFunc) *DialerCache {
 }
 
 type DialerCache struct {
-	dialerFunc  dialerFunc
 	ttl         time.Duration
 	flightGroup *singleflight.Group
 
@@ -98,7 +96,7 @@ func (d *DialerCache) evict() {
 
 // Dial returns a Dialer from the cache if one exists with the key provided,
 // or dials a new connection using the dialerFunc.
-func (d *DialerCache) Dial(ctx context.Context, key string) (*Dialer, bool, error) {
+func (d *DialerCache) Dial(ctx context.Context, key string, dialerFunc func() (*Dialer, error)) (*Dialer, bool, error) {
 	d.mut.RLock()
 	if dialer, ok := d.dialers[key]; ok {
 		closed := false
@@ -119,7 +117,7 @@ func (d *DialerCache) Dial(ctx context.Context, key string) (*Dialer, bool, erro
 	d.mut.RUnlock()
 
 	dialer, err, _ := d.flightGroup.Do(key, func() (interface{}, error) {
-		dialer, err := d.dialerFunc(ctx, key)
+		dialer, err := dialerFunc()
 		if err != nil {
 			return nil, err
 		}
diff --git a/wsnet/cache_test.go b/wsnet/cache_test.go
@@ -10,19 +10,23 @@ import (
 )
 
 func TestCache(t *testing.T) {
+	dialFunc := func(connectAddr string) func() (*Dialer, error) {
+		return func() (*Dialer, error) {
+			return DialWebsocket(context.Background(), connectAddr, nil)
+		}
+	}
+
 	t.Run("Caches", func(t *testing.T) {
 		connectAddr, listenAddr := createDumbBroker(t)
 		l, err := Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "")
 		require.NoError(t, err)
 		defer l.Close()
 
-		cache := DialCache(time.Hour, func(ctx context.Context, key string) (*Dialer, error) {
-			return DialWebsocket(ctx, connectAddr, nil)
-		})
-		_, cached, err := cache.Dial(context.Background(), "example")
+		cache := DialCache(time.Hour)
+		_, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr))
 		require.NoError(t, err)
 		require.Equal(t, cached, false)
-		_, cached, err = cache.Dial(context.Background(), "example")
+		_, cached, err = cache.Dial(context.Background(), "example", dialFunc(connectAddr))
 		require.NoError(t, err)
 		require.Equal(t, cached, true)
 	})
@@ -33,15 +37,13 @@ func TestCache(t *testing.T) {
 		require.NoError(t, err)
 		defer l.Close()
 
-		cache := DialCache(time.Hour, func(ctx context.Context, key string) (*Dialer, error) {
-			return DialWebsocket(ctx, connectAddr, nil)
-		})
+		cache := DialCache(time.Hour)
 
-		conn, cached, err := cache.Dial(context.Background(), "example")
+		conn, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr))
 		require.NoError(t, err)
 		require.Equal(t, cached, false)
 		require.NoError(t, conn.Close())
-		_, cached, err = cache.Dial(context.Background(), "example")
+		_, cached, err = cache.Dial(context.Background(), "example", dialFunc(connectAddr))
 		require.NoError(t, err)
 		require.Equal(t, cached, false)
 	})
@@ -52,15 +54,13 @@ func TestCache(t *testing.T) {
 		require.NoError(t, err)
 		defer l.Close()
 
-		cache := DialCache(0, func(ctx context.Context, key string) (*Dialer, error) {
-			return DialWebsocket(ctx, connectAddr, nil)
-		})
+		cache := DialCache(0)
 
-		_, cached, err := cache.Dial(context.Background(), "example")
+		_, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr))
 		require.NoError(t, err)
 		require.Equal(t, cached, false)
 		cache.evict()
-		_, cached, err = cache.Dial(context.Background(), "example")
+		_, cached, err = cache.Dial(context.Background(), "example", dialFunc(connectAddr))
 		require.NoError(t, err)
 		require.Equal(t, cached, false)
 	})
