/***************************************************************
 *
 * Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
 * University of Wisconsin-Madison, WI.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you
 * may not use this file except in compliance with the License.  You may
 * obtain a copy of the License at
 * 
 *    http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 ***************************************************************/


#include "condor_common.h"
#include "shadow.h"
#include "condor_daemon_core.h"
#include "condor_qmgr.h"         // need to talk to schedd's qmgr
#include "condor_attributes.h"   // for ATTR_ ClassAd stuff
#include "condor_email.h"        // for email.
#include "metric_units.h"
#include "ShadowHookMgr.h"
#include "store_cred.h"

extern "C" char* d_format_time(double);

UniShadow::UniShadow() : delayedExitReason( -1 ) {
		// pass RemoteResource ourself, so it knows where to go if
		// it has to call something like shutDown().
	remRes = new RemoteResource( this );
	remRes->setWaitOnKillFailure(true);
}

UniShadow::~UniShadow() {
	if ( remRes ) delete remRes;
	daemonCore->Cancel_Command( SHADOW_UPDATEINFO );
	daemonCore->Cancel_Command( CREDD_GET_CRED );
}


int
UniShadow::updateFromStarter(int /* command */, Stream *s)
{
	ClassAd update_ad;

	// get info from the starter encapsulated in a ClassAd
	s->decode();
	if( ! getClassAd(s, update_ad) ) {
		dprintf( D_ALWAYS, "ERROR in UniShadow::updateFromStarter:"
				 "Can't read ClassAd, aborting.\n" );
		return FALSE;
	}
	s->end_of_message();

	fix_update_ad(update_ad);
	return updateFromStarterClassAd(&update_ad);
}


int
UniShadow::updateFromStarterClassAd(ClassAd* update_ad) {
	ClassAd *job_ad = getJobAd();
	if (!job_ad) {
		// should never really happen...
		return FALSE;
	}


		// Save the current image size, so that if it changes as a
		// result of this update, we can log an event to the userlog.
	int64_t prev_usage = 0, prev_rss = 0, prev_pss = -1;
	int64_t prev_image = getImageSize(prev_usage, prev_rss, prev_pss);

		// Stick everything we care about into our RemoteResource.
		// For the UniShadow (with only 1 RemoteResource) it has a
		// pointer to the identical copy of the job classad, so when
		// it updates its copy, it's really updating ours, too.
	remRes->updateFromStarter(update_ad);

	int64_t cur_usage = 0, cur_rss = 0, cur_pss = -1;
	int64_t cur_image = getImageSize(cur_usage, cur_rss, cur_pss);
	if (cur_image != prev_image || 
		cur_usage != prev_usage ||
		cur_rss != prev_rss ||
		cur_pss != prev_pss) {
		JobImageSizeEvent event;
		event.image_size_kb = cur_image;
		event.memory_usage_mb = cur_usage;
		event.resident_set_size_kb = cur_rss;
		event.proportional_set_size_kb = cur_pss;

			// for performance, do not bother fsyncing this event
		if (!uLog.writeEventNoFsync(&event, job_ad)) {
			dprintf(D_ALWAYS, "Unable to log ULOG_IMAGE_SIZE event\n");
		}
	}

	return TRUE;
}


