Skip to content

Commit a5a634c

Browse files
arturopalaarturopala
committed
refactor out log capture function
1 parent fa1c06b commit a5a634c

File tree

5 files changed

+182
-145
lines changed

5 files changed

+182
-145
lines changed

‎.github/workflows/test.yaml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
name: Test
2+
3+
on:
4+
workflow_dispatch:
5+
push:
6+
7+
permissions:
8+
id-token: write # This is required for requesting the JWT
9+
contents: write # This is required for actions/checkout
10+
11+
jobs:
12+
runTests:
13+
runs-on: ubuntu-latest
14+
steps:
15+
- uses: actions/checkout@v4
16+
with:
17+
fetch-depth: 0
18+
- uses: coursier/cache-action@v6
19+
- uses: VirtusLab/scala-cli-setup@main
20+
with:
21+
jvm: adoptium:1.21
22+
apps: scala
23+
- name: Run tests
24+
run: scala test .
25+
26+

‎EventHandler.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ package org.encalmo.lambda
22

33
trait EventHandler {
44

5-
/** Abstract lambda invocation handler method. Provide your business logic here.
5+
/** Provide your lambda business logic here.
66
*/
77
def handleRequest(input: String)(using LambdaContext): String
88
}

‎LambdaRuntime.scala

Lines changed: 135 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -141,9 +141,9 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
141141
if (active.get)
142142
then
143143
try {
144-
reportError(createErrorMessage(e), lambdaEnvironment.initErrorUrl)
144+
reportError(LambdaRuntime.createErrorMessage(e), lambdaEnvironment.initErrorUrl)
145145
} catch {
146-
case e => error(createErrorMessage(e))
146+
case e => error(LambdaRuntime.createErrorMessage(e))
147147
}
148148
} finally {
149149
lambdaInvocationDebugMode = lambdaEnvironment.isDebugMode
@@ -200,17 +200,19 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
200200
case NonFatal(e) =>
201201
given le: LambdaEnvironment = new LambdaEnvironment()
202202
try {
203-
reportError(createErrorMessage(e), le.initErrorUrl)
203+
reportError(LambdaRuntime.createErrorMessage(e), le.initErrorUrl)
204204
} catch {
205-
case e => error(createErrorMessage(e))
205+
case e => error(LambdaRuntime.createErrorMessage(e))
206206
}
207207
le.resetOut()
208208
throw e
209209
}
210210
}
211211
}
212212

