Skip to content

Commit a1493a6

Browse files
asafchen-digasafchen-dig
authored
Merge pull request #40 from digma-ai/vertica-instrumentation
Vertica instrumentation
2 parents a9f3445 + 7fce50d commit a1493a6

File tree

5 files changed

+366
-185
lines changed

5 files changed

+366
-185
lines changed

‎src/OpenTelemetry.AutoInstrumentation.Digma/Instrumentation/SqlClientInstrumentation.cs‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Linq;
4+
using System.Reflection;
5+
using System.Reflection.Emit;
6+
using HarmonyLib;
7+
using OpenTelemetry.AutoInstrumentation.Digma.Utils;
8+
9+
namespace OpenTelemetry.AutoInstrumentation.Digma.Instrumentation;
10+
11+
public class SqlClientInstrumentation
12+
{
13+
private static readonly MethodInfo SqlTranspilerMethodInfo = typeof(SqlClientInstrumentation).GetMethod(nameof(Transpiler), BindingFlags.Static | BindingFlags.NonPublic);
14+
15+
private readonly Harmony _harmony;
16+
17+
public SqlClientInstrumentation(Harmony harmony)
18+
{
19+
_harmony = harmony;
20+
}
21+
22+
public void Instrument(Assembly systemDataAssembly)
23+
{
24+
try
25+
{
26+
var sqlCommandType = systemDataAssembly.GetType("System.Data.SqlClient.SqlCommand", throwOnError: false);
27+
if (sqlCommandType == null)
28+
{
29+
Logger.LogError("System.Data.SqlClient.SqlCommand not found.");
30+
return;
31+
}
32+
33+
var targetMethodInfo =
34+
sqlCommandType.GetMethod("WriteBeginExecuteEvent", BindingFlags.Instance | BindingFlags.NonPublic);
35+
if (targetMethodInfo == null)
36+
{
37+
Logger.LogError("WriteBeginExecuteEvent not found.");
38+
return;
39+
}
40+
41+
_harmony.Patch(targetMethodInfo, transpiler: new HarmonyMethod(SqlTranspilerMethodInfo));
42+
Logger.LogInfo("Patched SqlClient");
43+
}
44+
catch (Exception e)
45+
{
46+
Logger.LogError("Failed to patch System.Data.SqlClient.SqlCommand.WriteBeginExecuteEvent", e);
47+
}
48+
}
49+
50+
private static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
51+
{
52+
var instructionList = new List<CodeInstruction>(instructions);
53+
54+
Logger.LogDebug($"Before ({instructionList.Count}):\n"+
55+
string.Join("\n",instructionList.Select(x => x.ToString()).ToArray()));
56+
57+
for (var i = 0; i < instructionList.Count-1; i++)
58+
{
59+
if (instructionList[i].opcode == OpCodes.Ldarg_0 && // "this"
60+
instructionList[i+1].operand.ToString().Contains("System.Data.CommandType get_CommandType()"))
61+
{
62+
instructionList.RemoveRange(i, 6);
63+
break;
64+
}
65+
}
66+
67+
Logger.LogDebug($"After ({instructionList.Count}):\n"+
68+
string.Join("\n",instructionList.Select(x => x.ToString()).ToArray()));
69+
70+
return instructionList;
71+
}
72+
}

