otelgorm replaces (but does not restore) context from statement

[markhildreth-gravity]

The otelgorm package utilizes replacement the statement context when the before is called...

opentelemetry-go-extra/otelgorm/otelgorm.go

Lines 97 to 101 in fcc6618

	func (p *otelPlugin) before(spanName string) gormHookFunc {
	return func(tx *gorm.DB) {
	tx.Statement.Context, _ = p.tracer.Start(tx.Statement.Context, spanName, trace.WithSpanKind(trace.SpanKindClient))
	}
	}

However, the statement context is never restored to its original in the corresponding after() call. Thus, after the execution of the query finishes, if the db.Statement object is reused, then the context of the statement will still be using the context corresponding to the now-ended span.

The following is a failing test case:

  {
    do: func(ctx context.Context, db *gorm.DB) {
      var count int64
      query := db.WithContext(ctx).Table("generate_series(1, 10)")
      _, _ = query.Select("*").Rows()
      _ = query.Count(&count)
    },
    require: func(t *testing.T, spans []sdktrace.ReadOnlySpan) {
      require.Equal(t, 2, len(spans))
      require.Equal(
        t,
        spans[0].Parent().SpanID().String(),
        spans[1].Parent().SpanID().String(),
      )
    },
  }

In do, you can see two separate queries being run (with Rows() and Count()). To me, the spans generated for these two queries should have the same parent. Instead, otelgorm is making the span of the second query a child of the first. If I added a third execution using the query value, that span would be a child of the second. Ideally, the plugin would restore the statement back to it's original so the statement could be reused.

The only way I could think of doing this was to store the original context in the newly derived context, and then during the after callback pop it out. But that seemed error prone given that Gorm doesn't enforce the ordering of callbacks, so we'd have to hope another callback isn't trying to do the same thing but with a different ordering.

[vmihailenco]

The only way I could think of doing this was to store the original context in the newly derived context, and then during the after callback pop it out.

I don't see why this would not work. Do you want to give it a chance?

But that seemed error prone given that Gorm doesn't enforce the ordering of callbacks, so we'd have to hope another callback isn't trying to do the same thing but with a different ordering.

Is not *gorm.DB always executing a single query/callback on a *gorm.DB? Otherwise it would be unsafe to set tx.Statement.Context in before...

[markhildreth-gravity]

The only way I could think of doing this was to store the original context in the newly derived context, and then during the after callback pop it out.

I don't see why this would not work. Do you want to give it a chance?

I did try it out, and it did work. My concern is with plugin ordering (see below).

But that seemed error prone given that Gorm doesn't enforce the ordering of callbacks, so we'd have to hope another callback isn't trying to do the same thing but with a different ordering.

Is not *gorm.DB always executing a single query/callback on a *gorm.DB? Otherwise it would be unsafe to set tx.Statement.Context in before...

Not sure I fully understand the question. If you're asking if gorm.DB is executing the callbacks in a single goroutine, then yes, you're correct. My concern is what happens if otelgorm saves the current context as the first before callback method, then restores it as the last after callback method.

before: [otelgorm_before, some_other_plugin_before, ...builtins...]
after: [...builtins..., otelgorm_after]

In the example above, anything that some_other_plugin_before adds to the context would be wiped out by otelgorm_after. I can't think of a reason why someone would want to modify the context in a before and have it survive the callback chains, but I'm not confident enough to say such a use case doesn't exist.

[vmihailenco]

My concern is what happens if otelgorm saves the current context as the first before callback method, then restores it as the last after callback method.

Got you. The order of callbacks seems to be important to end spans in the correct order, but it probably should still work given that you pop the previous context from the current context like in Matryoshka dolls.

Anyway, this is up to you. Just trying to help.