void
UniShadow::init( ClassAd* job_ad, const char* schedd_addr, const char *xfer_queue_contact_info )
{
	if ( !job_ad ) {
		EXCEPT("No job_ad defined!");
	}

		// base init takes care of lots of stuff:
	baseInit( job_ad, schedd_addr, xfer_queue_contact_info );

		// we're only dealing with one host, so the rest is pretty
		// trivial.  we can just lookup everything we need in the job
		// ad, since it'll have the ClaimId, address (in the ClaimId)
		// startd's name (RemoteHost) and pool (RemotePool).
	remRes->setStartdInfo( jobAd );
	
		// In this case we just pass the pointer along...
	remRes->setJobAd( jobAd );
	
		// Register command which gets updates from the starter
		// on the job's image size, cpu usage, etc.  Each kind of
		// shadow implements it's own version of this to deal w/ it
		// properly depending on parallel vs. serial jobs, etc. 
	daemonCore->
		Register_Command( SHADOW_UPDATEINFO, "SHADOW_UPDATEINFO",
						  (CommandHandlercpp)&UniShadow::updateFromStarter, 
						  "UniShadow::updateFromStarter", this, DAEMON );

		// Register command which the starter uses to fetch a user's Kerberose/Afs auth credential
	daemonCore->
		Register_Command( CREDD_GET_CRED, "CREDD_GET_CRED",
						  &cred_get_cred_handler,
						  "cred_get_cred_handler", DAEMON,
						  true /*force authentication*/ );

		// Register our job hooks
	m_hook_mgr = std::unique_ptr<ShadowHookMgr>(new ShadowHookMgr());
	if (!m_hook_mgr->initialize(job_ad)) {
		m_hook_mgr.reset();
	}
}


void
UniShadow::spawnFinish()
{
	hookTimerCancel();
	if( ! remRes->activateClaim() ) {
			// we're screwed, give up:
		shutDown( JOB_NOT_STARTED );
	}
}


void
UniShadow::spawn()
{
	if (!m_hook_mgr) {
		dprintf(D_ALWAYS, "No hook manager available; will activate claim immediately.\n");
		spawnFinish();
	} else {
		auto rval = m_hook_mgr->tryHookPrepareJob();
		if (rval == -1) {
			dprintf(D_ALWAYS, "Prepare job hook has failed.  Will shutdown job.\n");
			BaseShadow::log_except("Submit-side job hook execution failed");
			shutDown(JOB_NOT_STARTED);
		} else if (rval == 0) {
			dprintf(D_FULLDEBUG, "No prepare job hook to run - activating job immediately.\n");
			spawnFinish();
		} else if (rval == 1) {
			dprintf(D_FULLDEBUG, "Hook successfully spawned.\n");
			m_exit_hook_timer_tid = daemonCore->Register_Timer(m_hook_mgr->getHookTimeout(HOOK_SHADOW_PREPARE_JOB, 120),
				(TimerHandlercpp)&UniShadow::hookTimeout,
				"hookTimeout",
				this);
		} else {
			EXCEPT("Hook manager returned an invalid code\n");
		}
	}
}


void
UniShadow::hookTimeout( int /* timerID */ )
{
	dprintf(D_ERROR, "Timed out waiting for a hook to exit\n");
	BaseShadow::log_except("Submit-side job hook execution timed out");
	shutDown(JOB_NOT_STARTED);
}


void
UniShadow::hookTimerCancel()
{
	if (m_exit_hook_timer_tid != -1) {
		daemonCore->Cancel_Timer(m_exit_hook_timer_tid);
		m_exit_hook_timer_tid = -1;
	}
}


void
UniShadow::reconnect( void )
{
	remRes->reconnect();
}


bool 
UniShadow::supportsReconnect( void )
{
		// For the UniShadow, the answer to this depends on our remote
		// starter.  If that supports it, so do we.  If not, we don't. 
	return remRes->supportsReconnect();
}


void
UniShadow::logExecuteEvent( void )
{
	ExecuteEvent event;
	char* sinful = NULL;
	remRes->getStartdAddress( sinful );
	event.setExecuteHost( sinful );
	free( sinful );

	remRes->populateExecuteEvent(event.slotName, event.setProp());

	if( !uLog.writeEvent(&event, getJobAd()) ) {
		dprintf( D_ALWAYS, "Unable to log ULOG_EXECUTE event: "
				 "can't write to UserLog!\n" );
	}
}


void
UniShadow::cleanUp( bool graceful )
{
		// Deactivate (fast) the claim
	if ( remRes ) {
		remRes->killStarter(graceful);
	}
}

void
UniShadow::gracefulShutDown( void )
{
	remRes->setExitReason(JOB_SHOULD_REQUEUE);
	remRes->killStarter(true);
}


int
UniShadow::getExitReason( void )
{
	if ( isDataflowJob ) {
		return JOB_EXITED;
	}
	if( remRes ) {
		return remRes->getExitReason();
	}
	return -1;
}

