Skip to content

Commit

Permalink
Atomically get Puppeteer utilities (#2236)
Browse files Browse the repository at this point in the history 
* Atomically get Puppeteer utilities

* Inject util in context

* cr
  • Loading branch information
@kblok
kblok committed Jun 14, 2023
1 parent a5ae465 commit 7cde7c4
Show file tree
Hide file tree
Showing 6 changed files with 45 additions and 55 deletions.
4 changes: 2 additions & 2 deletions lib/PuppeteerSharp.Tests/InjectedTests/InjectedTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ public async Task ShouldWork()
PuppeteerUtil => {
return typeof PuppeteerUtil === 'object';
}",
await world.GetPuppeteerUtilAsync());
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)));
Assert.True(result);
}

Expand All @@ -37,7 +37,7 @@ public async Task CreateFunctionShouldWork()
.PuppeteerWorld.EvaluateFunctionAsync<int>(@"({createFunction}, fnString) => {
return createFunction(fnString)(4);
}",
await world.GetPuppeteerUtilAsync(),
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
"() => 4");
Assert.Equal(4, result);
}
Expand Down
10 changes: 4 additions & 6 deletions lib/PuppeteerSharp/CustomQueriesManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ namespace PuppeteerSharp
{
internal class CustomQueriesManager
{
private static readonly string[] CustomQuerySeparators = new[] { "=", "/" };
private static readonly string[] _customQuerySeparators = new[] { "=", "/" };
private readonly Dictionary<string, PuppeteerQueryHandler> _internalQueryHandlers;
private readonly Dictionary<string, PuppeteerQueryHandler> _queryHandlers = new();
private readonly PuppeteerQueryHandler _pierceHandler = CreatePuppeteerQueryHandler(new CustomQueryHandler
Expand Down Expand Up @@ -194,7 +194,7 @@ internal void RegisterCustomQueryHandler(string name, CustomQueryHandler queryHa

foreach (var kv in handlers)
{
foreach (var separator in CustomQuerySeparators)
foreach (var separator in _customQuerySeparators)
{
var prefix = $"{kv.Key}{separator}";

Expand Down Expand Up @@ -231,11 +231,10 @@ private static PuppeteerQueryHandler CreatePuppeteerQueryHandler(CustomQueryHand
{
internalHandler.QueryOne = async (IElementHandle element, string selector) =>
{
var handle = element as JSHandle;
var jsHandle = await element.EvaluateFunctionHandleAsync(
handler.QueryOne,
selector,
await handle.ExecutionContext.World.GetPuppeteerUtilAsync().ConfigureAwait(false))
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)))
.ConfigureAwait(false);
if (jsHandle is ElementHandle elementHandle)
{
Expand Down Expand Up @@ -280,11 +279,10 @@ await handle.ExecutionContext.World.GetPuppeteerUtilAsync().ConfigureAwait(false
{
internalHandler.QueryAll = async (IElementHandle element, string selector) =>
{
var handle = element as JSHandle;
var jsHandle = await element.EvaluateFunctionHandleAsync(
handler.QueryAll,
selector,
await handle.ExecutionContext.World.GetPuppeteerUtilAsync().ConfigureAwait(false))
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)))
.ConfigureAwait(false);
var properties = await jsHandle.GetPropertiesAsync().ConfigureAwait(false);
var result = new List<ElementHandle>();
Expand Down
31 changes: 31 additions & 0 deletions lib/PuppeteerSharp/ExecutionContext.cs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Numerics;
using System.Reflection;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
Expand All @@ -16,8 +18,10 @@ public class ExecutionContext : IExecutionContext
internal const string EvaluationScriptUrl = "__puppeteer_evaluation_script__";

private static readonly Regex _sourceUrlRegex = new(@"^[\040\t]*\/\/[@#] sourceURL=\s*\S*?\s*$", RegexOptions.Multiline);
private static string _injectedSource;

private readonly string _evaluationScriptSuffix = $"//# sourceURL={EvaluationScriptUrl}";
private IJSHandle _puppeteerUtil;

internal ExecutionContext(
CDPSession client,
Expand Down Expand Up @@ -89,6 +93,17 @@ public async Task<IJSHandle> QueryObjectsAsync(IJSHandle prototypeHandle)
return CreateJSHandle(response.Objects);
}

internal async Task<IJSHandle> GetPuppeteerUtilAsync()
{
if (_puppeteerUtil == null)
{
var injectedSource = GetInjectedSource();
_puppeteerUtil = await EvaluateExpressionHandleAsync(injectedSource).ConfigureAwait(false);
}

return _puppeteerUtil;
}

internal IJSHandle CreateJSHandle(RemoteObject remoteObject)
=> remoteObject.Subtype == RemoteObjectSubtype.Node && Frame != null
? new ElementHandle(this, Client, remoteObject, Frame, ((Frame)Frame).FrameManager.Page, ((Frame)Frame).FrameManager)
Expand Down Expand Up @@ -120,6 +135,22 @@ internal async Task<IElementHandle> AdoptElementHandleAsync(IElementHandle eleme
return CreateJSHandle(obj.Object) as ElementHandle;
}

private static string GetInjectedSource()
{
if (string.IsNullOrEmpty(_injectedSource))
{
var assembly = Assembly.GetExecutingAssembly();
var resourceName = "PuppeteerSharp.Injected.injected.js";

using var stream = assembly.GetManifestResourceStream(resourceName);
using var reader = new StreamReader(stream);
var fileContent = reader.ReadToEnd();
_injectedSource = fileContent;
}

return _injectedSource;
}

private static string GetExceptionMessage(EvaluateExceptionResponseDetails exceptionDetails)
{
if (exceptionDetails.Exception != null)
Expand Down
4 changes: 2 additions & 2 deletions lib/PuppeteerSharp/Frame.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,7 @@ public async Task<IElementHandle> AddStyleTagAsync(AddTagOptions options)
await promise;
return element;
}",
await PuppeteerWorld.GetPuppeteerUtilAsync().ConfigureAwait(false),
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
options.Url,
options.Id,
options.Type,
Expand Down Expand Up @@ -283,7 +283,7 @@ public async Task<IElementHandle> AddScriptTagAsync(AddTagOptions options)
await promise;
return script;
}",
await PuppeteerWorld.GetPuppeteerUtilAsync().ConfigureAwait(false),
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
options.Url,
options.Id,
options.Type,
Expand Down
41 changes: 1 addition & 40 deletions lib/PuppeteerSharp/IsolatedWorld.cs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using Newtonsoft.Json.Linq;
Expand All @@ -21,7 +19,6 @@ namespace PuppeteerSharp

internal class IsolatedWorld
{
private static string _injectedSource;
private readonly FrameManager _frameManager;
private readonly TimeoutSettings _timeoutSettings;
private readonly CDPSession _client;
Expand Down Expand Up @@ -56,10 +53,6 @@ internal class IsolatedWorld

internal ConcurrentDictionary<string, Delegate> BoundFunctions { get; } = new();

internal TaskCompletionSource<IJSHandle> PuppeteerUtilTaskCompletionSource { get; private set; } = new(TaskCreationOptions.RunContinuationsAsynchronously);

internal Task<IJSHandle> GetPuppeteerUtilAsync() => PuppeteerUtilTaskCompletionSource.Task;

internal async Task AddBindingToContextAsync(ExecutionContext context, string name)
{
// Previous operation added the binding so we are done.
Expand Down Expand Up @@ -282,7 +275,7 @@ internal async Task<IElementHandle> WaitForSelectorInPageAsync(string queryOne,

var args = new List<object>
{
await GetPuppeteerUtilAsync().ConfigureAwait(false),
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
queryOne,
selector,
root,
Expand Down Expand Up @@ -438,7 +431,6 @@ internal void ClearContext()
{
_documentTask = null;
_contextResolveTaskWrapper = new TaskCompletionSource<ExecutionContext>(TaskCreationOptions.RunContinuationsAsynchronously);
PuppeteerUtilTaskCompletionSource = new TaskCompletionSource<IJSHandle>(TaskCreationOptions.RunContinuationsAsynchronously);
}

internal void SetContext(ExecutionContext context)
Expand All @@ -448,42 +440,11 @@ internal void SetContext(ExecutionContext context)
throw new ArgumentNullException(nameof(context));
}

_ = InjectPuppeteerUtil(context);
_ctxBindings.Clear();
_contextResolveTaskWrapper.TrySetResult(context);
TaskManager.RerunAll();
}

private static string GetInjectedSource()
{
if (string.IsNullOrEmpty(_injectedSource))
{
var assembly = Assembly.GetExecutingAssembly();
var resourceName = "PuppeteerSharp.Injected.injected.js";

using var stream = assembly.GetManifestResourceStream(resourceName);
using var reader = new StreamReader(stream);
var fileContent = reader.ReadToEnd();
_injectedSource = fileContent;
}

return _injectedSource;
}

private async Task InjectPuppeteerUtil(ExecutionContext context)
{
try
{
var injectedSource = GetInjectedSource();
var handle = await context.EvaluateExpressionHandleAsync(injectedSource).ConfigureAwait(false);
PuppeteerUtilTaskCompletionSource.TrySetResult(handle);
}
catch (Exception ex)
{
_logger.LogError(ex.ToString());
}
}

private async void Client_MessageReceived(object sender, MessageEventArgs e)
{
try
Expand Down
10 changes: 5 additions & 5 deletions lib/PuppeteerSharp/WaitTask.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,9 +105,9 @@ internal async Task Rerun()
}",
new object[]
{
await _isolatedWorld.GetPuppeteerUtilAsync().ConfigureAwait(false),
_pollingInterval,
_fn,
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
_pollingInterval,
_fn,
}.Concat(_args).ToArray()).ConfigureAwait(false);
}
else if (_polling == WaitForFunctionPollingOption.Raf)
Expand All @@ -122,7 +122,7 @@ internal async Task Rerun()
}",
new object[]
{
await _isolatedWorld.GetPuppeteerUtilAsync().ConfigureAwait(false),
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
_fn,
}.Concat(_args).ToArray()).ConfigureAwait(false);
}
Expand All @@ -138,7 +138,7 @@ internal async Task Rerun()
}",
new object[]
{
await _isolatedWorld.GetPuppeteerUtilAsync().ConfigureAwait(false),
new LazyArg(async context => await context.GetPuppeteerUtilAsync().ConfigureAwait(false)),
_root,
_fn,
}.Concat(_args).ToArray()).ConfigureAwait(false);
Expand Down

0 comments on commit 7cde7c4

Please sign in to comment.