/*-------------------------------------------------------------------------
 *
 * pg_snapclone_tools.c
 *
 *   Functions to publish and then clone a sessions transaction snapshot
 *
 * This module aims to support 8.1, 8.2 and 8.3 only. Support for pre-8.1
 * is possible, using the 8.1 version as a base, but is not provided.
 *
 * Note that the code differs a fair amount between releases since this
 * module uses a number of nifty Postgres features, but also because this
 * does some stuff with internals that mere mortals are not supposed to
 * play with.
 *
 * WARNING: Abuse of the code contained here can lead to SQL returning
 * wrong answers, provide security holes and worse. Treat these
 * functions with the caution and respect due to malignant biohazards.
 *
 * Original author simon@2ndquadrant.com, implementing the requirements
 * and original idea by Jan Wieck, for inclusion with the slony project.
 *
 * Copyright (c) 2008-, PostgreSQL Global Development Group
 *
 * Permission to use, copy, modify, and distribute this software and its 
 * documentation for any purpose, without fee, and without a written 
 * agreement is hereby granted, provided that the above copyright notice 
 * and this paragraph and the following two paragraphs appear in all copies.
 * 
 * IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA, AUTHOR OR COPYRIGHT
 * HOLDERS BE LIABLE TO ANY PARTY FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, 
 * OR CONSEQUENTIAL DAMAGES, INCLUDING LOST PROFITS, ARISING OUT OF THE USE 
 * OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF ANY OF THE UNIVERSITY OF 
 * CALIFORNIA, AUTHOR OR COPYRIGHT HOLDERS HAVE BEEN ADVISED OF 
 * THE POSSIBILITY OF SUCH DAMAGE.
 *
 * THE UNIVERSITY OF CALIFORNIA, AUTHOR AND COPYRIGHT HOLDER SPECIFICALLY 
 * DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. 
 * THE SOFTWARE PROVIDED HEREUNDER IS ON AN "AS IS" BASIS, AND THE 
 * UNIVERSITY OF CALIFORNIA, AUTHOR NOR COPYRIGHT HOLDERS HAVE ANY 
 * OBLIGATIONS TO PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, 
 * OR MODIFICATIONS.
 *
 * IDENTIFICATION
 *	  $PostgreSQL$
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include <fcntl.h>
#include <sys/stat.h>
#include <unistd.h>

#include "fmgr.h"
#include "funcapi.h"
#include "access/transam.h"
#include "access/xact.h"
#include "storage/fd.h"
#include "storage/procarray.h"
#include "miscadmin.h"

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

Datum		pg_snapclone_publish(PG_FUNCTION_ARGS);
Datum		pg_snapclone_unpublish_local(PG_FUNCTION_ARGS);
Datum		pg_snapclone_unpublish_specific(PG_FUNCTION_ARGS);
Datum		pg_snapclone_unpublish_internal(int remove_pid);
Datum		pg_snapclone_clone(PG_FUNCTION_ARGS);
Datum		pg_snapclone_clean_directory(PG_FUNCTION_ARGS);

#define	PG_SNAPCLONE_DIRPATH	"/pg_snapclone"

typedef struct SavedSnapshotHeaderData
{
	/*
	 * Fields to allow us to cross-check the snapshot
	 */
	int			  pid;
	Oid			  userid;

	/*
	 * Snapshot fields
	 */
	TransactionId xmin;         /* XID < xmin are visible to me */
	TransactionId xmax;         /* XID >= xmax are invisible to me */
	uint32        xcnt;         /* # of xact ids in xip[] */

	/* 
	 * we ignore subxcnt for pre-8.2, but it seems easier to have a common
	 * file format across versions
	 */
	int32         subxcnt;		/* # of xact ids in subxip[], -1 if overflow */

} SavedSnapshotHeaderData;

/* Double layering makes CRC calculation easier */
typedef struct SavedSnapshotHeader
{
	pg_crc32	  crc;
	SavedSnapshotHeaderData	hdr;
} SavedSnapshotHeader;

#define		FILE_READ	"read"
#define		FILE_WRITE	"write"