bool
UniShadow::claimIsClosing( void )
{
	if( remRes ) {
		return remRes->claimIsClosing();
	}
	return false;
}


void
UniShadow::emailTerminateEvent( int exitReason, update_style_t kind )
{
	Email mailer;
	double recvd_bytes = 0, sent_bytes = 0;

	if (kind == US_TERMINATE_PENDING) {
		/* I don't have a remote resource, so get the values directly from
			the jobad, and I know they should be present at this stage. */
		jobAd->LookupFloat(ATTR_BYTES_SENT, sent_bytes);
		jobAd->LookupFloat(ATTR_BYTES_RECVD, recvd_bytes);

		// note, we want to reverse the order of the send/recv in the
		// call, since we want the email from the job's perspective,
		// not the shadow's.. 

		mailer.sendExitWithBytes( jobAd, exitReason, 0, 0, 
						  sent_bytes, recvd_bytes);
		// all done.
		return;
	}

	// the default case of kind == US_NORMAL
	
		// note, we want to reverse the order of the send/recv in the
		// call, since we want the email from the job's perspective,
		// not the shadow's.. 
	mailer.sendExitWithBytes( jobAd, exitReason, 
							  bytesReceived(), bytesSent(), 
							  prev_run_bytes_sent + bytesReceived(), 
							  prev_run_bytes_recvd + bytesSent() );
}

void UniShadow::holdJob( const char* reason, int hold_reason_code, int hold_reason_subcode )
{
	remRes->setExitReason( JOB_SHOULD_HOLD );
	BaseShadow::holdJob(reason, hold_reason_code, hold_reason_subcode);
}

void UniShadow::removeJob( const char* reason )
{
	int iPrevExitReason=remRes->getExitReason();
	
	remRes->setExitReason( JOB_SHOULD_REMOVE );
	this->removeJobPre(reason);
	
	// exit immediately if the remote side is failing out or has already exited.
	if ( iPrevExitReason != JOB_SHOULD_REMOVE && iPrevExitReason != -1)
	{
		// don't wait for final update b/c there isn't one.
		DC_Exit( JOB_SHOULD_REMOVE );
	}
}

void
UniShadow::requestJobRemoval() {
	remRes->setExitReason( JOB_KILLED );
	bool job_wants_graceful_removal = jobWantsGracefulRemoval();
	dprintf(D_ALWAYS,"Requesting %s removal of job.\n",
			job_wants_graceful_removal ? "graceful" : "fast");
	remRes->killStarter( job_wants_graceful_removal );
}

int UniShadow::handleJobRemoval(int sig) {
    dprintf ( D_FULLDEBUG, "In handleJobRemoval(), sig %d\n", sig );
	remove_requested = true;
		// if we're not in the middle of trying to reconnect, we
		// should immediately kill the starter.  if we're
		// reconnecting, we'll do the right thing once a connection is
		// established now that the remove_requested flag is set... 
	if( remRes->getResourceState() != RR_RECONNECT ) {
		requestJobRemoval();
	}
		// more?
	return 0;
}


int UniShadow::JobSuspend( int sig )
{
	int iRet=1;
	
	if ( remRes->suspend() )
		iRet = 0;
	else
		dprintf ( D_ALWAYS, "JobSuspend() sig %d FAILED\n", sig );
	
	return iRet;
}

int UniShadow::JobResume( int sig )
{
	int iRet =1;
	
	if ( remRes->resume() )
		iRet =0;
	else
		dprintf ( D_ALWAYS, "JobResume() sig %d FAILED\n", sig );
	
	return iRet;
}

float
UniShadow::bytesSent()
{
	return remRes->bytesSent();
}


float
UniShadow::bytesReceived()
{
	return remRes->bytesReceived();
}

void
UniShadow::getFileTransferStats(ClassAd &upload_stats, ClassAd &download_stats)
{
	upload_stats = remRes->m_upload_file_stats;
	download_stats = remRes->m_download_file_stats;
}