‎src/OpenTelemetry.AutoInstrumentation.Digma/Instrumentation/UserCodeInstrumentation.cs‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
using System;
2+
using System.Diagnostics;
3+
using System.Linq;
4+
using System.Reflection;
5+
using HarmonyLib;
6+
using OpenTelemetry.AutoInstrumentation.Digma.Utils;
7+
8+
namespace OpenTelemetry.AutoInstrumentation.Digma.Instrumentation;
9+
10+
public class UserCodeInstrumentation
11+
{
12+
private static readonly ActivitySourceProvider ActivitySourceProvider = new();
13+
private static readonly MethodInfo PrefixMethodInfo = typeof(UserCodeInstrumentation).GetMethod(nameof(Prefix), BindingFlags.Static | BindingFlags.NonPublic);
14+
private static readonly MethodInfo FinalizerMethodInfo = typeof(UserCodeInstrumentation).GetMethod(nameof(Finalizer), BindingFlags.Static | BindingFlags.NonPublic);
15+
16+
private readonly Harmony _harmony;
17+
private readonly string[] _namespaces;
18+
private readonly bool _includePrivateMethods;
19+
20+
public UserCodeInstrumentation(Harmony harmony)
21+
{
22+
_harmony = harmony;
23+
var namespacesStr = Environment.GetEnvironmentVariable("OTEL_DOTNET_AUTO_NAMESPACES");
24+
_namespaces = namespacesStr?.Split(',')
25+
.Select(x => x.Trim())
26+
.Where(x => !string.IsNullOrWhiteSpace(x))
27+
.ToArray()
28+
?? Array.Empty<string>();
29+
_includePrivateMethods = Environment.GetEnvironmentVariable("OTEL_DOTNET_AUTO_INCLUDE_PRIVATE_METHODS")
30+
?.Equals("true", StringComparison.InvariantCultureIgnoreCase) == true;
31+
32+
Logger.LogInfo($"Requested to auto-instrument {_namespaces.Length} namespaces:\n"+
33+
string.Join("\n", _namespaces));
34+
}
35+
36+
public void Instrument(Assembly assembly)
37+
{
38+
var name = assembly.GetName().Name;
39+
if (ShouldInstrumentAssembly(name))
40+
{
41+
var relevantTypes = assembly.GetTypes().Where(ShouldInstrumentType).ToArray();
42+
var methods = relevantTypes
43+
.SelectMany(t => t
44+
.GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic)
45+
.Where(m => ShouldInstrumentMethod(t,m)))
46+
.ToArray();
47+
foreach (var method in methods)
48+
{
49+
PatchMethod(method);
50+
}
51+
}
52+
}
53+
54+
private void PatchMethod(MethodInfo originalMethodInfo)
55+
{
56+
var methodFullName = $"{originalMethodInfo.DeclaringType?.FullName}.{originalMethodInfo.Name}";
57+
try
58+
{
59+
var prefix = DoesAlreadyStartActivity(originalMethodInfo)
60+
? null
61+
: new HarmonyMethod(PrefixMethodInfo);
62+
63+
var finalizer = new HarmonyMethod(FinalizerMethodInfo);
64+
65+
_harmony.Patch(originalMethodInfo, prefix: prefix, finalizer: finalizer);
66+
Logger.LogInfo($"Patched method {methodFullName}"+(prefix==null?" (partially)":""));
67+
}
68+
catch (Exception e)
69+
{
70+
Logger.LogError($"Failed to patch {methodFullName}", e);
71+
}
72+
73+
}
74+
75+
private bool ShouldInstrumentAssembly(string assemblyName)
76+
{
77+
return _namespaces.Any(ns => ns.StartsWith(assemblyName, StringComparison.OrdinalIgnoreCase) ||
78+
assemblyName.StartsWith(ns, StringComparison.OrdinalIgnoreCase));
79+
}
80+
81+
private bool ShouldInstrumentType(Type type)
82+
{
83+
return !typeof(Delegate).IsAssignableFrom(type) &&
84+
!type.IsGenericType &&
85+
_namespaces.Any(ns => type.FullName?.StartsWith(ns, StringComparison.OrdinalIgnoreCase) == true);
86+
}
87+
88+
private bool ShouldInstrumentMethod(Type type, MethodInfo methodInfo)
89+
{
90+
return methodInfo.DeclaringType == type &&
91+
!methodInfo.IsAbstract &&
92+
!methodInfo.IsSpecialName && // property accessors and operator overloading methods
93+
!methodInfo.IsGenericMethod &&
94+
methodInfo.Name != "GetHashCode" &&
95+
methodInfo.Name != "Equals" &&
96+
methodInfo.Name != "ToString" &&
97+
methodInfo.Name != "Deconstruct" &&
98+
methodInfo.Name != "<Clone>$" &&
99+
(methodInfo.IsPublic || _includePrivateMethods);
100+
}
101+
102+
private bool DoesAlreadyStartActivity(MethodInfo methodInfo)
103+
{
104+
var instructions = PatchProcessor.GetOriginalInstructions(methodInfo);
105+
foreach (var instruction in instructions)
106+
{
107+
if (instruction.operand is MethodInfo call &&
108+
call.DeclaringType == typeof(ActivitySource) &&
109+
call.Name == nameof(System.Diagnostics.ActivitySource.StartActivity))
110+
{
111+
return true;
112+
}
113+
}
114+
115+
return false;
116+
}
117+
118+
private static void Prefix(MethodBase __originalMethod, out Activity __state)
119+
{
120+
var activitySource = ActivitySourceProvider.GetOrCreate(__originalMethod.DeclaringType!);
121+
var activity = activitySource.StartActivity(__originalMethod.Name);
122+
Logger.LogDebug($"Opened Activity: {activity?.Source.Name}.{activity?.OperationName}");
123+
activity?.SetTag(DigmaSemanticConventions.ExtendedObservabilityPackage, __originalMethod.DeclaringType?.Assembly.GetName().Name);
124+
activity?.SetCodeTags(__originalMethod);
125+
__state = activity;
126+
}
127+
128+
private static void Finalizer(MethodBase __originalMethod, Activity __state, Exception __exception)
129+
{
130+
var activity = __state;
131+
if (activity == null)
132+
return;
133+
134+
if (__exception != null)
135+
{
136+
activity.RecordException(__exception);
137+
activity.SetStatus(ActivityStatusCode.Error);
138+
}
139+
activity.Dispose();
140+
Logger.LogDebug($"Closed Activity: {activity.Source.Name}.{activity.OperationName}");
141+
}
142+
}

