20110113-copy_context.patch

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 00139b9..5ee6bb0 100644
*** a/src/backend/commands/copy.c
--- b/src/backend/commands/copy.c
*************** typedef struct CopyStateData
*** 124,129 ****
--- 124,134 ----
  	const char *cur_attval;		/* current att value for error messages */
  
  	/*
+ 	 * Working state for COPY TO/FROM
+ 	 */
+ 	MemoryContext copycontext;	/* per-copy execution context */
+ 
+ 	/*
  	 * Working state for COPY TO
  	 */
  	FmgrInfo   *out_functions;	/* lookup info for output functions */
*************** static const char BinarySignature[11] = 
*** 260,265 ****
--- 265,271 ----
  /* non-export function prototypes */
  static CopyState BeginCopy(bool is_from, Relation rel, Node *raw_query,
  				const char *queryString, List *attnamelist, List *options);
+ static void EndCopy(CopyState cstate);
  static CopyState BeginCopyTo(Relation rel, Node *query, const char *queryString,
  				const char *filename, List *attnamelist, List *options);
  static void EndCopyTo(CopyState cstate);
*************** BeginCopy(bool is_from,
*** 839,848 ****
--- 845,870 ----
  	ListCell   *option;
  	TupleDesc	tupDesc;
  	int			num_phys_attrs;
+ 	MemoryContext oldcontext;
  
  	/* Allocate workspace and zero all fields */
  	cstate = (CopyStateData *) palloc0(sizeof(CopyStateData));
  
+ 	/*
+ 	 * We allocate everything linked from cstate in new per-COPY context
+ 	 * which is created in caller's context.
+ 	 * This would avoid memory leak repeated COPY use in a query.
+ 	 *
+ 	 * Caller should MemoryContextDelete(cstate->copycontext) first, and then
+ 	 * pfree(cstate) to release memory.
+ 	 */
+ 	cstate->copycontext = AllocSetContextCreate(CurrentMemoryContext,
+ 												"COPY context",
+ 												ALLOCSET_DEFAULT_MINSIZE,
+ 												ALLOCSET_DEFAULT_INITSIZE,
+ 												ALLOCSET_DEFAULT_MAXSIZE);
+ 	oldcontext = MemoryContextSwitchTo(cstate->copycontext);
+ 
  	/* Extract options from the statement node tree */
  	foreach(option, options)
  	{
*************** BeginCopy(bool is_from,
*** 1242,1250 ****
--- 1264,1284 ----
  
  	cstate->copy_dest = COPY_FILE;		/* default */
  
+ 	/* switch back to caller's context */
+ 	MemoryContextSwitchTo(oldcontext);
+ 
  	return cstate;
  }
  
+ /*
+  * Release resources allocated in BeginCopu() for this copy execution.
+  */
+ static void
+ EndCopy(CopyState cstate)
+ {
+ 	MemoryContextDelete(cstate->copycontext);
+ 	pfree(cstate);
+ }
  
  /*
   * Setup CopyState to read tuples from a table or a query for COPY TO.
*************** BeginCopyTo(Relation rel,
*** 1259,1266 ****
--- 1293,1302 ----
  {
  	CopyState	cstate;
  	bool		pipe = (filename == NULL);
+ 	MemoryContext oldcontext;
  
  	cstate = BeginCopy(false, rel, query, queryString, attnamelist, options);
+ 	oldcontext = MemoryContextSwitchTo(cstate->copycontext);
  
  	if (cstate->rel)
  	{
*************** BeginCopyTo(Relation rel,
*** 1328,1333 ****
--- 1364,1372 ----
  					 errmsg("\"%s\" is a directory", cstate->filename)));
  	}
  
+ 	/* switch back to caller's context */
+ 	MemoryContextSwitchTo(oldcontext);
+ 
  	return cstate;
  }
  
*************** EndCopyTo(CopyState cstate)
*** 1395,1401 ****
  		PopActiveSnapshot();
  	}
  
! 	pfree(cstate);
  }
  
  /*
--- 1434,1441 ----
  		PopActiveSnapshot();
  	}
  
! 	/* Release memory resources */
! 	EndCopy(cstate);
  }
  
  /*
*************** BeginCopyFrom(Relation rel,
*** 2010,2017 ****
--- 2050,2059 ----
  	EState	   *estate = CreateExecutorState(); /* for ExecPrepareExpr() */
  	int		   *defmap;
  	ExprState **defexprs;
+ 	MemoryContext oldcontext;
  
  	cstate = BeginCopy(true, rel, NULL, NULL, attnamelist, options);
+ 	oldcontext = MemoryContextSwitchTo(cstate->copycontext);
  
  	/* Initialize state variables */
  	cstate->fe_eof = false;
*************** BeginCopyFrom(Relation rel,
*** 2172,2177 ****
--- 2214,2221 ----
  		cstate->raw_fields = (char **) palloc(nfields * sizeof(char *));
  	}
  
+ 	MemoryContextSwitchTo(oldcontext);
+ 
  	return cstate;
  }
  
*************** EndCopyFrom(CopyState cstate)
*** 2433,2448 ****
  	}
  
  	/* Clean up storage */
! 	if (!cstate->binary)
! 		pfree(cstate->raw_fields);
! 	pfree(cstate->attribute_buf.data);
! 	pfree(cstate->line_buf.data);
! 	pfree(cstate->raw_buf);
! 	pfree(cstate->in_functions);
! 	pfree(cstate->typioparams);
! 	pfree(cstate->defmap);
! 	pfree(cstate->defexprs);
! 	pfree(cstate);
  }
  
  
--- 2477,2483 ----
  	}
  
  	/* Clean up storage */
! 	EndCopy(cstate);
  }