void
UniShadow::getFileTransferStatus(FileTransferStatus &upload_status,FileTransferStatus &download_status)
{
	remRes->getFileTransferStatus(upload_status,download_status);
}

struct rusage
UniShadow::getRUsage( void ) 
{
	return remRes->getRUsage();
}


int64_t
UniShadow::getImageSize( int64_t & mem_usage, int64_t & rss, int64_t & pss )
{
	return remRes->getImageSize(mem_usage, rss, pss);
}


int64_t
UniShadow::getDiskUsage( void )
{
	return remRes->getDiskUsage();
}


bool
UniShadow::exitedBySignal( void )
{
	return remRes->exitedBySignal();
}


int
UniShadow::exitSignal( void )
{
	return remRes->exitSignal();
}


int
UniShadow::exitCode( void )
{
	if (isDataflowJob) {
		return 0;
	}
	return remRes->exitCode();
}


void
UniShadow::resourceBeganExecution( RemoteResource* rr )
{
	ASSERT( rr == remRes );

		// We've only got one remote resource, so if it's started
		// executing, we can safely log our execute event
	logExecuteEvent();

		// Invoke the base copy of this function to handle shared code.
	BaseShadow::resourceBeganExecution(rr);
}


void
UniShadow::resourceDisconnected( RemoteResource* rr )
{
	ASSERT( rr == remRes );


	// All of our children should be gone already, but let's be sure.
	daemonCore->kill_immediate_children();


	const char* txt = "Socket between submit and execute hosts "
		"closed unexpectedly";
	logDisconnectedEvent( txt );

	time_t now = time(NULL);
	jobAd->Assign(ATTR_JOB_DISCONNECTED_DATE, now);

	if (m_lazy_queue_update) {
			// For lazy update, we just want to make sure the
			// job_updater object knows about this attribute (which we
			// already updated our copy of).
		job_updater->watchAttribute(ATTR_JOB_DISCONNECTED_DATE);
	}
	else {
			// They want it now, so do the qmgmt operation directly.
		updateJobAttr(ATTR_JOB_DISCONNECTED_DATE, now);
	}
}

void
UniShadow::resourceReconnected( RemoteResource* rr )
{
	ASSERT( rr == remRes );

		// We've only got one remote resource, so if it successfully
		// reconnected, we can safely log our reconnect event
	logReconnectedEvent();

	// If the shadow started in reconnect mode, check the job ad to see
	// if we previously heard about the job starting execution, and set
	// up our state accordingly.
	if ( attemptingReconnectAtStartup ) {
		time_t job_execute_date = 0;
		time_t claim_start_date = 0;
		jobAd->LookupInteger(ATTR_JOB_CURRENT_START_EXECUTING_DATE, job_execute_date);
		jobAd->LookupInteger(ATTR_JOB_CURRENT_START_DATE, claim_start_date);
		if ( job_execute_date >= claim_start_date ) {
			began_execution = true;
		}
	}

		// Since our reconnect worked, clear attemptingReconnectAtStartup
		// flag so if we disconnect again and fail, we will exit
		// with JOB_SHOULD_REQUEUE instead of JOB_RECONNECT_FAILED.
		// See gt #4783.
	attemptingReconnectAtStartup = false;

		// Update NumJobReconnects in the schedd
		// TODO Should we do the update through the job_updater?
	int job_reconnect_cnt = 0;
	jobAd->LookupInteger(ATTR_NUM_JOB_RECONNECTS, job_reconnect_cnt);
	job_reconnect_cnt++;
	jobAd->Assign(ATTR_NUM_JOB_RECONNECTS, job_reconnect_cnt);
	jobAd->AssignExpr(ATTR_JOB_DISCONNECTED_DATE, "Undefined");

	if (m_lazy_queue_update) {
			// For lazy update, we just want to make sure the
			// job_updater object knows about this attribute (which we
			// already updated our copy of).
		job_updater->watchAttribute(ATTR_NUM_JOB_RECONNECTS);
		job_updater->watchAttribute(ATTR_JOB_DISCONNECTED_DATE);
	}
	else {
			// They want it now, so do the qmgmt operation directly.
		updateJobAttr(ATTR_NUM_JOB_RECONNECTS, job_reconnect_cnt);
		updateJobAttr(ATTR_JOB_DISCONNECTED_DATE, "Undefined");
	}

		// if we're trying to remove this job, now that connection is
		// reestablished, we can kill the starter, evict the job, and
		// handle any output/update messages from the starter.
	if( remove_requested ) {
		requestJobRemoval();
	}

		// If we know the job is already executing, ensure the timers
		// that are supposed to start then are running.
	if (began_execution) {
			// Start the timer for the periodic user job policy
		shadow_user_policy.startTimer();

			// Start the timer for updating the job queue for this job
		startQueueUpdateTimer();
	}
}