‎src/OpenTelemetry.AutoInstrumentation.Digma/Instrumentation/VerticaInstrumentation.cs‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
using System;
2+
using System.Diagnostics;
3+
using System.Linq;
4+
using System.Reflection;
5+
using HarmonyLib;
6+
using OpenTelemetry.AutoInstrumentation.Digma.Utils;
7+
8+
namespace OpenTelemetry.AutoInstrumentation.Digma.Instrumentation;
9+
10+
public class VerticaInstrumentation
11+
{
12+
private static readonly ActivitySource ActivitySource = new("OpenTelemetry.Instrumentation.Vertica");
13+
14+
private static readonly MethodInfo PrefixMethodInfo = typeof(VerticaInstrumentation).GetMethod(nameof(Prefix), BindingFlags.Static | BindingFlags.NonPublic);
15+
private static readonly MethodInfo FinalizerMethodInfo = typeof(VerticaInstrumentation).GetMethod(nameof(Finalizer), BindingFlags.Static | BindingFlags.NonPublic);
16+
17+
private static MethodInfo CommandTextGetter;
18+
private static MethodInfo ConnectionGetter;
19+
private static MethodInfo DatabaseGetter;
20+
21+
private readonly Harmony _harmony;
22+
23+
private readonly string[] TargetMethodNames =
24+
{
25+
"ExecuteReader",
26+
"ExecuteScalar",
27+
"ExecuteNonQuery",
28+
};
29+
30+
public VerticaInstrumentation(Harmony harmony)
31+
{
32+
_harmony = harmony;
33+
}
34+
35+
public void Instrument(Assembly verticaDataAssembly)
36+
{
37+
try
38+
{
39+
var verticaCommandType = verticaDataAssembly.GetType("Vertica.Data.VerticaClient.VerticaCommand", throwOnError: false);
40+
if (verticaCommandType == null)
41+
{
42+
Logger.LogError("Vertica.Data.VerticaClient.VerticaCommand not found.");
43+
return;
44+
}
45+
46+
var sCommandType = verticaDataAssembly.GetType("Vertica.Data.Internal.ADO.Net.SCommand", throwOnError: false);
47+
if (sCommandType == null)
48+
{
49+
Logger.LogError("Vertica.Data.Internal.ADO.Net.SCommand not found.");
50+
return;
51+
}
52+
53+
CommandTextGetter = sCommandType.GetProperty("CommandText")?.GetMethod;
54+
if (CommandTextGetter == null)
55+
{
56+
Logger.LogError("Vertica.Data.Internal.ADO.Net.SCommand.CommandText getter not found.");
57+
return;
58+
}
59+
60+
ConnectionGetter = sCommandType.GetProperties()
61+
.FirstOrDefault(x => x.Name == "Connection" && x.DeclaringType == sCommandType)?.GetMethod;
62+
if (ConnectionGetter == null)
63+
{
64+
Logger.LogError("Vertica.Data.Internal.ADO.Net.SCommand.Connection getter not found.");
65+
return;
66+
}
67+
68+
var sConnectionType = verticaDataAssembly.GetType("Vertica.Data.Internal.ADO.Net.SConnection", throwOnError: false);
69+
if (sConnectionType == null)
70+
{
71+
Logger.LogError("Vertica.Data.Internal.ADO.Net.SConnection not found.");
72+
return;
73+
}
74+
DatabaseGetter = sConnectionType.GetProperty("Database")?.GetMethod;
75+
if (ConnectionGetter == null)
76+
{
77+
Logger.LogError("Vertica.Data.Internal.ADO.Net.SConnection.Database getter not found.");
78+
return;
79+
}
80+
81+
var methodInfos = verticaCommandType.GetMethods(BindingFlags.Public | BindingFlags.Instance)
82+
.Where(x => TargetMethodNames.Contains(x.Name))
83+
.Where(x => x.DeclaringType == verticaCommandType)
84+
.ToArray();
85+
foreach (var methodInfo in methodInfos)
86+
{
87+
_harmony.Patch(methodInfo, prefix: PrefixMethodInfo, finalizer: FinalizerMethodInfo);
88+
Logger.LogInfo($"Patched {methodInfo.FullDescription()}");
89+
}
90+
}
91+
catch (Exception e)
92+
{
93+
Logger.LogError("Failed to patch Vertica.Data.VerticaClient.VerticaCommand", e);
94+
}
95+
}
96+
97+
private static void Prefix(MethodBase __originalMethod, object __instance, out Activity __state)
98+
{
99+
var connection = ConnectionGetter.Invoke(__instance, null);
100+
var database = DatabaseGetter.Invoke(connection, null)?.ToString();
101+
var sqlStatement = CommandTextGetter.Invoke(__instance, null);
102+
103+
var activity = ActivitySource.StartActivity(database ?? __originalMethod.Name, ActivityKind.Client);
104+
Logger.LogDebug($"Opened Activity: {activity?.Source.Name}.{activity?.OperationName}");
105+
activity?.SetTag("db.statement", sqlStatement);
106+
activity?.SetTag("db.system", "vertica");
107+
activity?.SetTag("db.name", database);
108+
109+
__state = activity;
110+
}
111+
112+
private static void Finalizer(MethodBase __originalMethod, Activity __state, Exception __exception)
113+
{
114+
var activity = __state;
115+
if (activity == null)
116+
return;
117+
118+
if (__exception != null)
119+
{
120+
activity.RecordException(__exception);
121+
activity.SetStatus(ActivityStatusCode.Error);
122+
}
123+
activity.Dispose();
124+
Logger.LogDebug($"Closed Activity: {activity.Source.Name}.{activity.OperationName}");
125+
}
126+
}

0 commit comments

Comments
 (0)