1package dataloader
2
3import (
4	"context"
5
6	opentracing "github.com/opentracing/opentracing-go"
7)
8
9type TraceLoadFinishFunc func(Thunk)
10type TraceLoadManyFinishFunc func(ThunkMany)
11type TraceBatchFinishFunc func([]*Result)
12
13// Tracer is an interface that may be used to implement tracing.
14type Tracer interface {
15	// TraceLoad will trace the calls to Load
16	TraceLoad(ctx context.Context, key Key) (context.Context, TraceLoadFinishFunc)
17	// TraceLoadMany will trace the calls to LoadMany
18	TraceLoadMany(ctx context.Context, keys Keys) (context.Context, TraceLoadManyFinishFunc)
19	// TraceBatch will trace data loader batches
20	TraceBatch(ctx context.Context, keys Keys) (context.Context, TraceBatchFinishFunc)
21}
22
23// OpenTracing Tracer implements a tracer that can be used with the Open Tracing standard.
24type OpenTracingTracer struct{}
25
26// TraceLoad will trace a call to dataloader.LoadMany with Open Tracing
27func (OpenTracingTracer) TraceLoad(ctx context.Context, key Key) (context.Context, TraceLoadFinishFunc) {
28	span, spanCtx := opentracing.StartSpanFromContext(ctx, "Dataloader: load")
29
30	span.SetTag("dataloader.key", key.String())
31
32	return spanCtx, func(thunk Thunk) {
33		// TODO: is there anything we should do with the results?
34		span.Finish()
35	}
36}
37
38// TraceLoadMany will trace a call to dataloader.LoadMany with Open Tracing
39func (OpenTracingTracer) TraceLoadMany(ctx context.Context, keys Keys) (context.Context, TraceLoadManyFinishFunc) {
40	span, spanCtx := opentracing.StartSpanFromContext(ctx, "Dataloader: loadmany")
41
42	span.SetTag("dataloader.keys", keys.Keys())
43
44	return spanCtx, func(thunk ThunkMany) {
45		// TODO: is there anything we should do with the results?
46		span.Finish()
47	}
48}
49
50// TraceBatch will trace a call to dataloader.LoadMany with Open Tracing
51func (OpenTracingTracer) TraceBatch(ctx context.Context, keys Keys) (context.Context, TraceBatchFinishFunc) {
52	span, spanCtx := opentracing.StartSpanFromContext(ctx, "Dataloader: batch")
53
54	span.SetTag("dataloader.keys", keys.Keys())
55
56	return spanCtx, func(results []*Result) {
57		// TODO: is there anything we should do with the results?
58		span.Finish()
59	}
60}
61
62// NoopTracer is the default (noop) tracer
63type NoopTracer struct{}
64
65// TraceLoad is a noop function
66func (NoopTracer) TraceLoad(ctx context.Context, key Key) (context.Context, TraceLoadFinishFunc) {
67	return ctx, func(Thunk) {}
68}
69
70// TraceLoadMany is a noop function
71func (NoopTracer) TraceLoadMany(ctx context.Context, keys Keys) (context.Context, TraceLoadManyFinishFunc) {
72	return ctx, func(ThunkMany) {}
73}
74
75// TraceBatch is a noop function
76func (NoopTracer) TraceBatch(ctx context.Context, keys Keys) (context.Context, TraceBatchFinishFunc) {
77	return ctx, func(result []*Result) {}
78}
79