void
UniShadow::logDisconnectedEvent( const char* reason )
{
	JobDisconnectedEvent event;
	event.disconnect_reason = reason;

	DCStartd* dc_startd = remRes->getDCStartd();
	if( ! dc_startd ) {
		EXCEPT( "impossible: remRes::getDCStartd() returned NULL" );
	}
	event.startd_addr = dc_startd->addr();
	event.startd_name = dc_startd->name();

	if( !uLog.writeEventNoFsync(&event,getJobAd()) ) {
		dprintf( D_ALWAYS, "Unable to log ULOG_JOB_DISCONNECTED event\n" );
	}
}


void
UniShadow::logReconnectedEvent( void )
{
	JobReconnectedEvent event;

	DCStartd* dc_startd = remRes->getDCStartd();
	if( ! dc_startd ) {
		EXCEPT( "impossible: remRes::getDCStartd() returned NULL" );
	}
	event.startd_addr = dc_startd->addr();
	event.startd_name = dc_startd->name();

	char* starter = NULL;
	remRes->getStarterAddress( starter );
	event.starter_addr = starter;
	free( starter );
	starter = NULL;

	if( !uLog.writeEventNoFsync(&event,getJobAd()) ) {
		dprintf( D_ALWAYS, "Unable to log ULOG_JOB_RECONNECTED event\n" );
	}
}


void
UniShadow::logReconnectFailedEvent( const char* reason )
{
	JobReconnectFailedEvent event;

	event.reason = reason;

	DCStartd* dc_startd = remRes->getDCStartd();
	if( ! dc_startd ) {
		EXCEPT( "impossible: remRes::getDCStartd() returned NULL" );
	}
	event.startd_name = dc_startd->name();

	if( !uLog.writeEventNoFsync(&event,getJobAd()) ) {
		dprintf( D_ALWAYS, "Unable to log ULOG_JOB_RECONNECT_FAILED event\n" );
	}
}

bool
UniShadow::getMachineName( std::string &machineName )
{
	if( remRes ) {
		char *name = NULL;
		remRes->getMachineName(name);
		if( name ) {
			machineName = name;
			free(name);
			return true;
		}
	}
	return false;
}

void
UniShadow::exitAfterEvictingJob( int reason ) {
	// Don't exit unless the starter has signalled it's done, or until
	// twenty seconds after we'd otherwise have exited.  This should
	// reduce the number of scary-looking error messages from the network
	// layer in both the shadow and the starter logs.
	//
	// This function should be called in BaseShadow functions which call
	// cleanUp() and then DC_Exit() before returning to the event loop.  It's
	// not called from UniShadow::cleanUp() because a bunch of those functions
	// do important-looking things between calling cleanUp() and calling
	// DC_Exit().
	if( remRes->gotJobExit() || remRes->getClaimSock() == NULL ) {
		DC_Exit( reason );
	} else {
		this->delayedExitReason = reason;
		remRes->setExitReason( reason );
		daemonCore->Register_Timer( 20, 0,
				(TimerHandlercpp)&UniShadow::exitLeaseHandler,
				"exit lease handler", this );
	}
}

bool
UniShadow::exitDelayed( int &reason ) {
	if ( delayedExitReason != -1 ) {
		reason = delayedExitReason;
		return true;
	}
	return false;
}

