transport: pool, store connections using atomic.Pointer[Conn]

Kulezi · Kulezi · commit 4c1e97e2924f · 2022-09-12T12:56:34.000+02:00
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/google/go-cmp v0.5.6
 	github.com/klauspost/compress v1.15.1
 	github.com/pierrec/lz4/v4 v4.1.14
-	go.uber.org/atomic v1.9.0
+	go.uber.org/atomic v1.10.0
 	go.uber.org/goleak v1.1.12
 )
 
diff --git a/go.sum b/go.sum
@@ -21,6 +21,8 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
 go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
+go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
 go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
 go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
diff --git a/transport/export_test.go b/transport/export_test.go
@@ -2,8 +2,8 @@ package transport
 
 func (p *ConnPool) AllConns() []*Conn {
 	var conns = make([]*Conn, len(p.conns))
-	for i, v := range p.conns {
-		conns[i], _ = v.Load().(*Conn)
+	for i := range conns {
+		conns[i] = p.loadConn(i)
 	}
 	return conns
 }
diff --git a/transport/pool.go b/transport/pool.go
@@ -19,7 +19,7 @@ type ConnPool struct {
 	host         string
 	nrShards     int
 	msbIgnore    uint8
-	conns        []atomic.Value
+	conns        []atomic.Pointer[Conn]
 	connClosedCh chan int // notification channel for when connection is closed
 	connObs      ConnObserver
 }
@@ -99,13 +99,11 @@ func (p *ConnPool) storeConn(conn *Conn) {
 }
 
 func (p *ConnPool) loadConn(shard int) *Conn {
-	conn, _ := p.conns[shard].Load().(*Conn)
-	return conn
+	return p.conns[shard].Load()
 }
 
 func (p *ConnPool) clearConn(shard int) bool {
-	conn, _ := p.conns[shard].Swap((*Conn)(nil)).(*Conn)
-	return conn != nil
+	return p.conns[shard].Swap(nil) != nil
 }
 
 func (p *ConnPool) Close() {
@@ -115,7 +113,7 @@ func (p *ConnPool) Close() {
 // closeAll is called by PoolRefiller.
 func (p *ConnPool) closeAll() {
 	for i := range p.conns {
-		if conn, ok := p.conns[i].Swap((*Conn)(nil)).(*Conn); ok {
+		if conn := p.conns[i].Swap(nil); conn != nil {
 			conn.Close()
 		}
 	}
@@ -168,7 +166,7 @@ func (r *PoolRefiller) init(ctx context.Context, host string) error {
 		host:         host,
 		nrShards:     int(ss.NrShards),
 		msbIgnore:    ss.MsbIgnore,
-		conns:        make([]atomic.Value, int(ss.NrShards)),
+		conns:        make([]atomic.Pointer[Conn], int(ss.NrShards)),
 		connClosedCh: make(chan int, int(ss.NrShards)+1),
 		connObs:      r.cfg.ConnObserver,
 	}

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@ require (`
`6`	`6`	`github.com/google/go-cmp v0.5.6`
`7`	`7`	`github.com/klauspost/compress v1.15.1`
`8`	`8`	`github.com/pierrec/lz4/v4 v4.1.14`
`9`		`- go.uber.org/atomic v1.9.0`
	`9`	`+ go.uber.org/atomic v1.10.0`
`10`	`10`	`go.uber.org/goleak v1.1.12`
`11`	`11`	`)`
`12`	`12`
Original file line number	Diff line number	Diff line change
`@@ -2,8 +2,8 @@ package transport`
`2`	`2`
`3`	`3`	`func (p ConnPool) AllConns() []Conn {`
`4`	`4`	`var conns = make([]*Conn, len(p.conns))`
`5`		`- for i, v := range p.conns {`
`6`		`- conns[i], _ = v.Load().(*Conn)`
	`5`	`+ for i := range conns {`
	`6`	`+ conns[i] = p.loadConn(i)`
`7`	`7`	`}`
`8`	`8`	`return conns`
`9`	`9`	`}`
Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@ type ConnPool struct {`
`19`	`19`	`host string`
`20`	`20`	`nrShards int`
`21`	`21`	`msbIgnore uint8`
`22`		`- conns []atomic.Value`
	`22`	`+ conns []atomic.Pointer[Conn]`
`23`	`23`	`connClosedCh chan int // notification channel for when connection is closed`
`24`	`24`	`connObs ConnObserver`
`25`	`25`	`}`
`@@ -99,13 +99,11 @@ func (p ConnPool) storeConn(conn Conn) {`
`99`	`99`	`}`
`100`	`100`
`101`	`101`	`func (p ConnPool) loadConn(shard int) Conn {`
`102`		`- conn, _ := p.conns[shard].Load().(*Conn)`
`103`		`- return conn`
	`102`	`+ return p.conns[shard].Load()`
`104`	`103`	`}`
`105`	`104`
`106`	`105`	`func (p *ConnPool) clearConn(shard int) bool {`
`107`		`- conn, _ := p.conns[shard].Swap((Conn)(nil)).(Conn)`
`108`		`- return conn != nil`
	`106`	`+ return p.conns[shard].Swap(nil) != nil`
`109`	`107`	`}`
`110`	`108`
`111`	`109`	`func (p *ConnPool) Close() {`
`@@ -115,7 +113,7 @@ func (p *ConnPool) Close() {`
`115`	`113`	`// closeAll is called by PoolRefiller.`
`116`	`114`	`func (p *ConnPool) closeAll() {`
`117`	`115`	`for i := range p.conns {`
`118`		`- if conn, ok := p.conns[i].Swap((Conn)(nil)).(Conn); ok {`
	`116`	`+ if conn := p.conns[i].Swap(nil); conn != nil {`
`119`	`117`	`conn.Close()`
`120`	`118`	`}`
`121`	`119`	`}`
`@@ -168,7 +166,7 @@ func (r *PoolRefiller) init(ctx context.Context, host string) error {`
`168`	`166`	`host: host,`
`169`	`167`	`nrShards: int(ss.NrShards),`
`170`	`168`	`msbIgnore: ss.MsbIgnore,`
`171`		`- conns: make([]atomic.Value, int(ss.NrShards)),`
	`169`	`+ conns: make([]atomic.Pointer[Conn], int(ss.NrShards)),`
`172`	`170`	`connClosedCh: make(chan int, int(ss.NrShards)+1),`
`173`	`171`	`connObs: r.cfg.ConnObserver,`
`174`	`172`	`}`