/*
 * pg_snapclone_publish
 *
 * Publishes the calling session's first transaction snapshot, referred
 * to internally as the Serializable snapshot. We return the backend id
 * of the publishing session. We can't return the TransactionId because
 * this won't have been assigned yet in 8.3
 *
 * This requires no acess to shared memory to access internal data
 */
PG_FUNCTION_INFO_V1(pg_snapclone_publish);

Datum
pg_snapclone_publish(PG_FUNCTION_ARGS)
{
	int		pid = MyProcPid;
	SavedSnapshotHeader snapHeader;

	/* Release specific declarations */
	struct	stat stat_buf;
	char	snapshotDirPath[MAXPGPATH];
	char	snapshotFilePath[MAXPGPATH];
	int		fd;
	Size	nbytes;

	if (!IsXactIsoLevelSerializable)
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TRANSACTION_STATE),
				 (errmsg("can only publish Serializable snapshots"))));

#ifdef SUPERUSER_ONLY
	if (!superuser())
		ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("must be superuser to publish snapshots"))));
#endif

	snapHeader.hdr.pid = pid;
	snapHeader.hdr.userid = GetUserId();
	snapHeader.hdr.xmin = SerializableSnapshot->xmin;
	snapHeader.hdr.xmax = SerializableSnapshot->xmax;
	snapHeader.hdr.xcnt = SerializableSnapshot->xcnt;
#ifdef PG82_OR_HIGHER
	snapHeader.hdr.subxcnt = SerializableSnapshot->subxcnt;
#else
	snapHeader.hdr.subxcnt = 0;
#endif
	/* calc CRC */
	snapHeader.crc = 0;

	/*
 	 * If pg_snapclone directory does not exist, create it
	 *
     * Don't check for error from mkdir; it could fail if someone else
     * just did the same thing.  If it doesn't work then we'll bomb out on
     * the second create attempt, instead.
	 */
	snprintf(snapshotDirPath, MAXPGPATH, "%s" PG_SNAPCLONE_DIRPATH, DataDir);
	if (stat(snapshotDirPath, &stat_buf) != 0)
		mkdir(snapshotDirPath, S_IRWXU);

	/*
 	 * Assign filename for new snapshot
	 */
	snprintf(snapshotFilePath, MAXPGPATH, "%s" PG_SNAPCLONE_DIRPATH "/snapshot.%d", DataDir, pid);

	/*
 	 * Remove stale file, if it exists. So no need to check return code
	 */
	unlink(snapshotFilePath);

	fd = BasicOpenFile(snapshotFilePath, 
				O_RDWR | O_CREAT | O_EXCL | PG_BINARY, 
				S_IRUSR | S_IWUSR);

	if (fd < 0)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not create snapshot file \"%s\": %m",
						snapshotFilePath)));

	if (write(fd, &snapHeader, sizeof(SavedSnapshotHeader)) != sizeof(SavedSnapshotHeader))
		ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_WRITE, "header to ", snapshotFilePath)));

	if (snapHeader.hdr.xcnt > 0)
	{
		nbytes = snapHeader.hdr.xcnt * sizeof(TransactionId);

		/*
		 * Write out the xip array
		 */
		if (write(fd, (char *)SerializableSnapshot->xip, nbytes) != nbytes)
			ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_WRITE, "xips to ", snapshotFilePath)));
	}

#ifdef PG82_OR_HIGHER
	if (snapHeader.hdr.subxcnt > 0)
	{
		nbytes = snapHeader.hdr.subxcnt * sizeof(TransactionId);

		/*
		 * Write out the subxip array
		 */
		if (write(fd, (char *)SerializableSnapshot->xip, nbytes) != nbytes)
			ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not %s %ssnapshot file \"%s\": %m",
						FILE_WRITE, "subxips to ", snapshotFilePath)));
	}
#endif

	/*
	 * Fsync and close
	 */
	if (pg_fsync(fd) != 0)
		ereport(ERROR,
			(errcode_for_file_access(),
			 errmsg("could not %s %ssnapshot file \"%s\": %m",
					FILE_WRITE, "subxips to ", snapshotFilePath)));

	if (close(fd))
		ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_WRITE, "close to ", snapshotFilePath)));

	ereport(LOG, (errmsg("published snapshot file \"%s\"", snapshotFilePath)));

	PG_RETURN_INT32(pid);
}

