使用GO对PostgreSQL进行有意思的多线程压测

前言

针对PostgreSQL进行压缩，有很多相关的工具。有同学又要问了，为何还要再搞一个？比如，pgbench, sysbench之类的，已经很强大了。是的，它们都很强大。但有时候，在一些特殊的场景，可能自己构造一个更能接近真实的生产环境。

这里，我半写，半借助于ChatGPT，搞出一个代码片段来模拟启动一段多线程并发SQL请求，作用于PostgreSQL数据库。然后，你可以对请求执行完以后的结果进行观测，尤其是表膨胀，受影响记录条数之类的。

基于此，我们还可以进行持续改造，快速用于工作之中。

实作

需求：

实现一段代码，读取一个sql文件，然后分段分批执行，并且是以多线程（比如10个线程，go里边可能就是协程，非常高效）去执行这个SQL中的所有SQL语句。再加一个时间限制，比如持续执行120秒。

实现:

package mainimport ("bufio""context""database/sql""fmt""io""os""strings""sync""time"_ "github.com/lib/pq"
)const (host     = "localhost"port     =  5555user     = "postgres"password = "password"dbname   = "mydb"
)func execute_sqls(ctx context.Context, sqls []string, wg *sync.WaitGroup, thread int) {defer wg.Done()psqlInfo := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable", host, port, user, password, dbname)db, err := sql.Open("postgres", psqlInfo)if err != nil {panic(err)}defer db.Close()start := time.Now()for {for _, sql := range sqls {select {case <-ctx.Done():elapsed := time.Since(start)fmt.Printf("Thread %d stopped. It executed SQLs for %s \n", thread, elapsed)returndefault:_, err := db.Exec(sql)if err != nil {fmt.Println(err)}}}}
}func read_sqls(file string) []string {f, err := os.Open(file)if err != nil {panic(err)}defer f.Close()sqls := make([]string, 0)r := bufio.NewReader(f)for {line, err := r.ReadString(';')if err == io.EOF {break} else if err != nil {panic(err)}sql := strings.TrimSpace(line)if sql != "" {sqls = append(sqls, sql)}}return sqls
}func main() {filepath := "file.sql" // Replace with your file pathnumThreads := 10  // Number of threadssqls := read_sqls(filepath)var wg sync.WaitGroupctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) // 60 secondsfor i := 0; i < numThreads; i++ {wg.Add(1)go execute_sqls(ctx, sqls, &wg, i)}wg.Wait()cancel()fmt.Println("All goroutines stopped")
}

上边的代码，关于输入文件：file.sql,  线程数：10， 运行时间：60秒，都是硬编码进去的。你可以根据实际情况，进行参数化。

体验：

在你的go环境已经安装了"github.com/lib/pq"等必备包之后（go get github.com/lib/pq)，就可以直接执行了。我们准备一个pg的基本环境。database: mydb,  端口：5555, 就用postgres用户及相应密码（仅用于测试目的），不缀述。

目标表的准备：

\c mydb
create table t(id int, col2 varchar(32));

file.sql文件内容如下：

insert into t values ((10000*random())::int, md5(random()::varchar));
with updates as (select (10000*random())::int as id) update t set col2 = 'update' || updates.id from updates where t.id=updates.id returning updates.id;

这个测试的代码片段，就是插入一条随机记录，并且再随机更新一条记录，使用CTE语法，把对应的id值返回来，有可能找不到对应的记录，就返回的是空值。在并发大的情况下，update语句慢慢就起作用了。这样就可以反复执行。

来看看效果：

go run ./stress.gohread 7 stopped. It executed SQLs for 59.999324s 
Thread 1 stopped. It executed SQLs for 1m0.000756208s 
Thread 2 stopped. It executed SQLs for 1m0.000604792s 
Thread 4 stopped. It executed SQLs for 1m0.001703583s 
Thread 0 stopped. It executed SQLs for 1m0.008518875s 
Thread 9 stopped. It executed SQLs for 1m0.008456083s 
Thread 5 stopped. It executed SQLs for 1m0.007964375s 
Thread 6 stopped. It executed SQLs for 1m0.007968292s 
Thread 3 stopped. It executed SQLs for 1m0.008145042s 
Thread 8 stopped. It executed SQLs for 1m0.008202209s 
All goroutines stopped

