Skip to content

Commit b01733c

Browse files
CitoCito
committed
inserttable: optimize memory allocation
1 parent f090e58 commit b01733c

3 files changed

Lines changed: 63 additions & 44 deletions

File tree

‎ext/pgconn.c‎

Lines changed: 19 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -783,17 +783,15 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
783783

784784
encoding = PQclientEncoding(self->cnx);
785785

786-
/* starts query */
787-
ext_char_buffer_s(&buffer, "copy ");
788-
789-
/* return early if there is no buffer */
790-
if (buffer.error) {
791-
if (buffer.data)
792-
PyMem_Free(buffer.data);
786+
/* pre-allocate some memory for the query buffer */
787+
if (!init_char_buffer(&buffer, 4096)) {
793788
Py_DECREF(iter_row);
794789
return PyErr_NoMemory();
795790
}
796791

792+
/* starts query */
793+
ext_char_buffer_s(&buffer, "copy ");
794+
797795
s = table;
798796
do {
799797
t = strchr(s, '.');
@@ -838,13 +836,14 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
838836
col = PQescapeIdentifier(self->cnx, col, (size_t)slen);
839837
Py_DECREF(obj);
840838
ext_char_buffer_s(&buffer, col);
841-
ext_char_buffer_c(&buffer, j == n - 1 ? ')' : ',');
842839
PQfreemem(col);
840+
ext_char_buffer_c(&buffer, j == n - 1 ? ')' : ',');
843841
}
844842
}
845843
ext_char_buffer_s(&buffer, " from stdin");
846844
if (freeze)
847845
ext_char_buffer_s(&buffer, " freeze");
846+
ext_char_buffer_c(&buffer, '\0');
848847

849848
if (buffer.error) {
850849
PyMem_Free(buffer.data);
@@ -897,10 +896,10 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
897896
return NULL;
898897
}
899898

900-
/* reset buffer to empty */
901-
buffer.len = 0;
899+
/* empty buffer while keeping allocated memory */
900+
buffer.size = 0;
902901

903-
/* builds insert line */
902+
/* build insert line */
904903

905904
for (j = 0; j < n; ++j) {
906905
if (j)
@@ -1015,22 +1014,21 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
10151014
}
10161015
Py_DECREF(s);
10171016
}
1018-
1019-
if (buffer.error) {
1020-
PQputCopyEnd(self->cnx, "Memory error");
1021-
PyMem_Free(buffer.data);
1022-
Py_DECREF(columns);
1023-
Py_DECREF(iter_row);
1024-
return PyErr_NoMemory();
1025-
}
10261017
}
10271018

10281019
Py_DECREF(columns);
10291020

1021+
/* terminate line */
10301022
ext_char_buffer_c(&buffer, '\n');
1023+
if (buffer.error) {
1024+
PQputCopyEnd(self->cnx, "Memory error");
1025+
PyMem_Free(buffer.data);
1026+
Py_DECREF(iter_row);
1027+
return PyErr_NoMemory();
1028+
}
10311029

1032-
/* sends data */
1033-
ret = PQputCopyData(self->cnx, buffer.data, (int)buffer.len);
1030+
/* send data */
1031+
ret = PQputCopyData(self->cnx, buffer.data, (int)buffer.size);
10341032
if (ret != 1) {
10351033
char *errormsg = ret == -1 ? PQerrorMessage(self->cnx)
10361034
: "Data cannot be queued";

‎ext/pginternal.c‎

Lines changed: 40 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1494,22 +1494,42 @@ notice_receiver(void *arg, const PGresult *res)
14941494
PyGILState_Release(gstate);
14951495
}
14961496

1497-
/* Extend char buffer with given string */
1497+
/* Pre-allocate some memory for a char buffer and return success status. */
1498+
static int
1499+
init_char_buffer(struct CharBuffer *buf, size_t initial_size)
1500+
{
1501+
buf->size = 0;
1502+
buf->data = PyMem_Malloc(initial_size);
1503+
if (buf->data) {
1504+
buf->max_size = initial_size;
1505+
buf->error = 0;
1506+
}
1507+
else {
1508+
buf->max_size = 0;
1509+
buf->error = 1;
1510+
}
1511+
return !buf->error;
1512+
}
1513+
1514+
/* Extend char buffer with given string.
1515+
Note: We do not assume or guarantee that the buffer is zero-terminated. */
14981516
static void
14991517
ext_char_buffer_s(struct CharBuffer *buf, const char *s)
15001518
{
1501-
size_t len = strlen(s);
1502-
size_t need = buf->len + len + 1;
1519+
size_t len = strlen(s), need;
15031520

15041521
if (!len || buf->error)
15051522
return;
15061523

1507-
if (need >= buf->max_len) {
1524+
if ((need = buf->size + len) >= buf->max_size) {
15081525
void *tmp;
15091526

1510-
// Allocate powers of two unless it's large
1511-
if (2 * buf->max_len >= need && buf->max_len < 1024 * 1024)
1512-
need = 2 * buf->max_len;
1527+
if (buf->max_size < 1024 * 1024) {
1528+
/* allocate powers of two unless it's large */
1529+
size_t double_size = 2 * buf->max_size;
1530+
if (double_size >= need) /*/ overflow check */
1531+
need = double_size;
1532+
}
15131533

15141534
tmp = PyMem_Realloc(buf->data, need);
15151535
if (!tmp) {
@@ -1518,26 +1538,27 @@ ext_char_buffer_s(struct CharBuffer *buf, const char *s)
15181538
}
15191539

15201540
buf->data = tmp;
1521-
buf->max_len = need;
1541+
buf->max_size = need;
15221542
}
15231543

1524-
memcpy(buf->data + buf->len, s, len + 1);
1525-
buf->len += len;
1544+
memcpy(buf->data + buf->size, s, len);
1545+
buf->size += len;
15261546
}
15271547

15281548
/* Extend char buffer with given character */
15291549
static void
15301550
ext_char_buffer_c(struct CharBuffer *buf, char c)
15311551
{
1532-
if (buf->len > buf->max_len - 2) {
1533-
// slow path dealing with reallocation
1534-
char tmp[2] = {c, '\0'};
1552+
if (buf->error)
1553+
return;
1554+
1555+
if (buf->size >= buf->max_size) { /* buffer is full? */
1556+
/* slow path dealing with reallocation */
1557+
char tmp[2] = {c ? c : '\n', '\0'}; /* allow adding a zero-byte */
15351558
ext_char_buffer_s(buf, tmp);
1559+
if (!c)
1560+
buf->data[buf->size - 1] = '\0'; /* fix zero-byte */
15361561
}
1537-
else {
1538-
if (buf->error)
1539-
return;
1540-
buf->data[buf->len++] = c;
1541-
buf->data[buf->len] = '\0';
1542-
}
1562+
else
1563+
buf->data[buf->size++] = c;
15431564
}

‎ext/pgmodule.c‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -163,10 +163,10 @@ typedef struct {
163163
The buffer can be extended with the extend_char_buffer_s/x() functions.
164164
*/
165165
struct CharBuffer {
166-
char *data; /* actual string data */
167-
size_t len; /* strlen() of data */
168-
size_t max_len; /* allocated size */
169-
int error; /* error flag (invalid string)*/
166+
char *data; /* actual string data */
167+
size_t size; /* current size of data */
168+
size_t max_size; /* allocated size */
169+
int error; /* error flag (invalid data) */
170170
};
171171

172172
/* Internal functions */

0 commit comments

Comments
 (0)