void
UniShadow::exitLeaseHandler( int /* timerID */ ) const {
	DC_Exit( delayedExitReason );
}

void
UniShadow::recordFileTransferStateChanges( ClassAd * jobAd, ClassAd * ftAd ) {
	bool tq = false; bool tqSet = ftAd->LookupBool( ATTR_TRANSFER_QUEUED, tq );
	bool ti = false; bool tiSet = ftAd->LookupBool( ATTR_TRANSFERRING_INPUT, ti );
	bool to = false; bool toSet = ftAd->LookupBool( ATTR_TRANSFERRING_OUTPUT, to );

	// If either ATTR_TRANSFER_QUEUED or ATTR_TRANSFERRING_INPUT hasn't
	// been set yet, file transfer hasn't done anything yet and there's
	// event to record.
	if( (!tqSet) || (!tiSet) ) {
		return;
	}

	// If all six of the booleans in the ftAd are the same as the booleans
	// in the job ad, then we've already written out the event.
	bool jtq = false; bool jtqSet = jobAd->LookupBool( ATTR_TRANSFER_QUEUED, jtq );
	bool jti = false; bool jtiSet = jobAd->LookupBool( ATTR_TRANSFERRING_INPUT, jti );
	bool jto = false; bool jtoSet = jobAd->LookupBool( ATTR_TRANSFERRING_OUTPUT, jto );
	if(  jtq == tq && jtqSet == tqSet
	  && jti == ti && jtiSet == tiSet
	  && jto == to && jtoSet == toSet ) {
		return;
	}

	//
	// I suppose for maximum geek cred I should shift tq, ti, toSet, and to
	// into a bitfield and switch on that, instead.
	//

	FileTransferEvent te;
	if( tq && ti && (!toSet) ) {
		te.setType( FileTransferEvent::IN_QUEUED );

		// There really ought to be a remote resource if we're doing
		// file transfer...
		if( remRes ) {
			char * starterAddr = NULL;
			remRes->getStarterAddress( starterAddr );
			if( starterAddr ) {
				te.setHost( starterAddr );
				free( starterAddr );
			}
		}

		jobAd->Assign( "TransferInQueued", (int)time(NULL) );
	} else if( (!tq) && ti && (!toSet) ) {
		te.setType( FileTransferEvent::IN_STARTED );

		time_t now = time(nullptr);
		jobAd->Assign( "TransferInStarted", now );

		time_t then;
		if( jobAd->LookupInteger( "TransferInQueued", then ) ) {
			te.setQueueingDelay( now - then );
		} else {
			if( remRes ) { // this should always be true...
				char * starterAddr = NULL;
				remRes->getStarterAddress( starterAddr );
				if( starterAddr ) {
					te.setHost( starterAddr );
					free( starterAddr );
				}
			}
		}
	} else if( (!tq) && (!ti) && (!toSet) ) {
		te.setType( FileTransferEvent::IN_FINISHED );
		// te.setSuccess( ... );

		jobAd->Assign( "TransferInFinished", (int)time(NULL) );
	} else if( tq && (!ti) && (toSet && to) ) {
		te.setType( FileTransferEvent::OUT_QUEUED );

		jobAd->Assign( "TransferOutQueued", (int)time(NULL) );
	} else if( (!tq) && (!ti) && (toSet && to) ) {
		te.setType( FileTransferEvent::OUT_STARTED );

		time_t now = (int)time(NULL);
		jobAd->Assign( "TransferOutStarted", now );

		time_t then;
		if( jobAd->LookupInteger( "TransferInQueued", then ) ) {
			te.setQueueingDelay( now - then );
		}
	} else if( (!tq) && (!ti) && (toSet && (!to)) ) {
		te.setType( FileTransferEvent::OUT_FINISHED );
		// te.setSuccess( ... );

		jobAd->Assign( "TransferOutFinished", (int)time(NULL) );
	}

	if(! uLog.writeEvent( &te, jobAd )) {
		dprintf( D_ALWAYS, "Unable to log file transfer event.\n" );
	}
}
