9
9
using System . IO ;
10
10
using System . Linq ;
11
11
using System . Reactive . Linq ;
12
+ using System . Text ;
12
13
using System . Threading . Tasks ;
13
14
using System . Threading . Tasks . Dataflow ;
14
15
using Microsoft . Azure . WebJobs . Script . Description ;
15
16
using Microsoft . Azure . WebJobs . Script . Diagnostics ;
16
17
using Microsoft . Azure . WebJobs . Script . Eventing ;
17
18
using Microsoft . Azure . WebJobs . Script . Grpc . Eventing ;
19
+ using Microsoft . Azure . WebJobs . Script . Grpc . Extensions ;
18
20
using Microsoft . Azure . WebJobs . Script . Grpc . Messages ;
19
21
using Microsoft . Azure . WebJobs . Script . ManagedDependencies ;
20
22
using Microsoft . Azure . WebJobs . Script . Workers ;
21
23
using Microsoft . Azure . WebJobs . Script . Workers . Rpc ;
24
+ using Microsoft . Azure . WebJobs . Script . Workers . SharedMemoryDataTransfer ;
25
+ using Microsoft . CodeAnalysis . VisualBasic . Syntax ;
22
26
using Microsoft . Extensions . Logging ;
23
27
using Microsoft . Extensions . Options ;
24
28
using static Microsoft . Azure . WebJobs . Script . Grpc . Messages . RpcLog . Types ;
25
29
using FunctionMetadata = Microsoft . Azure . WebJobs . Script . Description . FunctionMetadata ;
26
30
using MsgType = Microsoft . Azure . WebJobs . Script . Grpc . Messages . StreamingMessage . ContentOneofCase ;
31
+ using ParameterBindingType = Microsoft . Azure . WebJobs . Script . Grpc . Messages . ParameterBinding . RpcDataOneofCase ;
27
32
28
33
namespace Microsoft . Azure . WebJobs . Script . Grpc
29
34
{
@@ -35,6 +40,7 @@ internal class GrpcWorkerChannel : IRpcWorkerChannel, IDisposable
35
40
private readonly string _runtime ;
36
41
private readonly IEnvironment _environment ;
37
42
private readonly IOptionsMonitor < ScriptApplicationHostOptions > _applicationHostOptions ;
43
+ private readonly ISharedMemoryManager _sharedMemoryManager ;
38
44
39
45
private IDisposable _functionLoadRequestResponseEvent ;
40
46
private bool _disposed ;
@@ -60,6 +66,7 @@ internal class GrpcWorkerChannel : IRpcWorkerChannel, IDisposable
60
66
private TaskCompletionSource < bool > _reloadTask = new TaskCompletionSource < bool > ( TaskCreationOptions . RunContinuationsAsynchronously ) ;
61
67
private TaskCompletionSource < bool > _workerInitTask = new TaskCompletionSource < bool > ( TaskCreationOptions . RunContinuationsAsynchronously ) ;
62
68
private TimeSpan _functionLoadTimeout = TimeSpan . FromMinutes ( 10 ) ;
69
+ private bool _isSharedMemoryDataTransferEnabled ;
63
70
64
71
internal GrpcWorkerChannel (
65
72
string workerId ,
@@ -70,7 +77,8 @@ internal GrpcWorkerChannel(
70
77
IMetricsLogger metricsLogger ,
71
78
int attemptCount ,
72
79
IEnvironment environment ,
73
- IOptionsMonitor < ScriptApplicationHostOptions > applicationHostOptions )
80
+ IOptionsMonitor < ScriptApplicationHostOptions > applicationHostOptions ,
81
+ ISharedMemoryManager sharedMemoryManager )
74
82
{
75
83
_workerId = workerId ;
76
84
_eventManager = eventManager ;
@@ -81,6 +89,7 @@ internal GrpcWorkerChannel(
81
89
_metricsLogger = metricsLogger ;
82
90
_environment = environment ;
83
91
_applicationHostOptions = applicationHostOptions ;
92
+ _sharedMemoryManager = sharedMemoryManager ;
84
93
85
94
_workerCapabilities = new GrpcCapabilities ( _workerChannelLogger ) ;
86
95
@@ -101,7 +110,7 @@ internal GrpcWorkerChannel(
101
110
. Subscribe ( msg => _eventManager . Publish ( new HostRestartEvent ( ) ) ) ) ;
102
111
103
112
_eventSubscriptions . Add ( _inboundWorkerEvents . Where ( msg => msg . MessageType == MsgType . InvocationResponse )
104
- . Subscribe ( ( msg ) => InvokeResponse ( msg . Message . InvocationResponse ) ) ) ;
113
+ . Subscribe ( async ( msg ) => await InvokeResponse ( msg . Message . InvocationResponse ) ) ) ;
105
114
106
115
_inboundWorkerEvents . Where ( msg => msg . MessageType == MsgType . WorkerStatusResponse )
107
116
. Subscribe ( ( msg ) => ReceiveWorkerStatusResponse ( msg . Message . RequestId , msg . Message . WorkerStatusResponse ) ) ;
@@ -239,6 +248,7 @@ internal void WorkerInitResponse(GrpcEvent initEvent)
239
248
}
240
249
_state = _state | RpcWorkerChannelState . Initialized ;
241
250
_workerCapabilities . UpdateCapabilities ( _initMessage . Capabilities ) ;
251
+ _isSharedMemoryDataTransferEnabled = IsSharedMemoryDataTransferEnabled ( ) ;
242
252
_workerInitTask . SetResult ( true ) ;
243
253
}
244
254
@@ -406,7 +416,7 @@ internal async Task SendInvocationRequest(ScriptInvocationContext context)
406
416
context . ResultSource . SetCanceled ( ) ;
407
417
return ;
408
418
}
409
- var invocationRequest = await context . ToRpcInvocationRequest ( _workerChannelLogger , _workerCapabilities ) ;
419
+ var invocationRequest = await context . ToRpcInvocationRequest ( _workerChannelLogger , _workerCapabilities , _isSharedMemoryDataTransferEnabled , _sharedMemoryManager ) ;
410
420
_executingInvocations . TryAdd ( invocationRequest . InvocationId , context ) ;
411
421
412
422
SendStreamingMessage ( new StreamingMessage
@@ -421,7 +431,41 @@ internal async Task SendInvocationRequest(ScriptInvocationContext context)
421
431
}
422
432
}
423
433
424
- internal void InvokeResponse ( InvocationResponse invokeResponse )
434
+ private async Task < object > GetBindingDataAsync ( ParameterBinding binding , string invocationId )
435
+ {
436
+ switch ( binding . RpcDataCase )
437
+ {
438
+ case ParameterBindingType . RpcSharedMemory :
439
+ // Data was transferred by the worker using shared memory
440
+ return await binding . RpcSharedMemory . ToObjectAsync ( _workerChannelLogger , invocationId , _sharedMemoryManager ) ;
441
+ case ParameterBindingType . Data :
442
+ // Data was transferred by the worker using RPC
443
+ return binding . Data . ToObject ( ) ;
444
+ default :
445
+ throw new InvalidOperationException ( "Unknown ParameterBindingType" ) ;
446
+ }
447
+ }
448
+
449
+ /// <summary>
450
+ /// From the output data produced by the worker, get a list of the shared memory maps that were created for this invocation.
451
+ /// </summary>
452
+ /// <param name="bindings">List of <see cref="ParameterBinding"/> produced by the worker as output.</param>
453
+ /// <returns>List of names of shared memory maps produced by the worker.</returns>
454
+ private IList < string > GetOutputMaps ( IList < ParameterBinding > bindings )
455
+ {
456
+ IList < string > outputMaps = new List < string > ( ) ;
457
+ foreach ( ParameterBinding binding in bindings )
458
+ {
459
+ if ( binding . RpcSharedMemory != null )
460
+ {
461
+ outputMaps . Add ( binding . RpcSharedMemory . Name ) ;
462
+ }
463
+ }
464
+
465
+ return outputMaps ;
466
+ }
467
+
468
+ internal async Task InvokeResponse ( InvocationResponse invokeResponse )
425
469
{
426
470
_workerChannelLogger . LogDebug ( "InvocationResponse received for invocation id: {Id}" , invokeResponse . InvocationId ) ;
427
471
@@ -430,8 +474,29 @@ internal void InvokeResponse(InvocationResponse invokeResponse)
430
474
{
431
475
try
432
476
{
433
- IDictionary < string , object > bindingsDictionary = invokeResponse . OutputData
434
- . ToDictionary ( binding => binding . Name , binding => binding . Data . ToObject ( ) ) ;
477
+ StringBuilder logBuilder = new StringBuilder ( ) ;
478
+ bool usedSharedMemory = false ;
479
+
480
+ foreach ( ParameterBinding binding in invokeResponse . OutputData )
481
+ {
482
+ switch ( binding . RpcDataCase )
483
+ {
484
+ case ParameterBindingType . RpcSharedMemory :
485
+ logBuilder . AppendFormat ( "{0}:{1}," , binding . Name , binding . RpcSharedMemory . Count ) ;
486
+ usedSharedMemory = true ;
487
+ break ;
488
+ default :
489
+ break ;
490
+ }
491
+ }
492
+
493
+ if ( usedSharedMemory )
494
+ {
495
+ _workerChannelLogger . LogDebug ( "Shared memory usage for response of invocation Id: {Id} is {SharedMemoryUsage}" , invokeResponse . InvocationId , logBuilder . ToString ( ) ) ;
496
+ }
497
+
498
+ IDictionary < string , object > bindingsDictionary = await invokeResponse . OutputData
499
+ . ToDictionaryAsync ( binding => binding . Name , binding => GetBindingDataAsync ( binding , invokeResponse . InvocationId ) ) ;
435
500
436
501
var result = new ScriptInvocationResult ( )
437
502
{
@@ -444,9 +509,40 @@ internal void InvokeResponse(InvocationResponse invokeResponse)
444
509
{
445
510
context . ResultSource . TrySetException ( responseEx ) ;
446
511
}
512
+ finally
513
+ {
514
+ // Free memory allocated by the host (for input bindings)
515
+ if ( ! _sharedMemoryManager . TryFreeSharedMemoryMapsForInvocation ( invokeResponse . InvocationId ) )
516
+ {
517
+ _workerChannelLogger . LogWarning ( $ "Cannot free all shared memory resources for invocation: { invokeResponse . InvocationId } ") ;
518
+ }
519
+
520
+ // List of shared memory maps that were produced by the worker (for output bindings)
521
+ IList < string > outputMaps = GetOutputMaps ( invokeResponse . OutputData ) ;
522
+ if ( outputMaps . Count > 0 )
523
+ {
524
+ // If this invocation was using any shared memory maps produced by the worker, close them to free memory
525
+ SendCloseSharedMemoryResourcesForInvocationRequest ( outputMaps ) ;
526
+ }
527
+ }
447
528
}
448
529
}
449
530
531
+ /// <summary>
532
+ /// Request to free memory allocated by the worker (for output bindings)
533
+ /// </summary>
534
+ /// <param name="outputMaps">List of names of shared memory maps to close from the worker.</param>
535
+ internal void SendCloseSharedMemoryResourcesForInvocationRequest ( IList < string > outputMaps )
536
+ {
537
+ CloseSharedMemoryResourcesRequest closeSharedMemoryResourcesRequest = new CloseSharedMemoryResourcesRequest ( ) ;
538
+ closeSharedMemoryResourcesRequest . MapNames . AddRange ( outputMaps ) ;
539
+
540
+ SendStreamingMessage ( new StreamingMessage ( )
541
+ {
542
+ CloseSharedMemoryResourcesRequest = closeSharedMemoryResourcesRequest
543
+ } ) ;
544
+ }
545
+
450
546
internal void Log ( GrpcEvent msg )
451
547
{
452
548
var rpcLog = msg . Message . RpcLog ;
@@ -616,5 +712,44 @@ public bool TryFailExecutions(Exception workerException)
616
712
}
617
713
return true ;
618
714
}
715
+
716
+ /// <summary>
717
+ /// Determine if shared memory transfer is enabled.
718
+ /// The following conditions must be met:
719
+ /// 1) <see cref="RpcWorkerConstants.FunctionsWorkerSharedMemoryDataTransferEnabledSettingName"/> must be set in environment variable (AppSetting).
720
+ /// 2) Worker must have the capability <see cref="RpcWorkerConstants.SharedMemoryDataTransfer"/>.
721
+ /// </summary>
722
+ /// <returns><see cref="true"/> if shared memory data transfer is enabled, <see cref="false"/> otherwise.</returns>
723
+ internal bool IsSharedMemoryDataTransferEnabled ( )
724
+ {
725
+ // Check if the environment variable (AppSetting) has this feature enabled
726
+ string envVal = _environment . GetEnvironmentVariable ( RpcWorkerConstants . FunctionsWorkerSharedMemoryDataTransferEnabledSettingName ) ;
727
+ if ( string . IsNullOrEmpty ( envVal ) )
728
+ {
729
+ return false ;
730
+ }
731
+
732
+ bool envValEnabled = false ;
733
+ if ( bool . TryParse ( envVal , out bool boolResult ) )
734
+ {
735
+ // Check if value was specified as a bool (true/false)
736
+ envValEnabled = boolResult ;
737
+ }
738
+ else if ( int . TryParse ( envVal , out int intResult ) && intResult == 1 )
739
+ {
740
+ // Check if value was specified as an int (1/0)
741
+ envValEnabled = true ;
742
+ }
743
+
744
+ if ( ! envValEnabled )
745
+ {
746
+ return false ;
747
+ }
748
+
749
+ // Check if the worker supports this feature
750
+ bool capabilityEnabled = ! string . IsNullOrEmpty ( _workerCapabilities . GetCapabilityState ( RpcWorkerConstants . SharedMemoryDataTransfer ) ) ;
751
+ _workerChannelLogger . LogDebug ( "IsSharedMemoryDataTransferEnabled: {SharedMemoryDataTransferEnabled}" , capabilityEnabled ) ;
752
+ return capabilityEnabled ;
753
+ }
619
754
}
620
755
}
0 commit comments