/*
 * pg_snapclone_unpublish (various)
 *
 */
PG_FUNCTION_INFO_V1(pg_snapclone_unpublish_local);
Datum
pg_snapclone_unpublish_local(PG_FUNCTION_ARGS)
{
	return pg_snapclone_unpublish_internal(0);
}

PG_FUNCTION_INFO_V1(pg_snapclone_unpublish_specific);
Datum
pg_snapclone_unpublish_specific(PG_FUNCTION_ARGS)
{
	int		pid = PG_GETARG_INT32(0);

	return pg_snapclone_unpublish_internal(pid);
}

Datum
pg_snapclone_unpublish_internal(int remove_pid)
{
	int		pid;
	bool	removed = false;

	char	snapshotDirPath[MAXPGPATH];
	char	snapshotFilePath[MAXPGPATH];
	int		rc;

	if (remove_pid == 0)
		pid = MyProcPid;
	else
	{
		if (remove_pid != MyProcPid && IsBackendPid(remove_pid))
			ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("cannot remove published snapshot for another running backend"))));

		pid = remove_pid;
	}

#ifdef SUPERUSER_ONLY
	if (!superuser())
		ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("must be superuser to unpublish snapshots"))));
#endif

	/*
 	 * Assign filename for new snapshot
	 */
	snprintf(snapshotDirPath, MAXPGPATH, "%s" PG_SNAPCLONE_DIRPATH, DataDir);
	snprintf(snapshotFilePath, MAXPGPATH, "%s/snapshot.%d", snapshotDirPath, pid);

	/*
 	 * Remove stale file, if it exists
	 */
	rc = unlink(snapshotFilePath);

	if (rc == 0)
	{
		removed = true;
		ereport(LOG, (errmsg("removed snapshot file \"%s\"", snapshotFilePath)));
	}

	PG_RETURN_BOOL(removed);
}


/*
 * pg_snapclone_clone
 *
 * Reuses a published snapshot using the session's transactionid as a handle.
 * All other relevant information is not visible to the user. Many checks
 * are required at runtime to ensure this can be valid.
 *
 * - the pid must still be connected
 * - the calling transaction must be a SERIALIZABLE transaction
 * - the cloning session must be same userid as publisher, or superuser
 * - the cloning session must have started its transaction earlier than the
 *   publishing transaction, so that xmin always goes forwards, if at all,
 *   and so that the cloning transaction is present already within the snapshot
 *	 if that release of Postgres would expect that (pre-8.3 only)
 */
PG_FUNCTION_INFO_V1(pg_snapclone_clone);

Datum
pg_snapclone_clone(PG_FUNCTION_ARGS)
{
	int		pid = PG_GETARG_INT32(0);
	SavedSnapshotHeader snapHeader;

	char	snapshotFilePath[MAXPGPATH];
	int		fd;
	int		i;
	TransactionId	xread;

	if (!IsXactIsoLevelSerializable)
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TRANSACTION_STATE),
				 (errmsg("can only set snapshot if transaction is Serializable"))));

#ifdef SUPERUSER_ONLY
	if (!superuser())
		ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("must be superuser to clone snapshots"))));
