From c5b2db6e49cb8e5742557aeb920892bed59ca624 Mon Sep 17 00:00:00 2001 From: Marko Kreen Date: Sun, 23 May 2010 11:04:11 +0300 Subject: [PATCH] convert few static buffers into pktbuf --- include/bouncer.h | 6 ++-- include/pktbuf.h | 6 ++++ include/proto.h | 2 +- src/admin.c | 48 +++++++++++++++++------------- src/janitor.c | 3 ++ src/loader.c | 36 ++++++++++++++--------- src/objects.c | 7 ++++- src/pktbuf.c | 25 ++++++++++++++-- src/proto.c | 74 +++++++++++++++++++++++++++-------------------- src/server.c | 9 ++++-- src/varcache.c | 20 ++++++------- 11 files changed, 150 insertions(+), 86 deletions(-) diff --git a/include/bouncer.h b/include/bouncer.h index a6ab8c7..277f58d 100644 --- a/include/bouncer.h +++ b/include/bouncer.h @@ -180,8 +180,7 @@ struct PgPool { PgStats older_stats; /* database info to be sent to client */ - uint8_t welcome_msg[STARTUP_BUF]; /* ServerParams without VarCache ones */ - unsigned welcome_msg_len; + struct PktBuf *welcome_msg; /* ServerParams without VarCache ones */ VarCache orig_vars; /* default params from server */ @@ -236,8 +235,7 @@ struct PgDatabase { bool db_auto; /* is the database auto-created by autodb_connstr */ bool admin; /* internal console db */ - uint8_t startup_params[STARTUP_BUF]; /* partial StartupMessage (without user) be sent to server */ - unsigned startup_params_len; + struct PktBuf *startup_params; /* partial StartupMessage (without user) be sent to server */ PgUser *forced_user; /* if not NULL, the user/psw is forced */ diff --git a/include/pktbuf.h b/include/pktbuf.h index 638aaf7..465c6d7 100644 --- a/include/pktbuf.h +++ b/include/pktbuf.h @@ -41,6 +41,12 @@ struct PktBuf { PktBuf *pktbuf_dynamic(int start_len) _MUSTCHECK; void pktbuf_static(PktBuf *buf, uint8_t *data, int len); +void pktbuf_free(PktBuf *buf); + +void pktbuf_reset(struct PktBuf *pkt); +struct PktBuf *pktbuf_temp(void); + + /* * sending */ diff --git a/include/proto.h b/include/proto.h index f42aefd..74dc2d7 100644 --- a/include/proto.h +++ b/include/proto.h @@ -38,7 +38,7 @@ bool send_pooler_error(PgSocket *client, bool send_ready, const char *msg) /*_M void log_server_error(const char *note, PktHdr *pkt); void parse_server_error(PktHdr *pkt, const char **level_p, const char **msg_p); -void add_welcome_parameter(PgPool *pool, const char *key, const char *val); +bool add_welcome_parameter(PgPool *pool, const char *key, const char *val) _MUSTCHECK; void finish_welcome_msg(PgSocket *server); bool welcome_client(PgSocket *client) _MUSTCHECK; diff --git a/src/admin.c b/src/admin.c index 7940493..4b38e53 100644 --- a/src/admin.c +++ b/src/admin.c @@ -234,18 +234,19 @@ static bool send_one_fd(PgSocket *admin, { struct msghdr msg; struct cmsghdr *cmsg; - int res; struct iovec iovec; - uint8_t pktbuf[STARTUP_BUF * 2]; + int res; uint8_t cntbuf[CMSG_SPACE(sizeof(int))]; - iovec.iov_base = pktbuf; - BUILD_DataRow(res, pktbuf, sizeof(pktbuf), "issssiqissss", + struct PktBuf *pkt = pktbuf_temp(); + + pktbuf_write_DataRow(pkt, "issssiqissss", fd, task, user, db, addr, port, ckey, link, client_enc, std_strings, datestyle, timezone); - if (res < 0) + if (pkt->failed) return false; - iovec.iov_len = res; + iovec.iov_base = pkt->buf; + iovec.iov_len = pktbuf_written(pkt); /* sending fds */ memset(&msg, 0, sizeof(msg)); @@ -1142,7 +1143,7 @@ void admin_setup(void) PgDatabase *db; PgPool *pool; PgUser *user; - PktBuf msg; + PktBuf *msg; int res; /* fake database */ @@ -1169,21 +1170,28 @@ void admin_setup(void) fatal("cannot create admin user?"); /* prepare welcome */ - pktbuf_static(&msg, pool->welcome_msg, sizeof(pool->welcome_msg)); - pktbuf_write_AuthenticationOk(&msg); - pktbuf_write_ParameterStatus(&msg, "server_version", "8.0/bouncer"); - pktbuf_write_ParameterStatus(&msg, "client_encoding", "UNICODE"); - pktbuf_write_ParameterStatus(&msg, "server_encoding", "SQL_ASCII"); - pktbuf_write_ParameterStatus(&msg, "is_superuser", "on"); - - pool->welcome_msg_len = pktbuf_written(&msg); + msg = pktbuf_dynamic(128); + if (!msg) + fatal("cannot create admin welcome"); + pktbuf_write_AuthenticationOk(msg); + pktbuf_write_ParameterStatus(msg, "server_version", "8.0/bouncer"); + pktbuf_write_ParameterStatus(msg, "client_encoding", "UNICODE"); + pktbuf_write_ParameterStatus(msg, "server_encoding", "SQL_ASCII"); + pktbuf_write_ParameterStatus(msg, "is_superuser", "on"); + + if (msg->failed) + fatal("admin welcome failed"); + + pool->welcome_msg = msg; pool->welcome_msg_ready = 1; - pktbuf_static(&msg, db->startup_params, sizeof(db->startup_params)); - pktbuf_put_string(&msg, "database"); - db->dbname = (char *)db->startup_params + pktbuf_written(&msg); - pktbuf_put_string(&msg, "pgbouncer"); - db->startup_params_len = pktbuf_written(&msg); + msg = pktbuf_dynamic(128); + if (!msg) + fatal("cannot create admin startup pkt"); + db->startup_params = msg; + pktbuf_put_string(msg, "database"); + db->dbname = "pgbouncer"; + pktbuf_put_string(msg, db->dbname); /* initialize regexes */ res = regcomp(&rc_cmd, cmd_normal_rx, REG_EXTENDED | REG_ICASE); diff --git a/src/janitor.c b/src/janitor.c index 409167d..3fe49ba 100644 --- a/src/janitor.c +++ b/src/janitor.c @@ -576,6 +576,8 @@ static void kill_pool(PgPool *pool) close_server_list(&pool->tested_server_list, reason); close_server_list(&pool->new_server_list, reason); + pktbuf_free(pool->welcome_msg); + list_del(&pool->map_head); statlist_remove(&pool_list, &pool->head); varcache_clean(&pool->orig_vars); @@ -594,6 +596,7 @@ static void kill_database(PgDatabase *db) if (pool->db == db) kill_pool(pool); } + pktbuf_free(db->startup_params); if (db->forced_user) slab_free(user_cache, db->forced_user); if (db->connect_query) diff --git a/src/loader.c b/src/loader.c index 72be2fe..8f1c6b1 100644 --- a/src/loader.c +++ b/src/loader.c @@ -172,10 +172,11 @@ static void set_autodb(char *connstr) void parse_database(char *name, char *connstr) { char *p, *key, *val; - PktBuf buf; + PktBuf *msg; PgDatabase *db; int pool_size = -1; int res_pool_size = -1; + int dbname_ofs; char *dbname = name; char *host = NULL; @@ -336,29 +337,35 @@ void parse_database(char *name, char *connstr) /* assign connect_query */ set_connect_query(db, connect_query); - pktbuf_static(&buf, db->startup_params, sizeof(db->startup_params)); + if (db->startup_params) { + msg = db->startup_params; + pktbuf_reset(msg); + } else { + msg = pktbuf_dynamic(128); + if (!msg) + fatal("cannot allocate startup buf"); + db->startup_params = msg; + } - pktbuf_put_string(&buf, "database"); - db->dbname = (char *)db->startup_params + pktbuf_written(&buf); - pktbuf_put_string(&buf, dbname); + pktbuf_put_string(msg, "database"); + dbname_ofs = msg->write_pos; + pktbuf_put_string(msg, dbname); if (client_encoding) { - pktbuf_put_string(&buf, "client_encoding"); - pktbuf_put_string(&buf, client_encoding); + pktbuf_put_string(msg, "client_encoding"); + pktbuf_put_string(msg, client_encoding); } if (datestyle) { - pktbuf_put_string(&buf, "datestyle"); - pktbuf_put_string(&buf, datestyle); + pktbuf_put_string(msg, "datestyle"); + pktbuf_put_string(msg, datestyle); } if (timezone) { - pktbuf_put_string(&buf, "timezone"); - pktbuf_put_string(&buf, timezone); + pktbuf_put_string(msg, "timezone"); + pktbuf_put_string(msg, timezone); } - db->startup_params_len = pktbuf_written(&buf); - /* if user is forces, create fake object for it */ if (username != NULL) { if (!force_user(db, username, password)) @@ -366,6 +373,9 @@ void parse_database(char *name, char *connstr) } else if (db->forced_user) log_warning("losing forced user not supported," " keeping old setting"); + + /* remember dbname */ + db->dbname = (char *)msg->buf + dbname_ofs; } /* diff --git a/src/objects.c b/src/objects.c index 6251ef7..3a98a43 100644 --- a/src/objects.c +++ b/src/objects.c @@ -911,7 +911,8 @@ bool finish_client_login(PgSocket *client) fatal("bad client state"); } - if (!welcome_client(client)) { + /* check if we know server signature */ + if (!client->pool->welcome_msg_ready) { log_debug("finish_client_login: no welcome message, pause"); client->wait_for_welcome = 1; pause_client(client); @@ -921,6 +922,10 @@ bool finish_client_login(PgSocket *client) } client->wait_for_welcome = 0; + /* send the message */ + if (!welcome_client(client)) + return false; + slog_debug(client, "logged in"); return true; diff --git a/src/pktbuf.c b/src/pktbuf.c index be0a32e..bcd30cc 100644 --- a/src/pktbuf.c +++ b/src/pktbuf.c @@ -22,9 +22,9 @@ #include "bouncer.h" -static void pktbuf_free(PktBuf *buf) +void pktbuf_free(PktBuf *buf) { - if (buf->fixed_buf) + if (!buf || buf->fixed_buf) return; log_debug("pktbuf_free(%p)", buf); @@ -56,6 +56,15 @@ PktBuf *pktbuf_dynamic(int start_len) return buf; } +void pktbuf_reset(struct PktBuf *pkt) +{ + pkt->failed = 0; + pkt->write_pos = 0; + pkt->pktlen_pos = 0; + pkt->send_pos = 0; + pkt->sending = 0; +} + void pktbuf_static(PktBuf *buf, uint8_t *data, int len) { memset(buf, 0, sizeof(*buf)); @@ -64,6 +73,18 @@ void pktbuf_static(PktBuf *buf, uint8_t *data, int len) buf->fixed_buf = 1; } +struct PktBuf *pktbuf_temp(void) +{ + static PktBuf *temp_pktbuf; + + if (!temp_pktbuf) + temp_pktbuf = pktbuf_dynamic(512); + if (!temp_pktbuf) + fatal("failed to create temp pktbuf"); + pktbuf_reset(temp_pktbuf); + return temp_pktbuf; +} + bool pktbuf_send_immidiate(PktBuf *buf, PgSocket *sk) { int fd = sbuf_socket(&sk->sbuf); diff --git a/src/proto.c b/src/proto.c index 4ff15b3..3cc665c 100644 --- a/src/proto.c +++ b/src/proto.c @@ -168,25 +168,29 @@ void log_server_error(const char *note, PktHdr *pkt) */ /* add another server parameter packet to cache */ -void add_welcome_parameter(PgPool *pool, const char *key, const char *val) +bool add_welcome_parameter(PgPool *pool, const char *key, const char *val) { - PktBuf msg; + PktBuf *msg = pool->welcome_msg; if (pool->welcome_msg_ready) - return; + return true; - pktbuf_static(&msg, pool->welcome_msg + pool->welcome_msg_len, - sizeof(pool->welcome_msg) - pool->welcome_msg_len); + if (!msg) { + msg = pktbuf_dynamic(128); + if (!msg) + return false; + pool->welcome_msg = msg; + } /* first packet must be AuthOk */ - if (pool->welcome_msg_len == 0) - pktbuf_write_AuthenticationOk(&msg); + if (msg->write_pos == 0) + pktbuf_write_AuthenticationOk(msg); /* if not stored in ->orig_vars, write full packet */ - if (!varcache_set(&pool->orig_vars, key, val)) { - pktbuf_write_ParameterStatus(&msg, key, val); - pool->welcome_msg_len += pktbuf_written(&msg); - } + if (!varcache_set(&pool->orig_vars, key, val)) + pktbuf_write_ParameterStatus(msg, key, val); + + return !msg->failed; } /* all parameters processed */ @@ -201,30 +205,37 @@ void finish_welcome_msg(PgSocket *server) bool welcome_client(PgSocket *client) { int res; - uint8_t buf[STARTUP_BUF]; - PktBuf msg; PgPool *pool = client->pool; + const PktBuf *pmsg = pool->welcome_msg; + PktBuf *msg; slog_noise(client, "P: welcome_client"); - if (!pool->welcome_msg_ready) - return false; - pktbuf_static(&msg, buf, sizeof(buf)); - pktbuf_put_bytes(&msg, pool->welcome_msg, pool->welcome_msg_len); + /* copy prepared stuff around */ + msg = pktbuf_temp(); + pktbuf_put_bytes(msg, pmsg->buf, pmsg->write_pos); + /* fill vars */ varcache_fill_unset(&pool->orig_vars, client); - varcache_add_params(&msg, &client->vars); + varcache_add_params(msg, &client->vars); /* give each client its own cancel key */ get_random_bytes(client->cancel_key, 8); - pktbuf_write_BackendKeyData(&msg, client->cancel_key); - pktbuf_write_ReadyForQuery(&msg); + pktbuf_write_BackendKeyData(msg, client->cancel_key); - /* send all together */ - res = pktbuf_send_immidiate(&msg, client); - if (!res) - slog_warning(client, "unhandled failure to send welcome_msg"); + /* finish */ + pktbuf_write_ReadyForQuery(msg); + if (msg->failed) { + disconnect_client(client, true, "failed to prepare welcome message"); + return false; + } + /* send all together */ + res = pktbuf_send_immidiate(msg, client); + if (!res) { + disconnect_client(client, true, "failed to send welcome message"); + return false; + } return true; } @@ -330,14 +341,13 @@ bool send_startup_packet(PgSocket *server) { PgDatabase *db = server->pool->db; const char *username = server->pool->user->name; - PktBuf pkt; - uint8_t buf[STARTUP_BUF]; - - pktbuf_static(&pkt, buf, sizeof(buf)); - pktbuf_write_StartupMessage(&pkt, username, - db->startup_params, - db->startup_params_len); - return pktbuf_send_immidiate(&pkt, server); + PktBuf *pkt; + + pkt = pktbuf_temp(); + pktbuf_write_StartupMessage(pkt, username, + db->startup_params->buf, + db->startup_params->write_pos); + return pktbuf_send_immidiate(pkt, server); } int scan_text_result(struct MBuf *pkt, const char *tupdesc, ...) diff --git a/src/server.c b/src/server.c index 36b1e1f..01c6106 100644 --- a/src/server.c +++ b/src/server.c @@ -47,13 +47,18 @@ static bool load_parameter(PgSocket *server, PktHdr *pkt, bool startup) varcache_set(&client->vars, key, val); } - if (startup) - add_welcome_parameter(server->pool, key, val); + if (startup) { + if (!add_welcome_parameter(server->pool, key, val)) + goto failed_store; + } return true; failed: disconnect_server(server, true, "broken ParameterStatus packet"); return false; +failed_store: + disconnect_server(server, true, "failed to store ParameterStatus"); + return false; } /* we cannot log in at all, notify clients */ diff --git a/src/varcache.c b/src/varcache.c index ce3fc8e..54ead23 100644 --- a/src/varcache.c +++ b/src/varcache.c @@ -139,35 +139,33 @@ static int apply_var(PktBuf *pkt, const char *key, bool varcache_apply(PgSocket *server, PgSocket *client, bool *changes_p) { - PktBuf pkt; - uint8_t buf[STARTUP_BUF]; int changes = 0; struct PStr *cval, *sval; const struct var_lookup *lk; - uint8_t *debug_sql; + int sql_ofs; bool std_quote = is_std_quote(&server->vars); + struct PktBuf *pkt = pktbuf_temp(); - pktbuf_static(&pkt, buf, sizeof(buf)); - pktbuf_start_packet(&pkt, 'Q'); + pktbuf_start_packet(pkt, 'Q'); /* grab quory position inside pkt */ - debug_sql = pkt.buf + pkt.write_pos; + sql_ofs = pktbuf_written(pkt); for (lk = lookup; lk->name; lk++) { sval = get_value(&server->vars, lk); cval = get_value(&client->vars, lk); if (cval) - changes += apply_var(&pkt, lk->name, cval->str, sval->str, std_quote); + changes += apply_var(pkt, lk->name, cval->str, sval->str, std_quote); } *changes_p = changes > 0; if (!changes) return true; - pktbuf_put_char(&pkt, 0); - pktbuf_finish_packet(&pkt); + pktbuf_put_char(pkt, 0); + pktbuf_finish_packet(pkt); - slog_debug(server, "varcache_apply: %s", debug_sql); - return pktbuf_send_immidiate(&pkt, server); + slog_debug(server, "varcache_apply: %s", pkt->buf + sql_ofs); + return pktbuf_send_immidiate(pkt, server); } void varcache_fill_unset(VarCache *src, PgSocket *dst) -- 2.40.0