1分钟跑完之后，我们看到这样的部分记录结果：

mydb=# select * from t limit 10;id  |    col2
------+------------4792 | update47923416 | update34169290 | update9290887 | update8878778 | update87787472 | update74724602 | update46023454 | update34542604 | update26041990 | update1990
(10 rows)

总记录条数：

mydb=# select count(*) from t;count
--------126056
(1 row)

引申：可以认为单个C+U操作，10个线程并发，1分钟入库12.6万。

表大小：

mydb=# select pg_total_relation_size('t');pg_total_relation_size
------------------------8372224
(1 row)

使用下边的SQL看看相关指标：

WITH cteTableInfo AS 
(SELECT COUNT(1) AS ct,SUM(length(t::text)) AS TextLength  ,'public.t'::regclass AS TableName  FROM public.t AS t  
)
,cteRowSize AS 
(SELECT ARRAY [pg_relation_size(TableName), pg_relation_size(TableName, 'vm'), pg_relation_size(TableName, 'fsm'), pg_table_size(TableName), pg_indexes_size(TableName), pg_total_relation_size(TableName), TextLength] AS val, ARRAY ['Relation Size', 'Visibility Map', 'Free Space Map', 'Table Included Toast Size', 'Indexes Size', 'Total Relation Size', 'Live Row Byte Size'] AS NameFROM cteTableInfo
)
SELECT unnest(name) AS Description,unnest(val) AS Bytes,pg_size_pretty(unnest(val)) AS BytesPretty,unnest(val) / ct AS bytes_per_row
FROM cteTableInfo, cteRowSizeUNION ALL SELECT '------------------------------', NULL, NULL, NULL
UNION ALL SELECT 'TotalRows', ct, NULL, NULL FROM cteTableInfo
UNION ALL SELECT 'LiveTuples', pg_stat_get_live_tuples(TableName), NULL, NULL FROM cteTableInfo
UNION ALL SELECT 'DeadTuples', pg_stat_get_dead_tuples(TableName), NULL, NULL FROM cteTableInfo;

结果:

          description           |  bytes  | bytespretty | bytes_per_row
--------------------------------+---------+-------------+---------------Relation Size                  | 8339456 | 8144 kB     |            66Visibility Map                 |    8192 | 8192 bytes  |             0Free Space Map                 |   24576 | 24 kB       |             0Table Included Toast Size      | 8372224 | 8176 kB     |            66Indexes Size                   |       0 | 0 bytes     |             0Total Relation Size            | 8372224 | 8176 kB     |            66Live Row Byte Size             | 2338451 | 2284 kB     |            18------------------------------ |         |             |TotalRows                      |  126056 |             |LiveTuples                     |  126056 |             |DeadTuples                     |   12274 |             |
(11 rows)

里边有涉及到的死元组为12274行。

mydb=# create extension pgstattuple;
CREATE EXTENSIONmydb=# select * from pgstattuple('public.t') \gx
-[ RECORD 1 ]------+--------
table_len          | 8339456
tuple_count        | 126056
tuple_len          | 5125646
tuple_percent      | 61.46
dead_tuple_count   | 12110
dead_tuple_len     | 483172
dead_tuple_percent | 5.79
free_space         | 1451672
free_percent       | 17.41

这两种统计结果也都比较接近。

当你针对相同的表，进行随机多次测试，发现上边的值也会不断变化（update的命中率会越来越高）。

总结：

本文的目的，只是作一个抛砖引玉，可以随时使用go, python甚至rust去构建一个小的压缩环境，对各种复杂的压力环境进行模拟，并得出相关结论。当然，作为一个团队，可以开发出使用Java之类的接近业务逻辑的工具也是可以的。也有的测试团队，原意使用JMeter + jdbc来构建测试套集，都不失为一种方式。这类工具，是介于pgbench 和 真实业务场景压测之间的一种使用方式。哪个更方便，就可以用哪个。

上边的代码片段，稍加改造，就可以用到实际的实验当中。

关于表膨胀，可以看看我前边的文章：

也聊聊PostgreSQL中的空间膨胀与AutoVacuum

PG中的一例简单的update看表膨胀