From 8a0466e4c2d45521dc4b0ffabe07c89cec869bf7 Mon Sep 17 00:00:00 2001
From: kotsachin <kotsachin@gmail.com>
Date: Mon, 25 May 2015 18:14:47 +0900
Subject: [PATCH] Some improvements and fixes to previously submitted pull
 request for cleaning temporary objects

---
 bin/pg_repack.c |  5 ++++-
 bin/pgut/pgut.c | 21 +++++++++++++++------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/bin/pg_repack.c b/bin/pg_repack.c
index 7105f53..84f5dac 100644
--- a/bin/pg_repack.c
+++ b/bin/pg_repack.c
@@ -243,7 +243,7 @@ static char *
 utoa(unsigned int value, char *buffer)
 {
 	sprintf(buffer, "%u", value);
-	return strdup(buffer);
+	return pgut_strdup(buffer);
 }
 
 static pgut_option options[] =
@@ -1391,6 +1391,7 @@ repack_one_table(repack_table *table, const char *orderby)
 	params[1] = utoa(temp_obj_num, buffer);
 	command("SELECT repack.repack_drop($1, $2)", 2, params);
 	command("COMMIT", 0, NULL);
+	temp_obj_num = 0; /* reset temporary object counter after cleanup */
 
 	/*
 	 * 7. Analyze.
@@ -1706,6 +1707,7 @@ repack_cleanup_callback(bool fatal, void *userdata)
 
 		reconnect(ERROR);
 		command("SELECT repack.repack_drop($1, $2)", 2, params);
+		temp_obj_num = 0; /* reset temporary object counter after cleanup */
 	}
 }
 
@@ -1734,6 +1736,7 @@ repack_cleanup(bool fatal, const repack_table *table)
 		params[0] = utoa(table->target_oid, buffer);
 		params[1] =  utoa(temp_obj_num, buffer);
 		command("SELECT repack.repack_drop($1, $2)", 2, params);
+		temp_obj_num = 0; /* reset temporary object counter after cleanup */
 	}
 }
 
diff --git a/bin/pgut/pgut.c b/bin/pgut/pgut.c
index 37cc301..e533bef 100644
--- a/bin/pgut/pgut.c
+++ b/bin/pgut/pgut.c
@@ -64,7 +64,7 @@ static void on_before_exec(pgutConn *conn);
 static void on_after_exec(pgutConn *conn);
 static void on_interrupt(void);
 static void on_cleanup(void);
-static void exit_or_abort(int exitcode);
+static void exit_or_abort(int exitcode, int elevel);
 
 void
 pgut_init(int argc, char **argv)
@@ -872,7 +872,10 @@ pgut_errfinish(int dummy, ...)
 			edata->detail.data);
 
 	if (pgut_abort_level <= edata->elevel && edata->elevel <= PANIC)
-		exit_or_abort(edata->code);
+	{
+		in_cleanup = true; /* need to be set for cleaning temporary objects on error */
+		exit_or_abort(edata->code, edata->elevel);
+	}
 }
 
 #ifndef PGUT_OVERRIDE_ELOG
@@ -1195,13 +1198,19 @@ on_cleanup(void)
 }
 
 static void
-exit_or_abort(int exitcode)
+exit_or_abort(int exitcode, int elevel)
 {
-	call_atexit_callbacks(true);
-	if (in_cleanup)
+	
+	if (in_cleanup && FATAL > elevel)
 	{
 		/* oops, error in cleanup*/
-		in_cleanup = false;
+		call_atexit_callbacks(true);
+		exit(exitcode);
+	}
+	else if (FATAL <= elevel <= PANIC)
+	{
+		/* on FATAL or PANIC */
+		call_atexit_callbacks(true);
 		abort();
 	}
 	else