213-
/** The actual lambda invocation happens here. */
213+
/* ----------------------------------------------------------
214+
* THE ACTUAL LAMBDA BUSINESS METHOD INVOCATION HAPPENS HERE.
215+
* ---------------------------------------------------------- */
214216
private final def invokeHandleRequest(
215217
id: Int
216218
)(using lambdaEnvironment: LambdaEnvironment): Unit = {
@@ -265,7 +267,9 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
265267
.maxMemory()},"totalMemory":${Runtime
266268
.getRuntime()
267269
.totalMemory()},"freeMemory":${Runtime.getRuntime().freeMemory()}}"""
268-
else s"[$id]$tag ${REQUEST}LAMBDA REQUEST ${AnsiColor.BOLD}${event.body}"
270+
else {
271+
s"[$id]$tag ${REQUEST}LAMBDA REQUEST ${AnsiColor.BOLD}${event.body}"
272+
}
269273
)
270274

271275
try {
@@ -287,78 +291,44 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
287291

288292
val input = event.body
289293

290-
val logPrinter: PrintStream | NoAnsiColorJsonArray | NoAnsiColorJsonString | NoAnsiColorsSingleLine =
291-
val logPrefix =
292-
if (lambdaEnvironment.shouldLogStructuredJson)
293-
then s"""{"log":"LOGS",$structuredLogIntro,"""
294-
else s"[${lambdaEnvironment.getFunctionName()}] [$id]$tag LAMBDA LOGS {"
295-
296-
val logSuffix =
297-
if (lambdaEnvironment.shouldLogStructuredJson)
298-
then s""",$structuredLogEnd}"""
299-
else "}"
300-
301-
if (lambdaInvocationDebugMode) then {
302-
if (lambdaEnvironment.shouldDisplayAnsiColors)
303-
then LambdaEnvironment.originalOut
304-
else if (lambdaEnvironment.shouldLogInJsonArrayFormat)
305-
then new NoAnsiColorJsonArray(logPrefix, logSuffix, LambdaEnvironment.originalOut)
306-
else if (lambdaEnvironment.shouldLogInJsonStringFormat)
307-
then new NoAnsiColorJsonString(logPrefix, logSuffix, LambdaEnvironment.originalOut)
308-
else new NoAnsiColorsSingleLine(logPrefix, logSuffix, LambdaEnvironment.originalOut)
309-
} else NoOpPrinter.out
310-
311-
extension (
312-
v: PrintStream | NoAnsiColorJsonArray | NoAnsiColorsSingleLine | NoAnsiColorJsonString
313-
)
314-
inline def out: PrintStream = v match {
315-
case out: PrintStream => out
316-
case ps: NoAnsiColorJsonArray => ps.out
317-
case ps: NoAnsiColorJsonString => ps.out
318-
case ps: NoAnsiColorsSingleLine => ps.out
319-
}
320-
def close(): Unit = v match {
321-
case out: PrintStream => ()
322-
case ps: NoAnsiColorJsonArray => ps.close()
323-
case ps: NoAnsiColorJsonString => ps.close()
324-
case ps: NoAnsiColorsSingleLine => ps.close()
325-
}
326-
327-
val previousSystemOut = System.out
328-
var isError = false
294+
val logPrefix =
295+
if (lambdaEnvironment.shouldLogStructuredJson)
296+
then s"""{"log":"LOGS",$structuredLogIntro,"""
297+
else s"[${lambdaEnvironment.getFunctionName()}] [$id]$tag LAMBDA LOGS {"
329298

330-
val output =
331-
try {
332-
System.setOut(logPrinter.out)
333-
System.setErr(logPrinter.out)
334-
Console.withOut(logPrinter.out) {
335-
Console.withErr(logPrinter.out) {
336-
try {
337-
handleRequest(input)(using context).trim()
338-
} catch {
339-
case NonFatal(e) =>
340-
isError = true
341-
error(e.toString())
342-
val stackTrace = e
343-
.getStackTrace()
344-
.filterNot { s =>
345-
val n = s.getClassName()
346-
n.startsWith("scala") || n.startsWith("java")
347-
}
348-
if (stackTrace.size > 30) then
349-
stackTrace.take(15).foreach(println)
350-
println("...")
351-
stackTrace.takeRight(15).foreach(println)
352-
else stackTrace.foreach(println)
353-
createErrorMessage(e)
354-
}
299+
val logSuffix =
300+
if (lambdaEnvironment.shouldLogStructuredJson)
301+
then s""",$structuredLogEnd}"""
302+
else "}"
303+
304+
val result: Either[String, String] =
305+
LambdaRuntime
306+
.withLogCapture(
307+
lambdaEnvironment,
308+
logPrefix,
309+
logSuffix,
310+
lambdaInvocationDebugMode,
311+
try {
312+
Right(handleRequest(input)(using context).trim())
313+
} catch {
314+
case NonFatal(e) =>
315+
error(e.toString())
316+
val stackTrace = e
317+
.getStackTrace()
318+
.filterNot { s =>
319+
val n = s.getClassName()
320+
n.startsWith("scala") || n.startsWith("java")
321+
}
322+
if (stackTrace.size > 30) then
323+
stackTrace.take(15).foreach(println)
324+
println("...")
325+
stackTrace.takeRight(15).foreach(println)
326+
else stackTrace.foreach(println)
327+
Left(LambdaRuntime.createErrorMessage(e))
355328
}
356-
}
357-
} finally {
358-
logPrinter.close()
359-
System.setOut(previousSystemOut)
360-
System.setErr(previousSystemOut)
361-
}
329+
)
330+
331+
val output = result.fold(identity, identity)
362332

363333
val isJsonReponse =
364334
(output.startsWith("{") && output.endsWith("}"))
@@ -375,7 +345,7 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
375345
debug(
376346
if (lambdaEnvironment.shouldLogStructuredJson)
377347
then
378-
s"""{$awsEmbededMetric,"log":"RESPONSE",$structuredLogIntro,${
348+
s"""{"log":"RESPONSE",$structuredLogIntro,${
379349
if (lambdaEnvironment.shouldLogResponseIncludeRequest)
380350
then
381351
s""""request":${if (isJsonRequest) then event.body else s"\"${event.body.replace("\"", "\\\"")}\""},"""
@@ -390,12 +360,14 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
390360
.getRuntime()
391361
.maxMemory()},"totalMemory":${Runtime
392362
.getRuntime()
393-
.totalMemory()},"freeMemory":${Runtime.getRuntime().freeMemory()}}"""
394-
else s"[$id]$tag ${RESPONSE}LAMBDA RESPONSE [${t1 - t0}ms] ${AnsiColor.BOLD}$output"
363+
.totalMemory()},"freeMemory":${Runtime.getRuntime().freeMemory()},$awsEmbededMetric,}"""
364+
else {
365+
s"[$id]$tag ${RESPONSE}LAMBDA RESPONSE [${t1 - t0}ms] ${AnsiColor.BOLD}$output"
366+
}
395367
)
396368

397369
val responseResult =
398-
if (isError)
370+
if (result.isLeft)
399371
then reportError(output, lambdaEnvironment.errorUrl(requestId))
400372
else
401373
httpClient
@@ -410,42 +382,16 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
410382

411383
} catch {
412384
case NonFatal(e) =>
413-
val errorMessage = createErrorMessage(e)
385+
val errorMessage = LambdaRuntime.createErrorMessage(e)
414386
error(errorMessage)
415387
try {
416388
reportError(errorMessage, lambdaEnvironment.errorUrl(requestId))
417389
} catch {
418-
case e => error(createErrorMessage(e))
390+
case e => error(LambdaRuntime.createErrorMessage(e))
419391
}
420392
}
421393
}
422394

423-
/** Helper method to report error back to the AWS lambda host. */
424-
private final inline def reportError(
425-
errorMessage: String,
426-
errorUrl: URI
427-
)(using lambdaEnvironment: LambdaEnvironment): HttpResponse[String] =
428-
httpClient
429-
.send(
430-
HttpRequest
431-
.newBuilder(errorUrl)
432-
.POST(BodyPublishers.ofString(errorMessage))
433-
.setHeader("Lambda-Runtime-Function-Error-Type", "Runtime.UnknownReason")
434-
.build(),
435-
BodyHandlers.ofString()
436-
)
437-
438-
private final def createErrorMessage(e: Throwable): String =
439-
val stackTrace = e.getStackTrace().filterNot { s =>
440-
val n = s.getClassName()
441-
n.startsWith("scala") || n.startsWith("java")
442-
}
443-
s"{\"success\":false,\"errorMessage\":\"${e.getMessage()}\", \"error\":\"${e
444-
.getClass()
445-
.getName()}\", \"stackTrace\": [${(stackTrace.take(3).map(_.toString()) ++ Array("...") ++ stackTrace
446-
.takeRight(3)
447-
.map(_.toString())).map(s => s"\"$s\"").mkString(",")}]}"
448-
449395
// Lambda invocation for unit test purposes, does not interact with the AWS host, returns result directly. */
450396
final def test(
451397
input: String,
@@ -489,8 +435,23 @@ trait LambdaRuntime extends EventHandler, EventHandlerTag {
489435
output
490436
} catch {
491437
case NonFatal(e) =>
492-
createErrorMessage(e)
438+
LambdaRuntime.createErrorMessage(e)
493439
}
440+
441+
/** Report error back to the AWS lambda host. */
442+
final inline def reportError(
443+
errorMessage: String,
444+
errorUrl: URI
445+
)(using lambdaEnvironment: LambdaEnvironment): HttpResponse[String] =
446+
httpClient
447+
.send(
448+
HttpRequest
449+
.newBuilder(errorUrl)
450+
.POST(BodyPublishers.ofString(errorMessage))
451+
.setHeader("Lambda-Runtime-Function-Error-Type", "Runtime.UnknownReason")
452+
.build(),
453+
BodyHandlers.ofString()
454+
)
494455
}
495456

496457
object LambdaRuntime {
@@ -501,4 +462,69 @@ object LambdaRuntime {
501462
}
502463

503464
inline def durationMetric = """{"Name":"duration","Unit":"Milliseconds","StorageResolution":60}"""
465+
466+
final def createErrorMessage(e: Throwable): String =
467+
val stackTrace = e.getStackTrace().filterNot { s =>
468+
val n = s.getClassName()
469+
n.startsWith("scala") || n.startsWith("java")
470+
}
471+
s"{\"success\":false,\"errorMessage\":\"${e.getMessage()}\", \"error\":\"${e
472+
.getClass()
473+
.getName()}\", \"stackTrace\": [${(stackTrace.take(3).map(_.toString()) ++ Array("...") ++ stackTrace
474+
.takeRight(3)
475+
.map(_.toString())).map(s => s"\"$s\"").mkString(",")}]}"
476+
477+
final def withLogCapture(
478+
lambdaEnvironment: LambdaEnvironment,
479+
logPrefix: String,
480+
logSuffix: String,
481+
debugMode: Boolean,
482+
body: => Either[String, String]
483+
): Either[String, String] = {
484+
485+
val logPrinter: PrintStream | NoAnsiColorJsonArray | NoAnsiColorJsonString | NoAnsiColorsSingleLine =
486+
if (debugMode) then {
487+
if (lambdaEnvironment.shouldDisplayAnsiColors)
488+
then LambdaEnvironment.originalOut
489+
else if (lambdaEnvironment.shouldLogInJsonArrayFormat)
490+
then new NoAnsiColorJsonArray(logPrefix, logSuffix, LambdaEnvironment.originalOut)
491+
else if (lambdaEnvironment.shouldLogInJsonStringFormat)
492+
then new NoAnsiColorJsonString(logPrefix, logSuffix, LambdaEnvironment.originalOut)
493+
else new NoAnsiColorsSingleLine(logPrefix, logSuffix, LambdaEnvironment.originalOut)
494+
} else NoOpPrinter.out
495+
496+
extension (
497+
v: PrintStream | NoAnsiColorJsonArray | NoAnsiColorsSingleLine | NoAnsiColorJsonString
498+
)
499+
inline def out: PrintStream = v match {
500+
case out: PrintStream => out
501+
case ps: NoAnsiColorJsonArray => ps.out
502+
case ps: NoAnsiColorJsonString => ps.out
503+
case ps: NoAnsiColorsSingleLine => ps.out
504+
}
505+
506+
inline def close(): Unit = v match {
507+
case out: PrintStream => ()
508+
case ps: NoAnsiColorJsonArray => ps.close()
509+
case ps: NoAnsiColorJsonString => ps.close()
510+
case ps: NoAnsiColorsSingleLine => ps.close()
511+
}
512+
513+
val previousSystemOut = System.out
514+
var isError = false
515+
516+
try {
517+
System.setOut(logPrinter.out)
518+
System.setErr(logPrinter.out)
519+
Console.withOut(logPrinter.out) {
520+
Console.withErr(logPrinter.out) {
521+
body
522+
}
523+
}
524+
} finally {
525+
logPrinter.close()
526+
System.setOut(previousSystemOut)
527+
System.setErr(previousSystemOut)
528+
}
529+
}
504530
}

0 commit comments

Comments
 (0)