diff --git a/src/bin/pg_tde_archive_decrypt.c b/src/bin/pg_tde_archive_decrypt.c
index 786919a3..92c2f54f 100644
--- a/src/bin/pg_tde_archive_decrypt.c
+++ b/src/bin/pg_tde_archive_decrypt.c
@@ -8,6 +8,8 @@
 #include "access/pg_tde_fe_init.h"
 #include "access/pg_tde_xlog_smgr.h"
 
+#include <signal.h>
+
 #define TMPFS_DIRECTORY "/dev/shm"
 
 static bool
@@ -134,6 +136,27 @@ usage(const char *progname)
 			 "\n"), progname, progname);
 }
 
+static char tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_archiveXXXXXX";
+static char tmppath[MAXPGPATH];
+
+static void
+cleanup_tmpdir(void)
+{
+	if (unlink(tmppath) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove file \"%s\": %m", tmppath);
+	if (rmdir(tmpdir) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
+}
+
+/* Clean up temporary files and dirs on SIGINT/SIGTERM */
+static void
+trapsig(int signum)
+{
+	cleanup_tmpdir();
+	signal(signum, SIG_DFL);
+	kill(getpid(), signum);
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -143,8 +166,6 @@ main(int argc, char *argv[])
 	char	   *command;
 	char	   *sep;
 	char	   *sourcename;
-	char		tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_archiveXXXXXX";
-	char		tmppath[MAXPGPATH];
 	bool		issegment;
 	int			rc;
 
@@ -199,6 +220,11 @@ main(int argc, char *argv[])
 		s = stpcpy(s, "/");
 		stpcpy(s, sourcename);
 
+		/* There is a mostly harmless race condition with creating the dir */
+		signal(SIGTERM, trapsig);
+		signal(SIGINT, trapsig);
+		atexit(cleanup_tmpdir);
+
 		command = replace_percent_placeholders(command,
 											   "ARCHIVE-COMMAND", "fp",
 											   targetname, tmppath);
@@ -228,13 +254,5 @@ main(int argc, char *argv[])
 
 	free(command);
 
-	if (issegment)
-	{
-		if (unlink(tmppath) < 0)
-			pg_log_warning("could not remove file \"%s\": %m", tmppath);
-		if (rmdir(tmpdir) < 0)
-			pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
-	}
-
 	return 0;
 }
diff --git a/src/bin/pg_tde_restore_encrypt.c b/src/bin/pg_tde_restore_encrypt.c
index 1d5c51f4..16c60480 100644
--- a/src/bin/pg_tde_restore_encrypt.c
+++ b/src/bin/pg_tde_restore_encrypt.c
@@ -8,6 +8,8 @@
 #include "access/pg_tde_fe_init.h"
 #include "access/pg_tde_xlog_smgr.h"
 
+#include <signal.h>
+
 #define TMPFS_DIRECTORY "/dev/shm"
 
 /*
@@ -129,6 +131,27 @@ usage(const char *progname)
 			 "\n"), progname, progname);
 }
 
+static char tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_restoreXXXXXX";
+static char tmppath[MAXPGPATH];
+
+static void
+cleanup_tmpdir(void)
+{
+	if (unlink(tmppath) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove file \"%s\": %m", tmppath);
+	if (rmdir(tmpdir) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
+}
+
+/* Clean up temporary files and dirs on SIGINT/SIGTERM */
+static void
+trapsig(int signum)
+{
+	cleanup_tmpdir();
+	signal(signum, SIG_DFL);
+	kill(getpid(), signum);
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -138,8 +161,6 @@ main(int argc, char *argv[])
 	char	   *command;
 	char	   *sep;
 	char	   *targetname;
-	char		tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_restoreXXXXXX";
-	char		tmppath[MAXPGPATH];
 	bool		issegment;
 	int			rc;
 
@@ -194,6 +215,11 @@ main(int argc, char *argv[])
 		s = stpcpy(s, "/");
 		stpcpy(s, targetname);
 
+		/* There is a mostly harmless race condition with creating the dir */
+		signal(SIGTERM, trapsig);
+		signal(SIGINT, trapsig);
+		atexit(cleanup_tmpdir);
+
 		command = replace_percent_placeholders(command,
 											   "RESTORE-COMMAND", "fp",
 											   sourcename, tmppath);
@@ -222,14 +248,7 @@ main(int argc, char *argv[])
 	free(command);
 
 	if (issegment)
-	{
 		write_encrypted_segment(targetpath, sourcename, tmppath);
 
-		if (unlink(tmppath) < 0)
-			pg_log_warning("could not remove file \"%s\": %m", tmppath);
-		if (rmdir(tmpdir) < 0)
-			pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
-	}
-
 	return 0;
 }