github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/crdb/pool/fake_test.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/google/uuid"
     8  )
     9  
    10  var (
    11  	_ balanceablePool[*FakePoolConn[*FakeConn], *FakeConn] = &FakePool{}
    12  	_ balancePoolConn[*FakeConn]                           = &FakePoolConn[*FakeConn]{}
    13  )
    14  
    15  type FakeConn struct {
    16  	closed bool
    17  }
    18  
    19  func NewFakeConn() *FakeConn {
    20  	return &FakeConn{closed: false}
    21  }
    22  
    23  func (f FakeConn) IsClosed() bool {
    24  	return f.closed
    25  }
    26  
    27  type FakePoolConn[C balanceConn] struct {
    28  	conn C
    29  }
    30  
    31  func (f FakePoolConn[C]) Conn() C {
    32  	return f.conn
    33  }
    34  
    35  func (f FakePoolConn[C]) Release() {}
    36  
    37  type FakePool struct {
    38  	sync.Mutex
    39  	id          string
    40  	maxConns    uint32
    41  	gc          map[*FakeConn]uint32
    42  	nodeForConn map[*FakeConn]uint32
    43  }
    44  
    45  func NewFakePool(maxConns uint32) *FakePool {
    46  	return &FakePool{
    47  		id:          uuid.NewString(),
    48  		maxConns:    maxConns,
    49  		gc:          make(map[*FakeConn]uint32, 0),
    50  		nodeForConn: make(map[*FakeConn]uint32, 0),
    51  	}
    52  }
    53  
    54  func (f *FakePool) ID() string {
    55  	return f.id
    56  }
    57  
    58  func (f *FakePool) AcquireAllIdle(_ context.Context) []*FakePoolConn[*FakeConn] {
    59  	conns := make([]*FakePoolConn[*FakeConn], 0, len(f.nodeForConn))
    60  	f.Range(func(conn *FakeConn, value uint32) {
    61  		conns = append(conns, &FakePoolConn[*FakeConn]{conn: conn})
    62  	})
    63  	return conns
    64  }
    65  
    66  func (f *FakePool) Node(conn *FakeConn) uint32 {
    67  	f.Lock()
    68  	defer f.Unlock()
    69  	id := f.nodeForConn[conn]
    70  	return id
    71  }
    72  
    73  func (f *FakePool) GC(conn *FakeConn) {
    74  	f.Lock()
    75  	defer f.Unlock()
    76  	id := f.nodeForConn[conn]
    77  	delete(f.nodeForConn, conn)
    78  	f.gc[conn] = id
    79  }
    80  
    81  func (f *FakePool) MaxConns() uint32 {
    82  	return f.maxConns
    83  }
    84  
    85  func (f *FakePool) Range(fn func(conn *FakeConn, nodeID uint32)) {
    86  	f.Lock()
    87  	defer f.Unlock()
    88  	for k, v := range f.nodeForConn {
    89  		fn(k, v)
    90  	}
    91  }