#endif

	/*
 	 * Assign filename for new snapshot
	 */
	snprintf(snapshotFilePath, MAXPGPATH, "%s" PG_SNAPCLONE_DIRPATH "/snapshot.%d", DataDir, pid);

	fd = BasicOpenFile(snapshotFilePath, O_RDONLY | PG_BINARY, 0);

	if (fd < 0)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not create snapshot file \"%s\": %m",
						snapshotFilePath)));

	if (read(fd, &snapHeader, sizeof(SavedSnapshotHeader)) != sizeof(SavedSnapshotHeader))
		ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_READ, "header of ", snapshotFilePath)));

	ereport(NOTICE, (errmsg("reading snapshot file \"%s\"\n"
					"userid=%u, xcnt=%u, subxcnt=%d, xmin=%u, xmax=%u", 
			snapshotFilePath,
			snapHeader.hdr.userid,
			snapHeader.hdr.xcnt,
#ifdef PG82_OR_HIGHER
			snapHeader.hdr.subxcnt,
#else
			0,
#endif
			snapHeader.hdr.xmin,
			snapHeader.hdr.xmax)));

	/*
	 * - the transactionid must still be active
	 * - the calling transaction must be a SERIALIZABLE transaction
	 * - the cloning session must be same userid as publisher, or superuser
	 * - the cloning session must have started its transaction earlier than the
	 *   publishing transaction, so that xmin always goes forwardsi, if at all,
	 *   and so that the cloning transaction is present already within the snapshot
	 */
	if (!IsBackendPid(snapHeader.hdr.pid))
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TRANSACTION_STATE),
				 (errmsg("can only clone snapshot from a running transaction"))));

	if (!TransactionIdFollowsOrEquals(snapHeader.hdr.xmin, RecentXmin))
		ereport(ERROR,
				(errcode(ERRCODE_INVALID_TRANSACTION_STATE),
				 (errmsg("can only clone snapshot if publishing transaction started later"))));

	if (snapHeader.hdr.userid != GetUserId())
		ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("can only clone snapshot if same user as publishing transaction"))));

	/*
	 * Deep breath, then overwrite SerializableSnapshot
	 */
	SerializableSnapshot->xmin = snapHeader.hdr.xmin;
	SerializableSnapshot->xmax = snapHeader.hdr.xmax;
	SerializableSnapshot->xcnt = snapHeader.hdr.xcnt;
	SerializableSnapshot->subxcnt = snapHeader.hdr.subxcnt;

	for (i = 0; i < snapHeader.hdr.xcnt; i++)
	{
		if (read(fd, &xread, sizeof(TransactionId)) != sizeof(TransactionId))
			ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_READ, "transaction from ", snapshotFilePath)));
		SerializableSnapshot->xip[i++] = xread;
	}

#ifdef PG82_OR_HIGHER
	for (i = 0; i < snapHeader.hdr.subxcnt; i++)
	{
		if (read(fd, &xread, sizeof(TransactionId)) != sizeof(TransactionId))
			ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_READ, "subtransaction from ", snapshotFilePath)));
		SerializableSnapshot->subxip[i++] = xread;
	}
#endif

	if (close(fd))
		ereport(ERROR,
				(errcode_for_file_access(),
					 errmsg("could not %s %ssnapshot file \"%s\": %m",
							FILE_READ, "close to ", snapshotFilePath)));

	ereport(LOG, (errmsg("cloned snapshot using \"%s\"", snapshotFilePath)));

	PG_RETURN_VOID();
}

/*
 * pg_snapclone_clean_directory
 *
 */
PG_FUNCTION_INFO_V1(pg_snapclone_clean_directory);

Datum
pg_snapclone_clean_directory(PG_FUNCTION_ARGS)
{
	int		pid;
	bool	removed = false;
	char	snapshotDirPath[MAXPGPATH];
	DIR		*snapshotDir;
	struct dirent *snapshotDirEntry;
	char	snapshotFilePath[MAXPGPATH];
	int		rc;

	if (!superuser())
		ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("must be superuser to clean pg_snapclone directory"))));

	snprintf(snapshotDirPath, MAXPGPATH, "%s" PG_SNAPCLONE_DIRPATH, DataDir);
	snapshotDir = AllocateDir(snapshotDirPath);
	if (snapshotDir == NULL)
		ereport(ERROR,
				(errcode_for_file_access(),
				 errmsg("could not open pg_snapclone \"%s\": %m",
						snapshotDirPath)));

	while ((snapshotDirEntry = ReadDir(snapshotDir, snapshotDirPath)) != NULL)
	{
		/*
		 * ignore . and .. 
		 */
		if (strlen(snapshotDirEntry->d_name) <= 9)
			continue;

		pid = atoi(snapshotDirEntry->d_name + 9);
		if (!IsBackendPid(pid))
		{
			snprintf(snapshotFilePath, MAXPGPATH, "%s/%s", 
								snapshotDirPath, snapshotDirEntry->d_name);
			rc = unlink(snapshotFilePath);
			if (rc == 0)
			{
				removed = true;
				ereport(LOG, (errmsg("removed snapshot file \"%s\"", 
								snapshotFilePath)));
			}
		}
	}

	FreeDir(snapshotDir);

	PG_RETURN_BOOL(removed);
}
