From e1523e2699545201211442713ddb2c0d82c6a01e Mon Sep 17 00:00:00 2001
From: Garet Halliday <ghalliday@gfxlabs.io>
Date: Tue, 20 Sep 2022 15:02:56 -0500
Subject: [PATCH] that's hot

---
 go.mod                       |  5 ++---
 go.sum                       | 11 ++++-------
 lib/gat/pool/session/pool.go | 19 +++++--------------
 3 files changed, 11 insertions(+), 24 deletions(-)

diff --git a/go.mod b/go.mod
index 6779a323..a4a0b5e5 100644
--- a/go.mod
+++ b/go.mod
@@ -4,8 +4,8 @@ go 1.19
 
 require (
 	gfx.cafe/util/go/bufpool v0.0.0-20220906091724-3a24b7f40ccf
+	gfx.cafe/util/go/generic v0.0.0-20220917152604-80373e5a2c51
 	gfx.cafe/util/go/graceful v0.0.0-20220913082111-9770431e98e9
-	gfx.cafe/util/go/lambda v0.0.0-20220906200602-98a6b35a1b42
 	git.tuxpa.in/a/zlog v1.32.0
 	github.com/BurntSushi/toml v1.2.0
 	github.com/auxten/postgresql-parser v1.0.1
@@ -28,6 +28,7 @@ require (
 	github.com/getsentry/raven-go v0.2.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/google/go-cmp v0.5.8 // indirect
 	github.com/grpc-ecosystem/grpc-gateway v1.16.0 // indirect
 	github.com/konsorten/go-windows-terminal-sequences v1.0.3 // indirect
 	github.com/kr/pretty v0.2.0 // indirect
@@ -38,10 +39,8 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/sirupsen/logrus v1.6.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	github.com/stretchr/testify v1.7.2 // indirect
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
 	github.com/xdg-go/stringprep v1.0.3 // indirect
-	golang.org/x/exp v0.0.0-20220827204233-334a2380cb91 // indirect
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c // indirect
 	golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10 // indirect
 	golang.org/x/text v0.3.7 // indirect
diff --git a/go.sum b/go.sum
index ad65268b..892c276c 100644
--- a/go.sum
+++ b/go.sum
@@ -2,11 +2,10 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 gfx.cafe/util/go/bufpool v0.0.0-20220906091724-3a24b7f40ccf h1:ya4IK1D+Kq0DrFdrrZ7tjmp3BgoO4v5sCAeUytR6j1U=
 gfx.cafe/util/go/bufpool v0.0.0-20220906091724-3a24b7f40ccf/go.mod h1:+DiyiCOBGS9O9Ce4ewHQO3Y59h66WSWAbgZZ2O2AYYw=
+gfx.cafe/util/go/generic v0.0.0-20220917152604-80373e5a2c51 h1:4FEqqQR2Ruae4quQQn7lqLtLvQaH+dHmehpPHhiXm+Q=
+gfx.cafe/util/go/generic v0.0.0-20220917152604-80373e5a2c51/go.mod h1:0mgIMiX8+M1l4VHMaOqBaydZaztbFgblQAaEDOZpUhQ=
 gfx.cafe/util/go/graceful v0.0.0-20220913082111-9770431e98e9 h1:Z0RRPRbz5Bt7yu6v+RzLCAKm0YNcrwNGLMs9UVs8NsU=
 gfx.cafe/util/go/graceful v0.0.0-20220913082111-9770431e98e9/go.mod h1:sO44FAgBZXic9FwJaJHX1mI9vt1e3CRe0X/3bwnMRho=
-gfx.cafe/util/go/lambda v0.0.0-20220906200602-98a6b35a1b42 h1:8mKA+jVj7l3sM/s6CjqF/5DhAIWftIhit0XqkswzkAg=
-gfx.cafe/util/go/lambda v0.0.0-20220906200602-98a6b35a1b42/go.mod h1:+qAj+4kl6uABj/3RqR84AwHnm+vYyYoPwLJdTR3GgRc=
-git.tuxpa.in/a/lambda v0.0.0-20220903040836-72d2bd6dc070 h1:CuEIvPmhTWKmMMoJnsP2xOoWRFOFB7cZABONrSQaiDc=
 git.tuxpa.in/a/zlog v1.32.0 h1:KKXbRF1x8kJDSzUoGz/pivo+4TVY6xT5sVtdFZ6traY=
 git.tuxpa.in/a/zlog v1.32.0/go.mod h1:vUa2Qhu6DLPLqmfRy99FiPqaY2eb6/KQjtMekW3UNnA=
 github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
@@ -113,6 +112,7 @@ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.1/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg=
+github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-querystring v1.0.0/go.mod h1:odCYkC5MyYFN7vkCjXpyrEuKhc/BUO6wN/zVPAxq5ck=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
@@ -227,8 +227,7 @@ github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXf
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.7.2 h1:4jaiDzPyXQvSd7D0EjG45355tLlV3VOECpq10pLC+8s=
-github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals=
+github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/ugorji/go v1.1.4/go.mod h1:uQMGLiO92mf5W77hV/PUCpI3pbzQx3CRekS0kk+RGrc=
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
 github.com/urfave/negroni v1.0.0/go.mod h1:Meg73S6kFm/4PpbYdq35yYWoCZ9mS/YSx+lKnmiohz4=
@@ -258,8 +257,6 @@ golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4/go.mod h1:yigFU9vqHzYiE8U
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
-golang.org/x/exp v0.0.0-20220827204233-334a2380cb91 h1:tnebWN09GYg9OLPss1KXj8txwZc6X6uMr6VFdcGNbHw=
-golang.org/x/exp v0.0.0-20220827204233-334a2380cb91/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
diff --git a/lib/gat/pool/session/pool.go b/lib/gat/pool/session/pool.go
index 20c56aea..7db5b7f3 100644
--- a/lib/gat/pool/session/pool.go
+++ b/lib/gat/pool/session/pool.go
@@ -5,8 +5,8 @@ import (
 	"gfx.cafe/gfx/pggat/lib/config"
 	"gfx.cafe/gfx/pggat/lib/gat"
 	"gfx.cafe/gfx/pggat/lib/gat/protocol"
+	"gfx.cafe/util/go/generic"
 	"runtime"
-	"sync"
 	"sync/atomic"
 )
 
@@ -17,11 +17,9 @@ type Pool struct {
 
 	dialer gat.Dialer
 
-	assigned map[gat.ClientID]gat.Connection
+	assigned generic.Map[gat.ClientID, gat.Connection]
 
 	servers chan gat.Connection
-
-	mu sync.Mutex
 }
 
 func New(database gat.Database, dialer gat.Dialer, conf *config.Pool, user *config.User) *Pool {
@@ -31,8 +29,6 @@ func New(database gat.Database, dialer gat.Dialer, conf *config.Pool, user *conf
 
 		dialer: dialer,
 
-		assigned: make(map[gat.ClientID]gat.Connection),
-
 		servers: make(chan gat.Connection, 1+runtime.NumCPU()*4),
 	}
 	p.EnsureConfig(conf)
@@ -55,13 +51,11 @@ func (p *Pool) returnConnection(c gat.Connection) {
 }
 
 func (p *Pool) getOrAssign(client gat.Client) gat.Connection {
-	p.mu.Lock()
-	defer p.mu.Unlock()
 	cid := client.GetId()
-	c, ok := p.assigned[cid]
+	c, ok := p.assigned.Load(cid)
 	if !ok {
 		get := p.getConnection()
-		p.assigned[cid] = get
+		p.assigned.Store(cid, get)
 		return get
 	}
 	return c
@@ -76,14 +70,11 @@ func (p *Pool) EnsureConfig(c *config.Pool) {
 }
 
 func (p *Pool) OnDisconnect(client gat.Client) {
-	p.mu.Lock()
-	defer p.mu.Unlock()
 	cid := client.GetId()
-	c, ok := p.assigned[cid]
+	c, ok := p.assigned.LoadAndDelete(cid)
 	if !ok {
 		return
 	}
-	delete(p.assigned, cid)
 	p.servers <- c
 }
 
-- 
GitLab