Team Fortress 2 Source Code as on 22/4/2020
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

2478 lines
63 KiB

  1. //========= Copyright Valve Corporation, All rights reserved. ============//
  2. //
  3. // Purpose: This module implements the subset of MPI that VRAD and VVIS use.
  4. //
  5. // $NoKeywords: $
  6. //=============================================================================//
  7. #include <windows.h>
  8. #include <io.h>
  9. #include <conio.h>
  10. #include <sys/stat.h>
  11. #include <stdio.h>
  12. #include <direct.h>
  13. #include "iphelpers.h"
  14. #include "utlvector.h"
  15. #include "utllinkedlist.h"
  16. #include "vmpi.h"
  17. #include "bitbuf.h"
  18. #include "tier1/strtools.h"
  19. #include "threadhelpers.h"
  20. #include "IThreadedTCPSocket.h"
  21. #include "vstdlib/random.h"
  22. #include "vmpi_distribute_work.h"
  23. #include "filesystem.h"
  24. #include "checksum_md5.h"
  25. #include "tslist.h"
  26. #include "tier0/icommandline.h"
  27. #define DEFAULT_MAX_WORKERS 32 // Unless they specify -mpi_MaxWorkers, it will stop accepting workers after it gets this many.
  28. int g_nMaxWorkerCount = DEFAULT_MAX_WORKERS;
  29. #define VMPI_INTERNAL_PACKET_ID 27
  30. #define VMPI_INTERNAL_SUBPACKET_MACHINE_NAME 1
  31. #define VMPI_INTERNAL_SUBPACKET_COMMAND_LINE 2
  32. #define VMPI_INTERNAL_SUBPACKET_WAITING_FOR_COMMAND_LINE 3
  33. #define VMPI_INTERNAL_SUBPACKET_GROUPED_PACKET 4
  34. #define VMPI_INTERNAL_SUBPACKET_TIMING_WAIT_DONE 5
  35. #define VMPI_INTERNAL_SUBPACKET_VERIFY_EXE_NAME 6
  36. typedef CUtlVector<char> PersistentPacket;
  37. CCriticalSection g_PersistentPacketsCS;
  38. CUtlLinkedList<PersistentPacket*> g_PersistentPackets;
  39. // Command-line parameters list.
  40. #define VMPI_PARAM( paramName, paramFlags, helpText ) {paramName, paramFlags, "-"#paramName, helpText},
  41. class CVMPIParam
  42. {
  43. public:
  44. EVMPICmdLineParam m_eParam;
  45. int m_ParamFlags;
  46. const char *m_pName;
  47. const char *m_pHelpText;
  48. };
  49. static CVMPIParam g_VMPIParams[] =
  50. {
  51. {k_eVMPICmdLineParam_FirstParam, 0, "k_eVMPICmdLineParam_FirstParam", "unused"},
  52. {k_eVMPICmdLineParam_VMPIParam, 0, "mpi", "Enable VMPI."},
  53. #include "vmpi_parameters.h"
  54. };
  55. #undef VMPI_PARAM
  56. // ---------------------------------------------------------------------------------------- //
  57. // Globals.
  58. // ---------------------------------------------------------------------------------------- //
  59. class CVMPIConnection;
  60. // Used by -mpi_AutoRestart.
  61. CUtlVector<char*> g_OriginalCommandLineParameters;
  62. // This queues up all the incoming VMPI messages.
  63. CCriticalSection g_VMPIMessagesCS;
  64. CUtlLinkedList< CTCPPacket*, int > g_VMPIMessages;
  65. CEvent g_VMPIMessagesEvent; // This is set when there are messages in the queue.
  66. // These are used to notify the main thread when some socket had OnError() called on it.
  67. CUtlLinkedList<CVMPIConnection*,int> g_ErrorSockets;
  68. CEvent g_ErrorSocketsEvent;
  69. CCriticalSection g_ErrorSocketsCS;
  70. bool g_bTimingWaitDone = false;
  71. bool g_bGroupPackets = false;
  72. #define MAX_VMPI_CONNECTIONS 512
  73. CVMPIConnection *g_Connections[MAX_VMPI_CONNECTIONS];
  74. int g_nConnections = 0;
  75. CCriticalSection g_ConnectionsCS;
  76. // If true, then it will set certain thread priorities low.
  77. bool g_bSetThreadPriorities = true;
  78. VMPIDispatchFn g_VMPIDispatch[MAX_VMPI_PACKET_IDS];
  79. CTSList<MessageBuffer *> g_DispatchBuffers;
  80. VMPIRunMode g_VMPIRunMode = VMPI_RUN_NETWORKED;
  81. VMPIFileSystemMode g_VMPIFileSystemMode = VMPI_FILESYSTEM_TCP;
  82. static char g_GroupedPacketHeader[] = { VMPI_INTERNAL_PACKET_ID, VMPI_INTERNAL_SUBPACKET_GROUPED_PACKET };
  83. static unsigned long g_LastFlushGroupedPacketsTime = 0;
  84. // Set to true if we're running under the SDK (i.e. vmpi_transfer.exe is not found).
  85. bool g_bVMPISDKMode = false;
  86. bool g_bVMPISDKModeSet = false; // If g_bVMPISDKMode has not been set, then VMPI_IsSDKMode just looks for VMPI_Transfer (and doesn't check the command line).
  87. int g_nBytesSent = 0;
  88. int g_nMessagesSent = 0;
  89. int g_nBytesReceived = 0;
  90. int g_nMessagesReceived = 0;
  91. int g_nMulticastBytesSent = 0;
  92. int g_nMulticastBytesReceived = 0;
  93. CUtlLinkedList<VMPI_Disconnect_Handler,int> g_DisconnectHandlers;
  94. bool g_bUseMPI = false;
  95. int g_iVMPIVerboseLevel = 0;
  96. bool g_bMPIMaster = false;
  97. bool g_bMPI_Stats = false;
  98. bool g_bMPI_StatsTextOutput = false;
  99. char g_CurrentStageString[128] = "";
  100. CCriticalSection g_CurrentStageCS;
  101. char g_MasterExeName[MAX_PATH];
  102. bool g_bReceivedMasterExeName = false;
  103. // Change our window text.
  104. HINSTANCE g_hKernel32DLL = NULL;
  105. typedef HWND (*GetConsoleWndFn)();
  106. GetConsoleWndFn g_pConsoleWndFn = NULL;
  107. // ---------------------------------------------------------------------------------------- //
  108. // Classes.
  109. // ---------------------------------------------------------------------------------------- //
  110. // This class is used while discovering what files the workers need.
  111. class CDependencyInfo
  112. {
  113. public:
  114. class CDependencyFile
  115. {
  116. public:
  117. char m_Name[MAX_PATH];
  118. };
  119. // This is the directory where the dependency files live (i.e. all the binaries that the workers need to run the job).
  120. char m_DependencyFilesDir[MAX_PATH];
  121. // "vrad.exe", "vvis.exe", etc.
  122. char m_OriginalExeFilename[MAX_PATH];
  123. CUtlVector<CDependencyFile*> m_Files;
  124. public:
  125. CDependencyFile* FindFile( const char *pFilename )
  126. {
  127. for ( int i=0; i < m_Files.Count(); i++ )
  128. {
  129. if ( stricmp( pFilename, m_Files[i]->m_Name ) == 0 )
  130. return m_Files[i];
  131. }
  132. return NULL;
  133. }
  134. };
  135. class CVMPIConnection : public ITCPSocketHandler
  136. {
  137. public:
  138. CVMPIConnection( int iConnection )
  139. {
  140. m_iConnection = iConnection;
  141. m_pSocket = NULL;
  142. m_bIsAService = false;
  143. char str[512];
  144. Q_snprintf( str, sizeof( str ), "%d", iConnection );
  145. SetMachineName( str );
  146. m_JobWorkerID = 0xFFFFFFFF;
  147. m_bNameSet = false;
  148. }
  149. ~CVMPIConnection()
  150. {
  151. if ( m_pSocket )
  152. m_pSocket->Release();
  153. }
  154. public:
  155. void HandleDisconnect()
  156. {
  157. if ( m_pSocket )
  158. {
  159. // Copy out the error string.
  160. CCriticalSectionLock csLock( &g_ErrorSocketsCS );
  161. csLock.Lock();
  162. char str[512];
  163. Q_strncpy( str, m_ErrorString.Base(), sizeof( str ) );
  164. csLock.Unlock();
  165. // Tell the app.
  166. FOR_EACH_LL( g_DisconnectHandlers, i )
  167. g_DisconnectHandlers[i]( m_iConnection, str );
  168. // Free our socket.
  169. m_pSocket->Release();
  170. m_pSocket = NULL;
  171. }
  172. }
  173. IThreadedTCPSocket* GetSocket()
  174. {
  175. return m_pSocket;
  176. }
  177. void SetMachineName( const char *pName )
  178. {
  179. m_MachineName.CopyArray( pName, strlen( pName ) + 1 );
  180. m_bNameSet = true;
  181. }
  182. const char* GetMachineName()
  183. {
  184. return m_MachineName.Base();
  185. }
  186. bool HasMachineNameBeenSet()
  187. {
  188. return m_bNameSet;
  189. }
  190. // ITCPSocketHandler implementation (thread-safe stuff).
  191. public:
  192. virtual void Init( IThreadedTCPSocket *pSocket )
  193. {
  194. m_pSocket = pSocket;
  195. }
  196. virtual void OnPacketReceived( CTCPPacket *pPacket )
  197. {
  198. // Set who this message came from.
  199. pPacket->SetUserData( m_iConnection );
  200. Assert( m_iConnection >= 0 && m_iConnection < 2048 );
  201. // Store it in the global list.
  202. CCriticalSectionLock csLock( &g_VMPIMessagesCS );
  203. csLock.Lock();
  204. g_VMPIMessages.AddToTail( pPacket );
  205. if ( g_VMPIMessages.Count() == 1 )
  206. g_VMPIMessagesEvent.SetEvent();
  207. }
  208. virtual void OnError( int errorCode, const char *pErrorString )
  209. {
  210. if ( !g_bMPIMaster )
  211. {
  212. Msg( "%s - CVMPIConnection::OnError( %s )\n", GetMachineName(), pErrorString );
  213. }
  214. CCriticalSectionLock csLock( &g_ErrorSocketsCS );
  215. csLock.Lock();
  216. m_ErrorString.CopyArray( pErrorString, strlen( pErrorString ) + 1 );
  217. g_ErrorSockets.AddToTail( this );
  218. // Notify the main thread that a socket is in trouble!
  219. g_ErrorSocketsEvent.SetEvent();
  220. // Make sure the main thread picks up this error soon.
  221. InterlockedIncrement( &m_ErrorSignal );
  222. }
  223. public:
  224. unsigned long m_JobWorkerID;
  225. bool m_bIsAService; // If true, then this is just a service getting the files. Don't count it as an active worker.
  226. CUtlVector<int> m_GroupedChunkLengths;
  227. CUtlVector<void*> m_GroupedChunks;
  228. private:
  229. CUtlVector<char> m_MachineName;
  230. CUtlVector<char> m_ErrorString;
  231. long m_ErrorSignal;
  232. int m_iConnection;
  233. IThreadedTCPSocket *m_pSocket;
  234. bool m_bNameSet;
  235. };
  236. class CVMPIConnectionCreator : public IHandlerCreator
  237. {
  238. public:
  239. virtual ITCPSocketHandler* CreateNewHandler()
  240. {
  241. Assert( g_nConnections < MAX_VMPI_CONNECTIONS );
  242. CVMPIConnection *pRet = new CVMPIConnection( g_nConnections );
  243. g_Connections[g_nConnections++] = pRet;
  244. return pRet;
  245. }
  246. };
  247. // ---------------------------------------------------------------------------------------- //
  248. // Helpers.
  249. // ---------------------------------------------------------------------------------------- //
  250. const char* VMPI_FindArg( int argc, char **argv, const char *pName, const char *pDefault )
  251. {
  252. for ( int i=0; i < argc; i++ )
  253. {
  254. if ( stricmp( argv[i], pName ) == 0 )
  255. {
  256. if ( (i+1) < argc )
  257. return argv[i+1];
  258. else
  259. return pDefault;
  260. }
  261. }
  262. return NULL;
  263. }
  264. void ParseOptions( int argc, char **argv )
  265. {
  266. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_NoTimeout ) ) )
  267. ThreadedTCP_EnableTimeouts( false );
  268. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_DontSetThreadPriorities ) ) )
  269. {
  270. Msg( "%s found.\n", VMPI_GetParamString( mpi_DontSetThreadPriorities ) );
  271. g_bSetThreadPriorities = false;
  272. ThreadedTCP_SetTCPSocketThreadPriorities( false );
  273. }
  274. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_GroupPackets ) ) )
  275. {
  276. Msg( "%s found.\n", VMPI_GetParamString( mpi_GroupPackets ) );
  277. g_bGroupPackets = true;
  278. }
  279. const char *pTransmitRate = VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_FileTransmitRate ), "1" );
  280. if ( pTransmitRate )
  281. {
  282. extern int MULTICAST_TRANSMIT_RATE;
  283. MULTICAST_TRANSMIT_RATE = atoi( pTransmitRate ) * 1024;
  284. }
  285. const char *pVerbose = VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_Verbose ), "1" );
  286. if ( pVerbose )
  287. {
  288. if ( pVerbose[0] == '1' )
  289. g_iVMPIVerboseLevel = 1;
  290. else if ( pVerbose[0] == '2' )
  291. g_iVMPIVerboseLevel = 2;
  292. }
  293. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_Stats ) ) )
  294. g_bMPI_Stats = true;
  295. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_Stats_TextOutput ) ) )
  296. g_bMPI_StatsTextOutput = true;
  297. }
  298. void SetupDependencyFilename( CDependencyInfo *pInfo, const char *pPatchDirectory )
  299. {
  300. char baseExeFilename[512];
  301. if ( !GetModuleFileName( GetModuleHandle( NULL ), baseExeFilename, sizeof( baseExeFilename ) ) )
  302. Error( "GetModuleFileName failed." );
  303. // If they're in patch mode, then the dependency files come out of a directory they've passed in.
  304. // Otherwise, the files come from the same exe dir we're in (like c:\valve\game\bin).
  305. if ( pPatchDirectory )
  306. {
  307. V_strncpy( pInfo->m_DependencyFilesDir, pPatchDirectory, sizeof( pInfo->m_DependencyFilesDir ) );
  308. }
  309. else
  310. {
  311. V_strncpy( pInfo->m_DependencyFilesDir, baseExeFilename, sizeof( pInfo->m_DependencyFilesDir ) );
  312. V_StripLastDir( pInfo->m_DependencyFilesDir, sizeof( pInfo->m_DependencyFilesDir ) );
  313. }
  314. // Get the exe filename.
  315. V_strncpy( pInfo->m_OriginalExeFilename, V_UnqualifiedFileName( baseExeFilename ), sizeof( pInfo->m_OriginalExeFilename ) );
  316. }
  317. bool ReadString( char *pOut, int maxLen, FILE *fp )
  318. {
  319. if ( !fgets( pOut, maxLen, fp ) || pOut[0] == 0 )
  320. return false;
  321. int len = strlen( pOut );
  322. if ( pOut[len - 1] == '\n' )
  323. pOut[len - 1] = 0;
  324. return true;
  325. }
  326. void ParseDependencyFile( CDependencyInfo *pInfo, const char *pDepFilename )
  327. {
  328. FILE *fp = fopen( pDepFilename, "rt" );
  329. if ( !fp )
  330. Error( "Can't find %s.", pDepFilename );
  331. const char *pOptionalPrefix = "optional ";
  332. char tempStr[MAX_PATH];
  333. while ( ReadString( tempStr, sizeof( tempStr ), fp ) )
  334. {
  335. CDependencyInfo::CDependencyFile *pFile = new CDependencyInfo::CDependencyFile;
  336. bool bOptional = false;
  337. if ( strstr( tempStr, "optional " ) == tempStr )
  338. {
  339. bOptional = true;
  340. Q_strncpy( pFile->m_Name, tempStr + strlen( pOptionalPrefix ), sizeof( pFile->m_Name ) );
  341. }
  342. else
  343. {
  344. Q_strncpy( pFile->m_Name, tempStr, sizeof( pFile->m_Name ) );
  345. }
  346. // Now get the file info.
  347. char fullFilename[MAX_PATH];
  348. V_ComposeFileName( pInfo->m_DependencyFilesDir, pFile->m_Name, fullFilename, sizeof( fullFilename ) );
  349. if ( _access( fullFilename, 0 ) == 0 )
  350. {
  351. pInfo->m_Files.AddToTail( pFile );
  352. }
  353. else
  354. {
  355. delete pFile;
  356. if ( !bOptional )
  357. Error( "Can't find %s (listed in %s).", fullFilename, pDepFilename );
  358. }
  359. }
  360. fclose( fp );
  361. }
  362. void SetupDependenciesForPatch( CDependencyInfo *pInfo, const char *pPatchDirectory )
  363. {
  364. char searchStr[MAX_PATH];
  365. V_ComposeFileName( pPatchDirectory, "*.*", searchStr, sizeof( searchStr ) );
  366. _finddata_t data;
  367. long handle = _findfirst( searchStr, &data );
  368. if ( handle != -1 )
  369. {
  370. do
  371. {
  372. if ( data.name[0] == '.' || (data.attrib & _A_SUBDIR) != 0 )
  373. continue;
  374. CDependencyInfo::CDependencyFile *pFile = new CDependencyInfo::CDependencyFile;
  375. V_strncpy( pFile->m_Name, data.name, sizeof( pFile->m_Name ) );
  376. pInfo->m_Files.AddToTail( pFile );
  377. } while( _findnext( handle, &data ) == 0 );
  378. _findclose( handle );
  379. }
  380. }
  381. void SetupDependencyInfo( CDependencyInfo *pInfo, const char *pDependencyFilename, bool bPatchMode )
  382. {
  383. if ( bPatchMode )
  384. {
  385. const char *pPatchDirectory = pDependencyFilename;
  386. SetupDependencyFilename( pInfo, pPatchDirectory );
  387. SetupDependenciesForPatch( pInfo, pPatchDirectory );
  388. }
  389. else
  390. {
  391. SetupDependencyFilename( pInfo, NULL );
  392. // Parse the dependency file.
  393. char depFilename[MAX_PATH];
  394. V_ComposeFileName( pInfo->m_DependencyFilesDir, pDependencyFilename, depFilename, sizeof( depFilename ) );
  395. ParseDependencyFile( pInfo, depFilename );
  396. }
  397. }
  398. int GetCurMicrosecondsAndSleep( int sleepLen )
  399. {
  400. Sleep( sleepLen );
  401. CCycleCount cnt;
  402. cnt.Sample();
  403. return cnt.GetMicroseconds();
  404. }
  405. void CountActiveConnections( int *nRegularWorkers, int *nServiceDownloaders )
  406. {
  407. *nRegularWorkers = *nServiceDownloaders = 0;
  408. int nTotalConnections = g_nConnections;
  409. for ( int i=0; i < nTotalConnections; i++ )
  410. {
  411. if ( VMPI_IsProcConnected( i ) )
  412. {
  413. if ( VMPI_IsProcAService( i ) )
  414. (*nServiceDownloaders)++;
  415. else
  416. (*nRegularWorkers)++;
  417. }
  418. }
  419. }
  420. // In this function, we update the window text to tell how many active workers there are.
  421. void UpdateActiveConnectionsText()
  422. {
  423. if ( !g_bMPIMaster || !g_pConsoleWndFn )
  424. return;
  425. HWND hWnd = g_pConsoleWndFn();
  426. if ( !hWnd )
  427. return;
  428. int nRegularWorkers, nDownloaders;
  429. CountActiveConnections( &nRegularWorkers, &nDownloaders );
  430. char str[512];
  431. if ( g_bVMPISDKMode )
  432. {
  433. V_snprintf( str, sizeof( str ), "VMPI (SDK) - Workers: %d", nRegularWorkers );
  434. }
  435. else
  436. {
  437. V_snprintf( str, sizeof( str ), "VMPI - Workers: %d, Downloaders: %d", nRegularWorkers, nDownloaders );
  438. }
  439. SetWindowText( hWnd, str );
  440. }
  441. void VMPI_SendMachineNameTo( int iProc )
  442. {
  443. const char *pMyName = VMPI_GetLocalMachineName();
  444. unsigned char packetData[512];
  445. packetData[0] = VMPI_INTERNAL_PACKET_ID;
  446. packetData[1] = VMPI_INTERNAL_SUBPACKET_MACHINE_NAME;
  447. Q_strncpy( (char*)&packetData[2], pMyName, sizeof( packetData ) - 2 );
  448. VMPI_SendData( packetData, 2 + strlen( pMyName ) + 1, iProc );
  449. }
  450. static CVMPIConnection* FindConnectionBySocket( IThreadedTCPSocket *pSocket, bool bLockConnections )
  451. {
  452. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  453. if ( bLockConnections )
  454. connectionsLock.Lock();
  455. for ( int i=0; i < g_nConnections; i++ )
  456. if ( g_Connections[i]->GetSocket() == pSocket )
  457. return g_Connections[i];
  458. return NULL;
  459. }
  460. static char* CopyString( const char *pStr )
  461. {
  462. int len = V_strlen( pStr ) + 1;
  463. char *pArg = new char[len];
  464. Q_strncpy( pArg, pStr, len );
  465. return pArg;
  466. }
  467. // ---------------------------------------------------------------------------------------- //
  468. // Internal VMPI dispatch..
  469. // ---------------------------------------------------------------------------------------- //
  470. void VMPI_SetMachineName( int iProc, const char *pName );
  471. CUtlVector<char*> g_WorkerCommandLine;
  472. bool g_bReceivedWorkerCommandLine = false;
  473. bool VMPI_InternalDispatchFn( MessageBuffer *pBuf, int iSource, int iPacketID )
  474. {
  475. if ( pBuf->getLen() >= 2 )
  476. {
  477. if ( pBuf->data[1] == VMPI_INTERNAL_SUBPACKET_MACHINE_NAME )
  478. {
  479. if ( pBuf->getLen() >= 3 )
  480. {
  481. VMPI_SetMachineName( iSource, &pBuf->data[2] );
  482. return true;
  483. }
  484. }
  485. else if ( pBuf->data[1] == VMPI_INTERNAL_SUBPACKET_WAITING_FOR_COMMAND_LINE )
  486. {
  487. if ( !VMPI_IsSDKMode() )
  488. {
  489. Warning( "Worker %d is running in SDK mode (and the master is not)!\n", iSource );
  490. }
  491. return true;
  492. }
  493. else if ( pBuf->data[1] == VMPI_INTERNAL_SUBPACKET_COMMAND_LINE )
  494. {
  495. pBuf->setOffset( 2 );
  496. int nArgs;
  497. pBuf->read( &nArgs, sizeof( nArgs ) );
  498. for ( int i=0; i < nArgs; i++ )
  499. {
  500. char str[4096];
  501. if ( pBuf->ReadString( str, sizeof( str ) ) == -1 )
  502. Error( "Error in ReadString() while reading command line." );
  503. g_WorkerCommandLine.AddToTail( CopyString( str ) );
  504. }
  505. g_bReceivedWorkerCommandLine = true;
  506. return true;
  507. }
  508. else if ( pBuf->data[1] == VMPI_INTERNAL_SUBPACKET_VERIFY_EXE_NAME )
  509. {
  510. pBuf->setOffset( 2 );
  511. if ( pBuf->ReadString( g_MasterExeName, sizeof( g_MasterExeName ) ) == -1 )
  512. Error( "Error in ReadString() while reading VMPI_INTERNAL_SUBPACKET_VERIFY_EXE_NAME." );
  513. g_bReceivedMasterExeName = true;
  514. return true;
  515. }
  516. else if ( pBuf->data[1] == VMPI_INTERNAL_SUBPACKET_TIMING_WAIT_DONE )
  517. {
  518. g_bTimingWaitDone = true;
  519. return true;
  520. }
  521. }
  522. return false;
  523. }
  524. CDispatchReg g_VMPIInternalDispatchReg( VMPI_INTERNAL_PACKET_ID, VMPI_InternalDispatchFn ); // register to handle the messages we want
  525. void VMPI_SendCommandLine( int argc, char **argv )
  526. {
  527. MessageBuffer mb;
  528. char cPacketHeader[2] = {VMPI_INTERNAL_PACKET_ID, VMPI_INTERNAL_SUBPACKET_COMMAND_LINE};
  529. mb.write( cPacketHeader, sizeof( cPacketHeader ) );
  530. mb.write( &argc, sizeof( argc ) );
  531. for ( int i=0; i < argc; i++ )
  532. mb.WriteString( argv[i] );
  533. VMPI_SendData( mb.data, mb.getLen(), VMPI_PERSISTENT );
  534. }
  535. void VMPI_ReceiveCommandLine()
  536. {
  537. // For verification purposes, tell the master we're trying to get the command line.
  538. unsigned char chData[2] = {VMPI_INTERNAL_PACKET_ID, VMPI_INTERNAL_SUBPACKET_WAITING_FOR_COMMAND_LINE};
  539. VMPI_SendData( chData, sizeof( chData ), VMPI_MASTER_ID );
  540. double startTime = Plat_FloatTime();
  541. while ( !g_bReceivedWorkerCommandLine )
  542. {
  543. if ( Plat_FloatTime() - startTime > 30 )
  544. Error( "VMPI_ReceiveCommandLine: timeout. Is the master running in SDK mode?" );
  545. VMPI_DispatchNextMessage( 10 * 1000 );
  546. }
  547. }
  548. void VMPI_SendExeName()
  549. {
  550. MessageBuffer mb;
  551. char cPacketHeader[2] = {VMPI_INTERNAL_PACKET_ID, VMPI_INTERNAL_SUBPACKET_VERIFY_EXE_NAME};
  552. mb.write( cPacketHeader, sizeof( cPacketHeader ) );
  553. char baseExeFilename[MAX_PATH], fileBase[MAX_PATH];
  554. if ( !GetModuleFileName( GetModuleHandle( NULL ), baseExeFilename, sizeof( baseExeFilename ) ) )
  555. Error( "VMPI_CheckSDKMode -> GetModuleFileName failed." );
  556. V_FileBase( baseExeFilename, fileBase, sizeof( fileBase ) );
  557. mb.WriteString( fileBase );
  558. VMPI_SendData( mb.data, mb.getLen(), VMPI_PERSISTENT );
  559. }
  560. void VMPI_ReceiveExeName()
  561. {
  562. double startTime = Plat_FloatTime();
  563. while ( !g_bReceivedMasterExeName )
  564. {
  565. if ( Plat_FloatTime() - startTime > 30 )
  566. Error( "VMPI_ReceiveExeName: timeout." );
  567. VMPI_DispatchNextMessage( 10 * 1000 );
  568. }
  569. // Now compare the exe name we got with our own.
  570. char baseExeFilename[MAX_PATH], fileBase[MAX_PATH];
  571. if ( !GetModuleFileName( GetModuleHandle( NULL ), baseExeFilename, sizeof( baseExeFilename ) ) )
  572. Error( "VMPI_CheckSDKMode -> GetModuleFileName failed." );
  573. // Unless we're a vmpi_transfer.. vmpi_transfer can always connect.
  574. V_FileBase( baseExeFilename, fileBase, sizeof( fileBase ) );
  575. if ( V_stricmp( fileBase, "vmpi_transfer" ) != 0 )
  576. {
  577. if ( V_stricmp( fileBase, g_MasterExeName ) != 0 )
  578. {
  579. Error( "VMPI_ReceiveExeName: mismatched exe names (master: %s, me: %s).\nThis usually just means the master finished"
  580. " a job like vvis really fast and started a vrad immediately, and an old vvis worker connected to the new vrad job.",
  581. g_MasterExeName, fileBase );
  582. }
  583. }
  584. }
  585. // ---------------------------------------------------------------------------------------- //
  586. // CMasterBroadcaster
  587. // This class broadcasts messages looking for workers. The app updates it as often as possible
  588. // and it'll add workers as necessary.
  589. // ---------------------------------------------------------------------------------------- //
  590. #define MASTER_BROADCAST_INTERVAL 600 // Send every N milliseconds.
  591. class CMasterBroadcaster
  592. {
  593. public:
  594. CMasterBroadcaster();
  595. ~CMasterBroadcaster();
  596. bool Init( int argc, char **argv, const char *pDependencyFilename, int nMaxWorkers, VMPIRunMode runMode, bool bPatchMode );
  597. void Term();
  598. // What port is it listening on?
  599. int GetListenPort() const;
  600. // These can be used to allow more workers on or filter who's able to connect
  601. int GetMaxWorkers() const;
  602. void IncreaseMaxWorkers( int count );
  603. void SetPassword( const char *pPassword );
  604. void SetNoTimeoutOption();
  605. private:
  606. void GetPatchWorkerList( int argc, char **argv );
  607. private:
  608. class CMasterBroadcastInfo
  609. {
  610. public:
  611. int m_JobID[4];
  612. char m_Password[256];
  613. char m_WorkerExeFilename[MAX_PATH];
  614. CUtlVector<char*> m_Args;
  615. char m_PatchVersion[32]; // 0 if not patching.
  616. bool m_bForcePatch;
  617. };
  618. void ThreadFn();
  619. static DWORD WINAPI StaticThreadFn( LPVOID lpParameter );
  620. bool Update();
  621. void BuildBroadcastPacket( bf_write &buf );
  622. private:
  623. ITCPConnectSocket *m_pListenSocket;
  624. ITCPConnectSocket *m_pDownloaderListenSocket;
  625. ISocket *m_pSocket;
  626. DWORD m_LastSendTime;
  627. CMasterBroadcastInfo m_BroadcastInfo;
  628. CUtlVector<CIPAddr> m_PatchWorkerIPs; // If in patch mode, these are the IPs we send the job request to (instead of broadcasting).
  629. bool m_bPatching;
  630. CVMPIConnectionCreator m_ConnectionCreator;
  631. int m_nMaxWorkers;
  632. HANDLE m_hThread;
  633. CEvent m_hShutdownEvent;
  634. CEvent m_hShutdownReply;
  635. VMPIRunMode m_RunMode;
  636. int m_iListenPort;
  637. int m_iDownloaderListenPort;
  638. };
  639. CMasterBroadcaster::CMasterBroadcaster()
  640. {
  641. m_pListenSocket = NULL;
  642. m_pDownloaderListenSocket = NULL;
  643. m_pSocket = NULL;
  644. m_iListenPort = -1;
  645. m_iDownloaderListenPort = -1;
  646. }
  647. CMasterBroadcaster::~CMasterBroadcaster()
  648. {
  649. Term();
  650. }
  651. void CMasterBroadcaster::GetPatchWorkerList( int argc, char **argv )
  652. {
  653. m_PatchWorkerIPs.Purge();
  654. for ( int i=0; i < argc-1; i++ )
  655. {
  656. if ( V_stricmp( argv[i], "-mpi_PatchWorkers" ) == 0 )
  657. {
  658. int workerCount = atoi( argv[i+1] );
  659. for ( int iWorker=0; iWorker < workerCount; iWorker++ )
  660. {
  661. int iArg = i+2 + iWorker;
  662. if ( iArg >= argc )
  663. Error( "-mpi_PatchWorkers: %d specified for count, but not enough IPs following.\n", workerCount );
  664. int a, b, c, d;
  665. const char *pArg = argv[iArg];
  666. sscanf( pArg, "%d.%d.%d.%d", &a, &b, &c, &d );
  667. CIPAddr addr;
  668. addr.Init( a, b, c, d, 0 );
  669. m_PatchWorkerIPs.AddToTail( addr );
  670. }
  671. return;
  672. }
  673. }
  674. }
  675. bool CMasterBroadcaster::Init(
  676. int argc,
  677. char **argv,
  678. const char *pDependencyFilename,
  679. int nMaxWorkers,
  680. VMPIRunMode runMode,
  681. bool bPatchMode )
  682. {
  683. m_RunMode = runMode;
  684. m_nMaxWorkers = nMaxWorkers;
  685. // Open the file that tells us which binaries we depend on.
  686. CDependencyInfo dependencyInfo;
  687. if ( m_RunMode == VMPI_RUN_NETWORKED && !g_bVMPISDKMode )
  688. {
  689. SetupDependencyInfo( &dependencyInfo, pDependencyFilename, bPatchMode );
  690. }
  691. m_pListenSocket = NULL;
  692. m_pDownloaderListenSocket = NULL;
  693. const char *pPortStr = VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_Port ) );
  694. if ( pPortStr )
  695. {
  696. m_iListenPort = atoi( pPortStr );
  697. m_iDownloaderListenPort = m_iListenPort + 1;
  698. m_pListenSocket = ThreadedTCP_CreateListener( &m_ConnectionCreator, m_iListenPort );
  699. if ( !g_bVMPISDKMode )
  700. {
  701. m_pDownloaderListenSocket = ThreadedTCP_CreateListener( &m_ConnectionCreator, m_iDownloaderListenPort );
  702. }
  703. }
  704. else
  705. {
  706. // Create a socket to listen on.
  707. CCycleCount cnt;
  708. cnt.Sample();
  709. int iTime = (int)cnt.GetMicroseconds();
  710. srand( (unsigned)iTime );
  711. for ( int iTest=VMPI_MASTER_FIRST_PORT; iTest <= VMPI_MASTER_LAST_PORT; iTest++ )
  712. {
  713. m_iListenPort = iTest;
  714. m_pListenSocket = ThreadedTCP_CreateListener( &m_ConnectionCreator, m_iListenPort );
  715. if ( m_pListenSocket )
  716. break;
  717. }
  718. // No need to create the downloader in SDK mode.
  719. if ( m_pListenSocket && !g_bVMPISDKMode )
  720. {
  721. for ( int iTest=m_iListenPort+1; iTest <= VMPI_MASTER_LAST_PORT; iTest++ )
  722. {
  723. m_iDownloaderListenPort = iTest;
  724. if ( m_iDownloaderListenPort == m_iListenPort )
  725. continue;
  726. m_pDownloaderListenSocket = ThreadedTCP_CreateListener( &m_ConnectionCreator, m_iDownloaderListenPort );
  727. if ( m_pDownloaderListenSocket )
  728. break;
  729. }
  730. }
  731. }
  732. if ( !m_pListenSocket || (!g_bVMPISDKMode && !m_pDownloaderListenSocket) )
  733. {
  734. Error( "Can't bind a listen socket in port range [%d, %d].", VMPI_MASTER_PORT_FIRST, VMPI_MASTER_PORT_LAST );
  735. }
  736. // Create a socket to broadcast from unless we're in the SDK in which case we don't broadcast.
  737. m_bPatching = false;
  738. if ( m_RunMode == VMPI_RUN_NETWORKED && !g_bVMPISDKMode )
  739. {
  740. m_pSocket = CreateIPSocket();
  741. if ( !m_pSocket->BindToAny( 0 ) )
  742. Error( "MPI_Init_Master: can't bind a socket" );
  743. m_BroadcastInfo.m_bForcePatch = false;
  744. if ( bPatchMode )
  745. {
  746. m_bPatching = true;
  747. if ( VMPI_FindArg( argc, argv, "-mpi_ForcePatch", NULL ) )
  748. m_BroadcastInfo.m_bForcePatch = true;
  749. const char *pArg = VMPI_FindArg( argc, argv, "-mpi_PatchVersion", "0" );
  750. float iPatchVersion = atof( pArg );
  751. if ( iPatchVersion <= 0 || iPatchVersion >= ((1 << 15) - 1) )
  752. {
  753. Error( "-mpi_PatchVersion <val> - val must be between 1.0 and 32767.0" );
  754. }
  755. V_strncpy( m_BroadcastInfo.m_PatchVersion, pArg, sizeof( m_BroadcastInfo.m_PatchVersion ) );
  756. }
  757. else
  758. {
  759. m_BroadcastInfo.m_PatchVersion[0] = 0;
  760. }
  761. // Come up with a unique job ID.
  762. m_BroadcastInfo.m_JobID[0] = GetCurMicrosecondsAndSleep( 1 );
  763. m_BroadcastInfo.m_JobID[1] = GetCurMicrosecondsAndSleep( 1 );
  764. m_BroadcastInfo.m_JobID[2] = GetCurMicrosecondsAndSleep( 1 );
  765. m_BroadcastInfo.m_JobID[3] = GetCurMicrosecondsAndSleep( 1 );
  766. const char *pPassword = VMPI_FindArg( argc, argv, "-mpi_pw", "" );
  767. Q_strncpy( m_BroadcastInfo.m_Password, pPassword ? pPassword : "", sizeof( m_BroadcastInfo.m_Password ) );
  768. Q_strncpy( m_BroadcastInfo.m_WorkerExeFilename, dependencyInfo.m_OriginalExeFilename, sizeof( m_BroadcastInfo.m_WorkerExeFilename ) );
  769. // Store the command-line args.
  770. m_BroadcastInfo.m_Args.Purge();
  771. for ( int i=1; i < argc; i++ )
  772. {
  773. m_BroadcastInfo.m_Args.AddToTail( CopyString( argv[i] ) );
  774. }
  775. // 0th arg is the exe name.
  776. m_BroadcastInfo.m_Args.InsertBefore( 0, CopyString( m_BroadcastInfo.m_WorkerExeFilename ) );
  777. // Now add arguments for each file they need to transmit. The service will use this to get all the files from the master before it starts the app.
  778. for ( int i=0; i < dependencyInfo.m_Files.Count(); i++ )
  779. {
  780. m_BroadcastInfo.m_Args.InsertAfter( 0, "-mpi_file" );
  781. m_BroadcastInfo.m_Args.InsertAfter( 1, CopyString( dependencyInfo.m_Files[i]->m_Name ) );
  782. }
  783. // Add -mpi_filebase so it can use absolute paths with the filesystem so we get the exact right set of files.
  784. m_BroadcastInfo.m_Args.InsertAfter( 0, "-mpi_filebase" );
  785. m_BroadcastInfo.m_Args.InsertAfter( 1, CopyString( dependencyInfo.m_DependencyFilesDir ) );
  786. if ( bPatchMode )
  787. {
  788. GetPatchWorkerList( argc, argv );
  789. }
  790. }
  791. // Add ourselves as the first process (rank 0).
  792. m_ConnectionCreator.CreateNewHandler();
  793. // Initiate as many connections as we can for a few seconds.
  794. m_LastSendTime = Plat_MSTime() - MASTER_BROADCAST_INTERVAL*2;
  795. m_hShutdownEvent.Init( false, false );
  796. m_hShutdownReply.Init( false, false );
  797. DWORD dwThreadID = 0;
  798. m_hThread = CreateThread(
  799. NULL,
  800. 0,
  801. &CMasterBroadcaster::StaticThreadFn,
  802. this,
  803. 0,
  804. &dwThreadID );
  805. if ( m_hThread )
  806. {
  807. SetThreadPriority( m_hThread, THREAD_PRIORITY_HIGHEST );
  808. return true;
  809. }
  810. else
  811. {
  812. return false;
  813. }
  814. }
  815. void CMasterBroadcaster::BuildBroadcastPacket( bf_write &buf )
  816. {
  817. // Broadcast out to tell all the machines we want workers.
  818. buf.WriteByte( VMPI_PROTOCOL_VERSION );
  819. buf.WriteString( m_BroadcastInfo.m_Password );
  820. if ( m_BroadcastInfo.m_PatchVersion[0] == 0 )
  821. buf.WriteByte( VMPI_LOOKING_FOR_WORKERS );
  822. else
  823. buf.WriteByte( VMPI_SERVICE_PATCH );
  824. buf.WriteString( m_BroadcastInfo.m_PatchVersion );
  825. buf.WriteLong( m_iListenPort ); // Tell the port that we're listening on.
  826. buf.WriteLong( m_BroadcastInfo.m_JobID[0] );
  827. buf.WriteLong( m_BroadcastInfo.m_JobID[1] );
  828. buf.WriteLong( m_BroadcastInfo.m_JobID[2] );
  829. buf.WriteLong( m_BroadcastInfo.m_JobID[3] );
  830. buf.WriteWord( m_BroadcastInfo.m_Args.Count() + 2 );
  831. // Write the alternate exe name.
  832. buf.WriteString( m_BroadcastInfo.m_WorkerExeFilename );
  833. // Write the machine name of the master into the command line. It's ignored by the code, but it's useful
  834. // if a job crashes the workers - by looking at the command line in vmpi_service, you can see who ran the job.
  835. buf.WriteString( "-mpi_MasterName" );
  836. buf.WriteString( VMPI_GetLocalMachineName() );
  837. for ( int i=1; i < m_BroadcastInfo.m_Args.Count(); i++ )
  838. buf.WriteString( m_BroadcastInfo.m_Args[i] );
  839. buf.WriteByte( (unsigned char)m_BroadcastInfo.m_bForcePatch );
  840. buf.WriteShort( m_iDownloaderListenPort ); // Tell the port that we're listening for downloaders on.
  841. }
  842. bool CMasterBroadcaster::Update()
  843. {
  844. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  845. connectionsLock.Lock();
  846. // Don't accept any more connections when we've hit the limit.
  847. int nActiveConnections, nServiceDownloaders;
  848. CountActiveConnections( &nActiveConnections, &nServiceDownloaders );
  849. if ( nActiveConnections >= m_nMaxWorkers )
  850. return false;
  851. // Only broadcast our presence so often.
  852. if ( m_pSocket )
  853. {
  854. DWORD curTime = Plat_MSTime();
  855. if ( curTime - m_LastSendTime >= MASTER_BROADCAST_INTERVAL )
  856. {
  857. char packetData[512];
  858. bf_write packetBuf( "packetBuf", packetData, sizeof( packetData ) );
  859. BuildBroadcastPacket( packetBuf );
  860. for ( int iBroadcastPort=VMPI_SERVICE_PORT; iBroadcastPort <= VMPI_LAST_SERVICE_PORT; iBroadcastPort++ )
  861. {
  862. if ( m_bPatching )
  863. {
  864. // Only send to this specific list of workers if necessary.
  865. for ( int i=0; i < m_PatchWorkerIPs.Count(); i++ )
  866. {
  867. CIPAddr addr = m_PatchWorkerIPs[i];
  868. addr.port = iBroadcastPort;
  869. m_pSocket->SendTo( &addr, packetBuf.GetBasePointer(), packetBuf.GetNumBytesWritten() );
  870. }
  871. }
  872. else
  873. {
  874. m_pSocket->Broadcast( packetBuf.GetBasePointer(), packetBuf.GetNumBytesWritten(), iBroadcastPort );
  875. }
  876. }
  877. // We don't want them to keep patching over and over.
  878. if ( m_PatchWorkerIPs.Count() > 0 && m_BroadcastInfo.m_bForcePatch )
  879. m_PatchWorkerIPs.Purge();
  880. m_LastSendTime = curTime;
  881. }
  882. }
  883. // First look for normal workers.
  884. IThreadedTCPSocket *pNewConn = NULL;
  885. bool bRet = m_pListenSocket->Update( &pNewConn, 0 );
  886. // Now look for downloaders.
  887. if ( !bRet || !pNewConn )
  888. {
  889. if ( m_pDownloaderListenSocket )
  890. {
  891. int nDownloadersAllowed = (m_nMaxWorkers - nActiveConnections) + 8; // Don't allow too many downloaders.
  892. if ( nServiceDownloaders < nDownloadersAllowed )
  893. bRet = m_pDownloaderListenSocket->Update( &pNewConn, 0 );
  894. }
  895. }
  896. if ( bRet && pNewConn )
  897. {
  898. // Mark this guy as a downloader if necessary.
  899. CIPAddr remoteAddr = pNewConn->GetRemoteAddr();
  900. if ( remoteAddr.port >= VMPI_SERVICE_DOWNLOADER_PORT_FIRST && remoteAddr.port <= VMPI_SERVICE_DOWNLOADER_PORT_LAST )
  901. {
  902. CVMPIConnection *pVMPIConnection = FindConnectionBySocket( pNewConn, false );
  903. if ( pVMPIConnection )
  904. pVMPIConnection->m_bIsAService = true;
  905. }
  906. // Send this guy all the persistent packets.
  907. CCriticalSectionLock csLock( &g_PersistentPacketsCS );
  908. csLock.Lock();
  909. FOR_EACH_LL( g_PersistentPackets, iPacket )
  910. {
  911. PersistentPacket *pPacket = g_PersistentPackets[iPacket];
  912. VMPI_SendData( pPacket->Base(), pPacket->Count(), g_nConnections-1 );
  913. }
  914. UpdateActiveConnectionsText();
  915. return true;
  916. }
  917. else
  918. {
  919. return false;
  920. }
  921. }
  922. void CMasterBroadcaster::ThreadFn()
  923. {
  924. // Update every 100ms or until the main thread tells us to go away.
  925. while ( WaitForSingleObject( m_hShutdownEvent.GetEventHandle(), 20 ) == WAIT_TIMEOUT )
  926. {
  927. DWORD startTime = GetTickCount();
  928. while ( Update() && (GetTickCount() - startTime) < 500 )
  929. {
  930. }
  931. }
  932. m_hShutdownReply.SetEvent();
  933. }
  934. DWORD CMasterBroadcaster::StaticThreadFn( LPVOID lpParameter )
  935. {
  936. ((CMasterBroadcaster*)lpParameter)->ThreadFn();
  937. return 0;
  938. }
  939. void CMasterBroadcaster::Term()
  940. {
  941. // Shutdown the update thread.
  942. if ( m_hThread )
  943. {
  944. m_hShutdownEvent.SetEvent();
  945. WaitForSingleObject( m_hThread, INFINITE );
  946. CloseHandle( m_hThread );
  947. m_hThread = 0;
  948. }
  949. if ( m_pSocket )
  950. {
  951. m_pSocket->Release();
  952. m_pSocket = NULL;
  953. }
  954. if ( m_pListenSocket )
  955. {
  956. m_pListenSocket->Release();
  957. m_pListenSocket = NULL;
  958. }
  959. if ( m_pDownloaderListenSocket )
  960. {
  961. m_pDownloaderListenSocket->Release();
  962. m_pDownloaderListenSocket = NULL;
  963. }
  964. m_iListenPort = -1;
  965. m_iDownloaderListenPort = -1;
  966. }
  967. int CMasterBroadcaster::GetListenPort() const
  968. {
  969. return m_iListenPort;
  970. }
  971. int CMasterBroadcaster::GetMaxWorkers() const
  972. {
  973. return m_nMaxWorkers;
  974. }
  975. void CMasterBroadcaster::IncreaseMaxWorkers( int count )
  976. {
  977. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  978. connectionsLock.Lock();
  979. m_nMaxWorkers = min( MAX_VMPI_CONNECTIONS, m_nMaxWorkers + count );
  980. }
  981. void CMasterBroadcaster::SetPassword( const char *pPassword )
  982. {
  983. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  984. connectionsLock.Lock();
  985. Q_strncpy( m_BroadcastInfo.m_Password, pPassword, sizeof( m_BroadcastInfo.m_Password ) );
  986. }
  987. void CMasterBroadcaster::SetNoTimeoutOption()
  988. {
  989. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  990. connectionsLock.Lock();
  991. // Don't re-add the option if it's already there.
  992. for ( int i=1; i < m_BroadcastInfo.m_Args.Count(); i++ )
  993. {
  994. if ( Q_stricmp( m_BroadcastInfo.m_Args[i], VMPI_GetParamString( mpi_NoTimeout ) ) == 0 )
  995. return;
  996. }
  997. m_BroadcastInfo.m_Args.InsertAfter( 0, (char*)VMPI_GetParamString( mpi_NoTimeout ) );
  998. }
  999. CMasterBroadcaster g_MasterBroadcaster;
  1000. // ---------------------------------------------------------------------------------------- //
  1001. // CDispatchReg.
  1002. // ---------------------------------------------------------------------------------------- //
  1003. CDispatchReg::CDispatchReg( int iPacketID, VMPIDispatchFn fn )
  1004. {
  1005. Assert( iPacketID >= 0 && iPacketID < MAX_VMPI_PACKET_IDS );
  1006. Assert( !g_VMPIDispatch[iPacketID] );
  1007. g_VMPIDispatch[iPacketID] = fn;
  1008. }
  1009. void VMPI_HandleTimingWait_Worker()
  1010. {
  1011. if ( VMPI_IsParamUsed( mpi_TimingWait ) )
  1012. {
  1013. Msg( "-mpi_TimingWait specified. Waiting for master to start..." );
  1014. // Wait for the signal to go.
  1015. while ( !g_bTimingWaitDone )
  1016. {
  1017. VMPI_DispatchNextMessage( 50 );
  1018. }
  1019. Msg( "\n ");
  1020. }
  1021. }
  1022. void VMPI_HandleTimingWait_Master()
  1023. {
  1024. if ( VMPI_IsParamUsed( mpi_TimingWait ) )
  1025. {
  1026. Msg( "-mpi_TimingWait specified. Waiting for a keypress to continue... " );
  1027. getch();
  1028. Msg( "\n" );
  1029. unsigned char cPacket[2] = { VMPI_INTERNAL_PACKET_ID, VMPI_INTERNAL_SUBPACKET_TIMING_WAIT_DONE };
  1030. VMPI_SendData( cPacket, sizeof( cPacket ), VMPI_PERSISTENT );
  1031. }
  1032. }
  1033. // ---------------------------------------------------------------------------------------- //
  1034. // Helpers.
  1035. // ---------------------------------------------------------------------------------------- //
  1036. bool MPI_Init_Worker( int &argc, char **&argv, const CIPAddr &masterAddr, bool bConnectingAsService )
  1037. {
  1038. g_bMPIMaster = false;
  1039. // Make a connector to try connect to the master.
  1040. CVMPIConnectionCreator connectionCreator;
  1041. int iFirstPort = VMPI_WORKER_PORT_FIRST;
  1042. int iLastPort = VMPI_WORKER_PORT_LAST;
  1043. if ( bConnectingAsService )
  1044. {
  1045. iFirstPort = VMPI_SERVICE_DOWNLOADER_PORT_FIRST;
  1046. iLastPort = VMPI_SERVICE_DOWNLOADER_PORT_LAST;
  1047. }
  1048. // Now wait for a connection.
  1049. int nAttempts = 1;
  1050. Retry:;
  1051. ITCPConnectSocket *pConnectSocket = NULL;
  1052. int iPort;
  1053. for ( iPort=iFirstPort; iPort <= iLastPort; iPort++ )
  1054. {
  1055. pConnectSocket = ThreadedTCP_CreateConnector(
  1056. masterAddr,
  1057. CIPAddr( 0, 0, 0, 0, iPort ),
  1058. &connectionCreator );
  1059. if ( pConnectSocket )
  1060. break;
  1061. }
  1062. if ( !pConnectSocket )
  1063. {
  1064. Error( "Can't bind a port in range [%d, %d].", iFirstPort, iLastPort );
  1065. }
  1066. CWaitTimer wait( 3 );
  1067. while ( 1 )
  1068. {
  1069. IThreadedTCPSocket *pSocket = NULL;
  1070. if ( pConnectSocket->Update( &pSocket, 100 ) )
  1071. {
  1072. if ( pSocket )
  1073. {
  1074. // Send the master our machine name.
  1075. VMPI_SendMachineNameTo( VMPI_MASTER_ID );
  1076. // Verify that the exe is correct.
  1077. VMPI_ReceiveExeName();
  1078. if ( g_bVMPISDKMode )
  1079. {
  1080. VMPI_ReceiveCommandLine();
  1081. CommandLine()->CreateCmdLine( g_WorkerCommandLine.Count(), g_WorkerCommandLine.Base() );
  1082. argc = g_WorkerCommandLine.Count();
  1083. argv = g_WorkerCommandLine.Base();
  1084. }
  1085. ParseOptions( g_WorkerCommandLine.Count(), g_WorkerCommandLine.Base() );
  1086. for ( int i=0; i < g_WorkerCommandLine.Count(); i++ )
  1087. {
  1088. Msg( "arg %d: %s\n", i, g_WorkerCommandLine[i] );
  1089. }
  1090. VMPI_HandleTimingWait_Worker();
  1091. return true;
  1092. }
  1093. }
  1094. else
  1095. {
  1096. pConnectSocket->Release();
  1097. Error( "ITCPConnectSocket::Update() errored out" );
  1098. }
  1099. if( wait.ShouldKeepWaiting() )
  1100. Sleep( 100 );
  1101. else
  1102. break;
  1103. };
  1104. // Never made a connection, shucks.
  1105. pConnectSocket->Release();
  1106. if ( VMPI_IsParamUsed( mpi_Retry ) )
  1107. {
  1108. Msg( "%s found. Retrying connection to %d.%d.%d.%d:%d (attempt %d).\n", VMPI_GetParamString( mpi_Retry ), masterAddr.ip[0], masterAddr.ip[1], masterAddr.ip[2], masterAddr.ip[3], masterAddr.port, nAttempts++ );
  1109. goto Retry;
  1110. }
  1111. return false;
  1112. }
  1113. bool SpawnLocalWorker( int argc, char **argv, int iListenPort, bool bShowConsoleWindow )
  1114. {
  1115. char commandLine[4096];
  1116. commandLine[0] = 0;
  1117. // Add the -mpi_worker argument in, then launch the process.
  1118. for ( int i=0; i < 9999999; i++ )
  1119. {
  1120. char argStr[512];
  1121. if ( i == 1 )
  1122. {
  1123. Q_snprintf( argStr, sizeof( argStr ), "-mpi_worker 127.0.0.1:%d ", iListenPort );
  1124. Q_strncat( commandLine, argStr, sizeof( commandLine ), COPY_ALL_CHARACTERS );
  1125. Q_strncat( commandLine, "-allowdebug ", sizeof( commandLine ), COPY_ALL_CHARACTERS );
  1126. // Add -mpi_SDKMode if it's needed. This would mostly only occur in a debugging situation
  1127. // (someone running out of rel using -mpi_AutoLocalWorker).
  1128. if ( VMPI_IsSDKMode() && !VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_SDKMode ), "" ) )
  1129. {
  1130. Q_strncat( commandLine, VMPI_GetParamString( mpi_SDKMode ), sizeof( commandLine ), COPY_ALL_CHARACTERS );
  1131. }
  1132. }
  1133. if ( i >= argc )
  1134. break;
  1135. Q_snprintf( argStr, sizeof( argStr ), "\"%s\" ", argv[i] );
  1136. Q_strncat( commandLine, argStr, sizeof( commandLine ), COPY_ALL_CHARACTERS );
  1137. }
  1138. char workingDir[1024];
  1139. if ( !_getcwd( workingDir, sizeof( workingDir ) ) )
  1140. {
  1141. Warning( "_getcwd() failed.\n" );
  1142. return false;
  1143. }
  1144. STARTUPINFO si;
  1145. memset( &si, 0, sizeof( si ) );
  1146. si.cb = sizeof( si );
  1147. PROCESS_INFORMATION pi;
  1148. memset( &pi, 0, sizeof( pi ) );
  1149. if ( CreateProcess(
  1150. NULL,
  1151. commandLine,
  1152. NULL, // security
  1153. NULL,
  1154. TRUE,
  1155. (bShowConsoleWindow ? CREATE_NEW_CONSOLE : CREATE_NO_WINDOW) | IDLE_PRIORITY_CLASS, // flags
  1156. NULL, // environment
  1157. workingDir, // current directory (use c:\\ because we don't want it to accidentally share
  1158. // DLLs like vstdlib with us).
  1159. &si,
  1160. &pi ) )
  1161. {
  1162. return true;
  1163. }
  1164. else
  1165. {
  1166. char errStr[1024];
  1167. IP_GetLastErrorString( errStr, sizeof( errStr ) );
  1168. Warning( " - ERROR in CreateProcess (%s)!\n", errStr );
  1169. return false;
  1170. }
  1171. }
  1172. bool InitMaster( int argc, char **argv, const char *pDependencyFilename, VMPIRunMode runMode, bool bPatchMode )
  1173. {
  1174. int nMaxWorkers = -1;
  1175. const char *pProcCount = VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_WorkerCount ) );
  1176. if ( pProcCount )
  1177. {
  1178. nMaxWorkers = atoi( pProcCount );
  1179. Warning( "%s: waiting for %d processes to join.\n", VMPI_GetParamString( mpi_WorkerCount ), nMaxWorkers );
  1180. }
  1181. else
  1182. {
  1183. nMaxWorkers = DEFAULT_MAX_WORKERS;
  1184. }
  1185. nMaxWorkers = clamp( nMaxWorkers, 2, MAX_VMPI_CONNECTIONS );
  1186. g_bMPIMaster = true;
  1187. g_nMaxWorkerCount = nMaxWorkers;
  1188. if ( argc <= 0 )
  1189. Error( "MPI_Init_Master: argc <= 0!" );
  1190. ParseOptions( argc, argv );
  1191. // Send the base filename of the exe we're running. Sometimes if we run vvis followed by vrad
  1192. // really quickly, the old vvis workers can connect to the vrad process and mess with it.
  1193. VMPI_SendExeName();
  1194. // In SDK mode, the master sends the command line to the workers since
  1195. // the workers weren't given a full command line by vmpi_service.
  1196. if ( VMPI_IsSDKMode() )
  1197. {
  1198. VMPI_SendCommandLine( argc, argv );
  1199. }
  1200. if ( !g_MasterBroadcaster.Init( argc, argv, pDependencyFilename, nMaxWorkers, runMode, bPatchMode ) )
  1201. return false;
  1202. bool bRet;
  1203. if ( runMode == VMPI_RUN_LOCAL )
  1204. {
  1205. bRet = SpawnLocalWorker( argc, argv, g_MasterBroadcaster.GetListenPort(), false );
  1206. }
  1207. else
  1208. {
  1209. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_AutoLocalWorker ), "" ) )
  1210. {
  1211. Msg( "%s found. Spawning a local worker automatically.\n", VMPI_GetParamString( mpi_AutoLocalWorker ) );
  1212. SpawnLocalWorker( 1, argv, g_MasterBroadcaster.GetListenPort(), true );
  1213. }
  1214. bRet = true;
  1215. }
  1216. VMPI_HandleTimingWait_Master();
  1217. return bRet;
  1218. }
  1219. void VMPI_InitGlobals( int argc, char **argv, VMPIRunMode runMode )
  1220. {
  1221. g_bUseMPI = true;
  1222. g_VMPIRunMode = runMode;
  1223. // Init event objects.
  1224. g_VMPIMessagesEvent.Init( false, false );
  1225. g_ErrorSocketsEvent.Init( false, false );
  1226. // Load this for GetConsoleWindow().
  1227. g_hKernel32DLL = LoadLibrary( "kernel32.dll" );
  1228. if ( g_hKernel32DLL )
  1229. {
  1230. g_pConsoleWndFn = (GetConsoleWndFn)GetProcAddress( g_hKernel32DLL, "GetConsoleWindow" );
  1231. }
  1232. #if defined( _DEBUG )
  1233. for ( int iArg=0; iArg < argc; iArg++ )
  1234. {
  1235. Warning( "%s\n", argv[iArg] );
  1236. }
  1237. Warning( "\n" );
  1238. #endif
  1239. }
  1240. bool VMPI_CheckForNonSDKExecutables()
  1241. {
  1242. char baseExeFilename[512];
  1243. if ( !GetModuleFileName( GetModuleHandle( NULL ), baseExeFilename, sizeof( baseExeFilename ) ) )
  1244. Error( "VMPI_CheckSDKMode -> GetModuleFileName failed." );
  1245. V_StripLastDir( baseExeFilename, sizeof( baseExeFilename ) );
  1246. V_AppendSlash( baseExeFilename, sizeof( baseExeFilename ) );
  1247. V_strncat( baseExeFilename, "mysql_wrapper.dll", sizeof( baseExeFilename ) );
  1248. // If vmpi_transfer.exe doesn't exist, then we assume we're in SDK mode.
  1249. return ( _access( baseExeFilename, 0 ) == 0 );
  1250. }
  1251. bool IsValidSDKBinPath( CUtlVector< char* > &outStrings, int *pError )
  1252. {
  1253. *pError = 0;
  1254. // Minimum must have drive:/basedir/steamapps/name/sourcesdk/bin/[ep1|orangebox]/bin/exename
  1255. if ( outStrings.Count() < 9 )
  1256. {
  1257. *pError = 0;
  1258. return false;
  1259. }
  1260. if ( V_stricmp( outStrings[outStrings.Count()-2], "bin" ) != 0 )
  1261. {
  1262. *pError = 1;
  1263. return false;
  1264. }
  1265. if ( V_stricmp( outStrings[outStrings.Count()-5], "sourcesdk" ) != 0 )
  1266. {
  1267. *pError = 2;
  1268. return false;
  1269. }
  1270. if ( V_stricmp( outStrings[outStrings.Count()-7], "steamapps" ) != 0 )
  1271. {
  1272. *pError = 3;
  1273. return false;
  1274. }
  1275. // Check the last-access date on clientregistry.blob
  1276. char baseSteamPath[MAX_PATH];
  1277. V_strncpy( baseSteamPath, outStrings[0], sizeof( baseSteamPath) );
  1278. for ( int i=1; i < outStrings.Count() - 7; i++ )
  1279. {
  1280. V_AppendSlash( baseSteamPath, sizeof( baseSteamPath ) );
  1281. V_strncat( baseSteamPath, outStrings[i], sizeof( baseSteamPath ) );
  1282. }
  1283. char blobPath[MAX_PATH];
  1284. V_ComposeFileName( baseSteamPath, "ClientRegistry.blob", blobPath, sizeof( blobPath ) );
  1285. struct _stat results;
  1286. if ( _stat( blobPath, &results ) != 0 )
  1287. {
  1288. *pError = 4;
  1289. return false;
  1290. }
  1291. long curTime;
  1292. VCRHook_Time( &curTime );
  1293. int nSecondsSinceLastSteamAccess = curTime - results.st_mtime;
  1294. int nSecondsPerDay = 60 * 60 * 24;
  1295. int nMaxDaysUnaccessed = 10;
  1296. if ( nSecondsSinceLastSteamAccess > nSecondsPerDay*nMaxDaysUnaccessed )
  1297. {
  1298. *pError = 5; // NOTE: don't change this error code because the outer function checks for it.
  1299. return false;
  1300. }
  1301. // Check for some of the files under sourcesdk_content.
  1302. char sourcesdkContentPath[MAX_PATH];
  1303. V_strncpy( sourcesdkContentPath, outStrings[0], sizeof( sourcesdkContentPath ) );
  1304. for ( int i=1; i < outStrings.Count() - 5; i++ )
  1305. {
  1306. V_AppendSlash( sourcesdkContentPath, sizeof( sourcesdkContentPath ) );
  1307. V_strncat( sourcesdkContentPath, outStrings[i], sizeof( sourcesdkContentPath ) );
  1308. }
  1309. V_AppendSlash( sourcesdkContentPath, sizeof( sourcesdkContentPath ) );
  1310. V_strncat( sourcesdkContentPath, "sourcesdk_content", sizeof( sourcesdkContentPath ) );
  1311. char tempFilename[MAX_PATH], mapsrcFilename[MAX_PATH];
  1312. V_snprintf( tempFilename, sizeof( tempFilename ), "cstrike%cmapsrc", CORRECT_PATH_SEPARATOR );
  1313. V_ComposeFileName( sourcesdkContentPath, tempFilename, mapsrcFilename, sizeof( mapsrcFilename ) );
  1314. if ( _access( mapsrcFilename, 0 ) != 0 )
  1315. {
  1316. *pError = 6;
  1317. return false;
  1318. }
  1319. return true;
  1320. }
  1321. void VerifyValidSDKMode()
  1322. {
  1323. // Make sure we're running out of the SourceSDK directory and that our SDK directories are filled out.
  1324. char baseExeFilename[MAX_PATH];
  1325. if ( !GetModuleFileName( GetModuleHandle( NULL ), baseExeFilename, sizeof( baseExeFilename ) ) )
  1326. Error( "VerifyValidSDKMode: GetModuleFileName failed." );
  1327. V_FixSlashes( baseExeFilename );
  1328. CUtlVector< char* > outStrings;
  1329. char strSlash[2] = {CORRECT_PATH_SEPARATOR, 0};
  1330. V_SplitString( baseExeFilename, strSlash, outStrings );
  1331. int err;
  1332. if ( !IsValidSDKBinPath( outStrings, &err ) )
  1333. {
  1334. outStrings.PurgeAndDeleteElements();
  1335. if ( err == 5 )
  1336. Error( "VMPI running in SDK mode but Steam hasn't been run recently. Please run Steam and retry." );
  1337. else
  1338. Error( "VMPI running in SDK mode but incorrect SDK install detected (error %d).", err );
  1339. }
  1340. }
  1341. void VMPI_CheckSDKMode( int argc, char **argv )
  1342. {
  1343. g_bVMPISDKMode = !VMPI_CheckForNonSDKExecutables();
  1344. g_bVMPISDKModeSet = true;
  1345. // Also check for -mpi_sdkmode (only used in testing).
  1346. if ( !g_bVMPISDKMode )
  1347. {
  1348. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_SDKMode ), "" ) )
  1349. g_bVMPISDKMode = true;
  1350. }
  1351. if ( g_bVMPISDKMode )
  1352. {
  1353. VerifyValidSDKMode();
  1354. }
  1355. if ( g_bVMPISDKMode )
  1356. {
  1357. Msg( "VMPI running in SDK mode.\n" );
  1358. }
  1359. }
  1360. void VMPI_SetupAutoRestartParameters( int argc, char **argv )
  1361. {
  1362. if ( VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_AutoRestart ) ) )
  1363. {
  1364. g_OriginalCommandLineParameters.SetSize( argc );
  1365. for ( int i=0; i < argc; i++ )
  1366. {
  1367. g_OriginalCommandLineParameters[i] = CopyString( argv[i] );
  1368. }
  1369. }
  1370. }
  1371. bool VMPI_HandleAutoRestart()
  1372. {
  1373. if ( g_OriginalCommandLineParameters.Count() == 0 )
  1374. return true;
  1375. Msg( "%s found. Auto-restarting.\n", VMPI_GetParamString( mpi_AutoRestart ) );
  1376. DWORD curPriority = GetPriorityClass( GetCurrentProcess() );
  1377. char commandLine[1024*8];
  1378. commandLine[0] = 0;
  1379. // Add the -mpi_worker argument in, then launch the process.
  1380. for ( int i=0; i < g_OriginalCommandLineParameters.Count(); i++ )
  1381. {
  1382. char argStr[512];
  1383. Q_snprintf( argStr, sizeof( argStr ), "\"%s\" ", g_OriginalCommandLineParameters[i] );
  1384. Q_strncat( commandLine, argStr, sizeof( commandLine ), COPY_ALL_CHARACTERS );
  1385. }
  1386. STARTUPINFO si;
  1387. memset( &si, 0, sizeof( si ) );
  1388. si.cb = sizeof( si );
  1389. PROCESS_INFORMATION pi;
  1390. memset( &pi, 0, sizeof( pi ) );
  1391. if ( CreateProcess(
  1392. NULL,
  1393. commandLine,
  1394. NULL, // security
  1395. NULL,
  1396. TRUE,
  1397. CREATE_NEW_CONSOLE | curPriority, // flags
  1398. NULL, // environment
  1399. NULL,
  1400. &si,
  1401. &pi ) )
  1402. {
  1403. g_OriginalCommandLineParameters.Purge();
  1404. return true;
  1405. }
  1406. else
  1407. {
  1408. char errStr[1024];
  1409. IP_GetLastErrorString( errStr, sizeof( errStr ) );
  1410. Warning( " - ERROR in CreateProcess (%s)!\n", errStr );
  1411. return false;
  1412. }
  1413. }
  1414. bool VMPI_Init(
  1415. int &argc,
  1416. char **&argv,
  1417. const char *pDependencyFilename,
  1418. VMPI_Disconnect_Handler handler,
  1419. VMPIRunMode runMode,
  1420. bool bConnectingAsService
  1421. )
  1422. {
  1423. if ( handler )
  1424. VMPI_AddDisconnectHandler( handler );
  1425. VMPI_SetupAutoRestartParameters( argc, argv );
  1426. VMPI_CheckSDKMode( argc, argv );
  1427. VMPI_InitGlobals( argc, argv, runMode );
  1428. // Were we launched by the vmpi service as a worker?
  1429. const char *pMasterIP = VMPI_FindArg( argc, argv, VMPI_GetParamString( mpi_Worker ), NULL );
  1430. if ( pMasterIP )
  1431. {
  1432. CIPAddr addr;
  1433. addr.port = VMPI_MASTER_FIRST_PORT;
  1434. if ( !ConvertStringToIPAddr( pMasterIP, &addr ) )
  1435. Error( "Unable to parse or resolve master IP (%s).\n", pMasterIP );
  1436. return MPI_Init_Worker( argc, argv, addr, bConnectingAsService );
  1437. }
  1438. else
  1439. {
  1440. if ( !pDependencyFilename )
  1441. {
  1442. Error( "VMPI started as master, but no dependency filename specified.\n" );
  1443. return false;
  1444. }
  1445. return InitMaster( argc, argv, pDependencyFilename, runMode, false );
  1446. }
  1447. }
  1448. void VMPI_Init_PatchMaster( int argc, char **argv )
  1449. {
  1450. const char *pPatchDirectory = VMPI_FindArg( argc, argv, "-mpi_PatchDirectory", NULL );
  1451. if ( !pPatchDirectory )
  1452. Error( "-mpi_PatchDirectory <dir> must be specified if using -PatchHost mode." );
  1453. VMPI_InitGlobals( argc, argv, VMPI_RUN_NETWORKED );
  1454. InitMaster( argc, argv, pPatchDirectory, VMPI_RUN_NETWORKED, true );
  1455. }
  1456. void VMPI_Finalize()
  1457. {
  1458. g_MasterBroadcaster.Term();
  1459. DistributeWork_Cancel();
  1460. // Get rid of all the sockets.
  1461. for ( int iConn=0; iConn < g_nConnections; iConn++ )
  1462. delete g_Connections[iConn];
  1463. g_nConnections = 0;
  1464. // Get rid of all the packets.
  1465. FOR_EACH_LL( g_VMPIMessages, i )
  1466. {
  1467. g_VMPIMessages[i]->Release();
  1468. }
  1469. g_VMPIMessages.Purge();
  1470. g_PersistentPackets.PurgeAndDeleteElements();
  1471. // Get rid of the message buffers
  1472. g_DispatchBuffers.Purge();
  1473. if ( g_hKernel32DLL )
  1474. {
  1475. FreeLibrary( g_hKernel32DLL );
  1476. g_hKernel32DLL = NULL;
  1477. }
  1478. g_WorkerCommandLine.PurgeAndDeleteElements();
  1479. VMPI_HandleAutoRestart();
  1480. }
  1481. VMPIRunMode VMPI_GetRunMode()
  1482. {
  1483. return g_VMPIRunMode;
  1484. }
  1485. VMPIFileSystemMode VMPI_GetFileSystemMode()
  1486. {
  1487. return g_VMPIFileSystemMode;
  1488. }
  1489. int VMPI_GetCurrentNumberOfConnections()
  1490. {
  1491. return g_nConnections;
  1492. }
  1493. void InternalHandleSocketErrors()
  1494. {
  1495. // Copy the list of sockets with errors into a local array so we can handle all the errors outside
  1496. // the mutex, thus avoiding potential deadlock if any error handlers call Error().
  1497. CUtlVector<CVMPIConnection*> errorSockets;
  1498. CCriticalSectionLock csLock( &g_ErrorSocketsCS );
  1499. csLock.Lock();
  1500. errorSockets.SetSize( g_ErrorSockets.Count() );
  1501. int iCur = 0;
  1502. FOR_EACH_LL( g_ErrorSockets, i )
  1503. {
  1504. errorSockets[iCur++] = g_ErrorSockets[i];
  1505. }
  1506. g_ErrorSockets.Purge();
  1507. csLock.Unlock();
  1508. // Handle the errors.
  1509. for ( int i=0; i < errorSockets.Count(); i++ )
  1510. {
  1511. errorSockets[i]->HandleDisconnect();
  1512. }
  1513. UpdateActiveConnectionsText();
  1514. }
  1515. void VMPI_HandleSocketErrors( unsigned long timeout )
  1516. {
  1517. DWORD ret = WaitForSingleObject( g_ErrorSocketsEvent.GetEventHandle(), timeout );
  1518. if ( ret == WAIT_OBJECT_0 )
  1519. {
  1520. InternalHandleSocketErrors();
  1521. }
  1522. }
  1523. // If bWait is false, then this function returns false immediately if there are no messages waiting.
  1524. bool VMPI_GetNextMessage( MessageBuffer *pBuf, int *pSource, unsigned long startTimeout )
  1525. {
  1526. HANDLE handles[2] = { g_ErrorSocketsEvent.GetEventHandle(), g_VMPIMessagesEvent.GetEventHandle() };
  1527. DWORD startTime = Plat_MSTime();
  1528. DWORD timeout = startTimeout;
  1529. while ( 1 )
  1530. {
  1531. DWORD ret = WaitForMultipleObjects( ARRAYSIZE( handles ), handles, FALSE, timeout );
  1532. if ( ret == WAIT_TIMEOUT )
  1533. {
  1534. return false;
  1535. }
  1536. else if ( ret == WAIT_OBJECT_0 )
  1537. {
  1538. // A socket had an error. Handle all socket errors.
  1539. InternalHandleSocketErrors();
  1540. // Update the timeout.
  1541. DWORD delta = Plat_MSTime() - startTime;
  1542. if ( delta >= startTimeout )
  1543. return false;
  1544. timeout = startTimeout - delta;
  1545. continue;
  1546. }
  1547. else if ( ret == (WAIT_OBJECT_0 + 1) )
  1548. {
  1549. // Read out the next message.
  1550. CCriticalSectionLock csLock( &g_VMPIMessagesCS );
  1551. csLock.Lock();
  1552. GrabNextMessage:;
  1553. int iHead = g_VMPIMessages.Head();
  1554. CTCPPacket *pPacket = g_VMPIMessages[iHead];
  1555. g_VMPIMessages.Remove( iHead );
  1556. // Set the event again if there are more messages waiting.
  1557. const char *pBase = pPacket->GetData();
  1558. if ( pPacket->GetLen() >= 6 && (unsigned char)pBase[0] == VMPI_INTERNAL_PACKET_ID && (unsigned char)pBase[1] == VMPI_INTERNAL_SUBPACKET_GROUPED_PACKET )
  1559. {
  1560. // Ok, this is a grouped packet. Split it out into a bunch of separate packets.
  1561. CUtlVector<CTCPPacket*> groupedPackets;
  1562. int iCurOffset = 2;
  1563. while ( (iCurOffset+4) <= pPacket->GetLen() )
  1564. {
  1565. int curPacketLen = *((int*)&pBase[iCurOffset]);
  1566. if ( iCurOffset + curPacketLen > pPacket->GetLen() )
  1567. Error( "Invalid chunked packet\n" );
  1568. iCurOffset += 4;
  1569. CTCPPacket *pChunkPacket = (CTCPPacket*)malloc( sizeof( CTCPPacket ) + curPacketLen - 1 );
  1570. pChunkPacket->m_Len = curPacketLen;
  1571. pChunkPacket->m_UserData = pPacket->m_UserData;
  1572. memcpy( pChunkPacket->m_Data, &pBase[iCurOffset], curPacketLen );
  1573. groupedPackets.AddToTail( pChunkPacket );
  1574. iCurOffset += curPacketLen;
  1575. }
  1576. for ( int i=0; i < groupedPackets.Count(); i++ )
  1577. {
  1578. g_VMPIMessages.AddToHead( groupedPackets[groupedPackets.Count() - i - 1] );
  1579. }
  1580. pPacket->Release();
  1581. goto GrabNextMessage;
  1582. }
  1583. else
  1584. {
  1585. if ( g_VMPIMessages.Count() > 0 )
  1586. g_VMPIMessagesEvent.SetEvent();
  1587. }
  1588. csLock.Unlock();
  1589. // Copy it into their message buffer.
  1590. pBuf->setLen( pPacket->GetLen() );
  1591. memcpy( pBuf->data, pPacket->GetData(), pPacket->GetLen() );
  1592. *pSource = pPacket->GetUserData();
  1593. Assert( *pSource >= 0 && *pSource < g_nConnections );
  1594. // Update global stats about how much data we've received.
  1595. ++g_nMessagesReceived;
  1596. g_nBytesReceived += pPacket->GetLen() + 4; // (4 bytes extra for the packet length)
  1597. // Free the memory associated with the packet.
  1598. pPacket->Release();
  1599. return true;
  1600. }
  1601. else
  1602. {
  1603. Error( "VMPI_GetNextMessage: WaitForSingleObject returned %lu", ret );
  1604. return false;
  1605. }
  1606. }
  1607. }
  1608. bool VMPI_InternalDispatch( MessageBuffer *pBuf, int iSource )
  1609. {
  1610. if ( pBuf->getLen() >= 1 &&
  1611. pBuf->data[0] >= 0 && pBuf->data[0] < MAX_VMPI_PACKET_IDS &&
  1612. g_VMPIDispatch[pBuf->data[0]] )
  1613. {
  1614. return g_VMPIDispatch[ pBuf->data[0] ]( pBuf, iSource, pBuf->data[0] );
  1615. }
  1616. else
  1617. {
  1618. return false;
  1619. }
  1620. }
  1621. bool VMPI_DispatchNextMessage( unsigned long timeout )
  1622. {
  1623. MessageBuffer *pBuf = NULL;
  1624. if ( !g_DispatchBuffers.PopItem( &pBuf ) )
  1625. {
  1626. pBuf = new MessageBuffer();
  1627. }
  1628. bool bRetval = true;
  1629. while ( 1 )
  1630. {
  1631. int iSource;
  1632. if ( VMPI_GetNextMessage( pBuf, &iSource, timeout ) )
  1633. {
  1634. if ( VMPI_InternalDispatch( pBuf, iSource ) )
  1635. {
  1636. break;
  1637. }
  1638. else
  1639. {
  1640. // Workers running in service mode don't hook anything except filesystem stuff, so if they happen to be sent something, no problem.
  1641. if ( !VMPI_IsProcAService( iSource ) )
  1642. {
  1643. // Oops! What is this packet?
  1644. Assert( false );
  1645. }
  1646. }
  1647. }
  1648. else
  1649. {
  1650. bRetval = false;
  1651. break;
  1652. }
  1653. }
  1654. g_DispatchBuffers.PushItem( pBuf );
  1655. return bRetval;
  1656. }
  1657. bool VMPI_DispatchUntil( MessageBuffer *pBuf, int *pSource, int packetID, int subPacketID, bool bWait )
  1658. {
  1659. while ( 1 )
  1660. {
  1661. if ( !VMPI_GetNextMessage( pBuf, pSource, bWait ? VMPI_TIMEOUT_INFINITE : 0 ) )
  1662. return false;
  1663. if ( !VMPI_InternalDispatch( pBuf, *pSource ) )
  1664. {
  1665. if ( pBuf->getLen() >= 1 && (unsigned char)pBuf->data[0] == packetID )
  1666. {
  1667. if ( subPacketID == -1 )
  1668. return true;
  1669. if ( pBuf->getLen() >= 2 && (unsigned char)pBuf->data[1] == subPacketID )
  1670. return true;
  1671. }
  1672. // Oops! What is this packet?
  1673. // Note: the most common case where this happens is if it finishes a BuildFaceLights run
  1674. // and is in an AppBarrier and one of the workers is still finishing up some work given to it.
  1675. // It'll be waiting for a barrier packet, and it'll get results. In that case, the packet should
  1676. // be discarded like we do here, so maybe this assert won't be necessary.
  1677. //Assert( false );
  1678. }
  1679. }
  1680. }
  1681. bool VMPI_SendData( void *pData, int nBytes, int iDest, int fVMPISendFlags )
  1682. {
  1683. return VMPI_SendChunks( &pData, &nBytes, 1, iDest, fVMPISendFlags );
  1684. }
  1685. inline bool VMPI_FilterPacketsForServiceDownloader( CVMPIConnection *pConnection, void const * const *pChunks, const int *pChunkLengths, int nChunks )
  1686. {
  1687. if ( pConnection->m_bIsAService )
  1688. {
  1689. // Find the first byte and treat that as the packet ID.
  1690. for ( int i=0; i < nChunks; i++ )
  1691. {
  1692. if ( pChunkLengths[i] > 0 )
  1693. {
  1694. unsigned char cPacketID = *((unsigned char*)pChunks[i]);
  1695. if ( cPacketID == VMPI_INTERNAL_PACKET_ID || cPacketID == VMPI_SHARED_PACKET_ID || cPacketID == VMPI_PACKETID_FILESYSTEM )
  1696. return false;
  1697. else
  1698. return true;
  1699. }
  1700. }
  1701. }
  1702. return false;
  1703. }
  1704. void VMPI_GroupPackets( CVMPIConnection *pConn, void const * const *pChunks, const int *pChunkLengths, int nChunks )
  1705. {
  1706. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  1707. connectionsLock.Lock();
  1708. // First add the header.
  1709. if ( pConn->m_GroupedChunks.Count() == 0 )
  1710. {
  1711. pConn->m_GroupedChunks.AddToTail( g_GroupedPacketHeader );
  1712. pConn->m_GroupedChunkLengths.AddToTail( sizeof( g_GroupedPacketHeader ) );
  1713. }
  1714. // Collate the chunks.
  1715. int nTotalLength = 0;
  1716. for ( int i=0; i < nChunks; i++ )
  1717. nTotalLength += pChunkLengths[i];
  1718. char *pOut = new char[nTotalLength + 4];
  1719. *((int*)pOut) = nTotalLength;
  1720. int iOutByte = 4;
  1721. for ( int i=0; i < nChunks; i++ )
  1722. {
  1723. memcpy( &pOut[iOutByte], pChunks[i], pChunkLengths[i] );
  1724. iOutByte += pChunkLengths[i];
  1725. }
  1726. pConn->m_GroupedChunks.AddToTail( pOut );
  1727. pConn->m_GroupedChunkLengths.AddToTail( nTotalLength + 4 );
  1728. }
  1729. void VMPI_FlushGroupedPackets( unsigned long msInterval )
  1730. {
  1731. if ( msInterval != 0 )
  1732. {
  1733. unsigned long curTime = Plat_MSTime();
  1734. if ( curTime - g_LastFlushGroupedPacketsTime < msInterval )
  1735. return;
  1736. g_LastFlushGroupedPacketsTime = curTime;
  1737. }
  1738. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  1739. connectionsLock.Lock();
  1740. for ( int i=0; i < g_nConnections; i++ )
  1741. {
  1742. CVMPIConnection *pConn = g_Connections[i];
  1743. if ( !pConn )
  1744. continue;
  1745. IThreadedTCPSocket *pSocket = pConn->GetSocket();
  1746. if ( !pSocket || pConn->m_GroupedChunks.Count() == 0 )
  1747. continue;
  1748. pSocket->SendChunks( pConn->m_GroupedChunks.Base(), pConn->m_GroupedChunkLengths.Base(), pConn->m_GroupedChunks.Count() );
  1749. // Free the chunks.
  1750. for ( int i=1; i < pConn->m_GroupedChunks.Count(); i++ )
  1751. {
  1752. free( pConn->m_GroupedChunks[i] );
  1753. }
  1754. pConn->m_GroupedChunks.RemoveAll();
  1755. pConn->m_GroupedChunkLengths.RemoveAll();
  1756. }
  1757. }
  1758. bool VMPI_SendChunks( void const * const *pChunks, const int *pChunkLengths, int nChunks, int iDest, int fVMPISendFlags )
  1759. {
  1760. if ( iDest == VMPI_SEND_TO_ALL )
  1761. {
  1762. // Don't want new connections while in here!
  1763. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  1764. connectionsLock.Lock();
  1765. for ( int i=0; i < g_nConnections; i++ )
  1766. VMPI_SendChunks( pChunks, pChunkLengths, nChunks, i );
  1767. return true;
  1768. }
  1769. else if ( iDest == VMPI_PERSISTENT )
  1770. {
  1771. // Don't want new connections while in here!
  1772. CCriticalSectionLock connectionsLock( &g_ConnectionsCS );
  1773. connectionsLock.Lock();
  1774. CCriticalSectionLock csLock( &g_PersistentPacketsCS );
  1775. csLock.Lock();
  1776. // Send the packet to everyone.
  1777. for ( int i=0; i < g_nConnections; i++ )
  1778. VMPI_SendChunks( pChunks, pChunkLengths, nChunks, i );
  1779. // Remember to send it to the new workers.
  1780. if ( iDest == VMPI_PERSISTENT )
  1781. {
  1782. PersistentPacket *pNew = new PersistentPacket;
  1783. for ( int i=0; i < nChunks; i++ )
  1784. pNew->AddMultipleToTail( pChunkLengths[i], (const char*)pChunks[i] );
  1785. g_PersistentPackets.AddToTail( pNew );
  1786. }
  1787. return true;
  1788. }
  1789. else
  1790. {
  1791. g_nMessagesSent++;
  1792. g_nBytesSent += 4; // for message tag.
  1793. for ( int i=0; i < nChunks; i++ )
  1794. g_nBytesSent += pChunkLengths[i];
  1795. CVMPIConnection *pConnection = g_Connections[iDest];
  1796. if ( pConnection )
  1797. {
  1798. // If it's a service downloader, only send certain packet IDs.
  1799. if ( VMPI_FilterPacketsForServiceDownloader( pConnection, pChunks, pChunkLengths, nChunks ) )
  1800. return true;
  1801. IThreadedTCPSocket *pSocket = pConnection->GetSocket();
  1802. if ( !pSocket )
  1803. return false;
  1804. if ( g_bGroupPackets && (fVMPISendFlags & k_eVMPISendFlags_GroupPackets) )
  1805. {
  1806. VMPI_GroupPackets( pConnection, pChunks, pChunkLengths, nChunks );
  1807. return true;
  1808. }
  1809. else
  1810. {
  1811. return pSocket->SendChunks( pChunks, pChunkLengths, nChunks );
  1812. }
  1813. }
  1814. else
  1815. {
  1816. return false;
  1817. }
  1818. }
  1819. }
  1820. bool VMPI_Send2Chunks( const void *pChunk1, int chunk1Len, const void *pChunk2, int chunk2Len, int iDest, int fVMPISendFlags )
  1821. {
  1822. const void *pChunks[2] = { pChunk1, pChunk2 };
  1823. int len[2] = { chunk1Len, chunk2Len };
  1824. return VMPI_SendChunks( pChunks, len, ARRAYSIZE( pChunks ), iDest, fVMPISendFlags );
  1825. }
  1826. bool VMPI_Send3Chunks( const void *pChunk1, int chunk1Len, const void *pChunk2, int chunk2Len, const void *pChunk3, int chunk3Len, int iDest, int fVMPISendFlags )
  1827. {
  1828. const void *pChunks[3] = { pChunk1, pChunk2, pChunk3 };
  1829. int len[3] = { chunk1Len, chunk2Len, chunk3Len };
  1830. return VMPI_SendChunks( pChunks, len, ARRAYSIZE( pChunks ), iDest, fVMPISendFlags );
  1831. }
  1832. void VMPI_AddDisconnectHandler( VMPI_Disconnect_Handler handler )
  1833. {
  1834. g_DisconnectHandlers.AddToTail( handler );
  1835. }
  1836. CVMPIConnection* GetConnection( int procID )
  1837. {
  1838. Assert( procID >= 0 && procID < g_nConnections );
  1839. return g_Connections[procID];
  1840. }
  1841. bool VMPI_IsProcConnected( int procID )
  1842. {
  1843. if ( procID < 0 || procID >= g_nConnections )
  1844. {
  1845. Assert( false );
  1846. return false;
  1847. }
  1848. return g_Connections[procID]->GetSocket() != NULL;
  1849. }
  1850. bool VMPI_IsProcAService( int procID )
  1851. {
  1852. if ( procID < 0 || procID >= g_nConnections )
  1853. {
  1854. Assert( false );
  1855. return false;
  1856. }
  1857. return g_Connections[procID]->m_bIsAService;
  1858. }
  1859. void VMPI_Sleep( unsigned long ms )
  1860. {
  1861. Sleep( ms );
  1862. }
  1863. const char* VMPI_GetMachineName( int iProc )
  1864. {
  1865. if ( g_bMPIMaster && iProc == VMPI_MASTER_ID )
  1866. return VMPI_GetLocalMachineName();
  1867. if ( iProc < 0 || iProc >= g_nConnections )
  1868. {
  1869. Assert( false );
  1870. return "invalid index";
  1871. }
  1872. return g_Connections[iProc]->GetMachineName();
  1873. }
  1874. void VMPI_SetMachineName( int iProc, const char *pName )
  1875. {
  1876. if ( iProc < 0 || iProc >= g_nConnections )
  1877. {
  1878. Assert( false );
  1879. return;
  1880. }
  1881. g_Connections[iProc]->SetMachineName( pName );
  1882. }
  1883. bool VMPI_HasMachineNameBeenSet( int iProc )
  1884. {
  1885. if ( iProc < 0 || iProc >= g_nConnections )
  1886. {
  1887. Assert( false );
  1888. return false;
  1889. }
  1890. return g_Connections[iProc]->HasMachineNameBeenSet();
  1891. }
  1892. const char* VMPI_GetLocalMachineName()
  1893. {
  1894. static char cName[MAX_COMPUTERNAME_LENGTH+1];
  1895. DWORD len = sizeof( cName );
  1896. if ( GetComputerName( cName, &len ) )
  1897. return cName;
  1898. else
  1899. return "(error in GetComputerName)";
  1900. }
  1901. unsigned long VMPI_GetJobWorkerID( int iProc )
  1902. {
  1903. return GetConnection( iProc )->m_JobWorkerID;
  1904. }
  1905. void VMPI_SetJobWorkerID( int iProc, unsigned long jobWorkerID )
  1906. {
  1907. GetConnection( iProc )->m_JobWorkerID = jobWorkerID;
  1908. }
  1909. void VMPI_GetCurrentStage( char *pOut, int strLen )
  1910. {
  1911. CCriticalSectionLock csLock( &g_CurrentStageCS );
  1912. csLock.Lock();
  1913. Q_strncpy( pOut, g_CurrentStageString, strLen );
  1914. }
  1915. void VMPI_SetCurrentStage( const char *pCurStage )
  1916. {
  1917. CCriticalSectionLock csLock( &g_CurrentStageCS );
  1918. csLock.Lock();
  1919. Q_strncpy( g_CurrentStageString, pCurStage, sizeof( g_CurrentStageString ) );
  1920. }
  1921. void VMPI_InviteDebugWorkers()
  1922. {
  1923. // Only allow workers with password set to debugworker.
  1924. g_MasterBroadcaster.SetPassword( "debugworker" );
  1925. // Disable timeouts so they can sit in the debugger.
  1926. g_MasterBroadcaster.SetNoTimeoutOption();
  1927. ThreadedTCP_EnableTimeouts( false );
  1928. // Let in some more workers.
  1929. g_MasterBroadcaster.IncreaseMaxWorkers( 25 );
  1930. }
  1931. bool VMPI_IsSDKMode()
  1932. {
  1933. if ( g_bVMPISDKModeSet )
  1934. return g_bVMPISDKMode;
  1935. else
  1936. return !VMPI_CheckForNonSDKExecutables();
  1937. }
  1938. const char* VMPI_GetParamString( EVMPICmdLineParam eParam )
  1939. {
  1940. if ( eParam <= k_eVMPICmdLineParam_FirstParam || eParam >= k_eVMPICmdLineParam_LastParam )
  1941. {
  1942. Assert( false );
  1943. Warning( "Invalid call: VMPI_GetParamString( %d )\n", eParam );
  1944. return "unknown";
  1945. }
  1946. else
  1947. {
  1948. return g_VMPIParams[eParam].m_pName;
  1949. }
  1950. }
  1951. int VMPI_GetParamFlags( EVMPICmdLineParam eParam )
  1952. {
  1953. if ( eParam <= k_eVMPICmdLineParam_FirstParam || eParam >= k_eVMPICmdLineParam_LastParam )
  1954. {
  1955. Assert( false );
  1956. Warning( "Invalid call: VMPI_GetParamString( %d )\n", eParam );
  1957. return 0;
  1958. }
  1959. else
  1960. {
  1961. return g_VMPIParams[eParam].m_ParamFlags;
  1962. }
  1963. }
  1964. bool VMPI_IsParamUsed( EVMPICmdLineParam eParam )
  1965. {
  1966. int iParam = CommandLine()->FindParm( VMPI_GetParamString( eParam ) );
  1967. return iParam != 0;
  1968. }
  1969. const char* VMPI_GetParamHelpString( EVMPICmdLineParam eParam )
  1970. {
  1971. if ( eParam <= k_eVMPICmdLineParam_FirstParam || eParam >= k_eVMPICmdLineParam_LastParam )
  1972. {
  1973. Assert( false );
  1974. Warning( "Invalid call: VMPI_GetParamHelpString( %d )\n", eParam );
  1975. return "unknown vmpi param";
  1976. }
  1977. else
  1978. {
  1979. return g_VMPIParams[eParam].m_pHelpText;
  1980. }
  1981. }