Add tracing support to LDAP requests

ldap/pool.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff"
+	"go.opencensus.io/trace"
 	"gopkg.in/ldap.v2"
 )
 
@@ -147,9 +148,18 @@ func NewConnectionPool(uri, bindDN, bindPw string, cacheSize int) (*ConnectionPo
 	}, nil
 }
 
-func (p *ConnectionPool) doRequest(ctx context.Context, fn func(*ldap.Conn) error) error {
-	return backoff.Retry(func() error {
-		conn, err := p.Get(ctx)
+func (p *ConnectionPool) doRequest(ctx context.Context, name string, attrs []trace.Attribute, fn func(*ldap.Conn) error) error {
+	// Tracing: initialize a new client span.
+	sctx, span := trace.StartSpan(ctx, name,
+		trace.WithSpanKind(trace.SpanKindClient))
+	defer span.End()
+
+	if len(attrs) > 0 {
+		span.AddAttributes(attrs...)
+	}
+
+	rerr := backoff.Retry(func() error {
+		conn, err := p.Get(sctx)
 		if err != nil {
 			// Here conn is nil, so we don't need to Release it.
 			if isTemporaryLDAPError(err) {
@@ -158,7 +168,7 @@ func (p *ConnectionPool) doRequest(ctx context.Context, fn func(*ldap.Conn) erro
 			return backoff.Permanent(err)
 		}
 
-		if deadline, ok := ctx.Deadline(); ok {
+		if deadline, ok := sctx.Deadline(); ok {
 			conn.SetTimeout(time.Until(deadline))
 		}
 
@@ -169,30 +179,42 @@ func (p *ConnectionPool) doRequest(ctx context.Context, fn func(*ldap.Conn) erro
 		}
 		return err
 	}, backoff.WithContext(newExponentialBackOff(), ctx))
+
+	// Tracing: set the final status.
+	span.SetStatus(errorToTraceStatus(rerr))
+
+	return rerr
 }
 
 // Search performs the given search request. It will retry the request
 // on temporary errors.
 func (p *ConnectionPool) Search(ctx context.Context, searchRequest *ldap.SearchRequest) (*ldap.SearchResult, error) {
 	var result *ldap.SearchResult
-	err := p.doRequest(ctx, func(conn *ldap.Conn) error {
-		var err error
-		result, err = conn.Search(searchRequest)
-		return err
+	err := p.doRequest(ctx, "ldap.Search", []trace.Attribute{
+		trace.StringAttribute("ldap.base", searchRequest.BaseDN),
+		trace.StringAttribute("ldap.filter", searchRequest.Filter),
+		trace.Int64Attribute("ldap.scope", int64(searchRequest.Scope)),
+	}, func(conn *ldap.Conn) (cerr error) {
+		result, cerr = conn.Search(searchRequest)
+		return
 	})
 	return result, err
 }
 
 // Modify issues a ModifyRequest to the LDAP server.
 func (p *ConnectionPool) Modify(ctx context.Context, modifyRequest *ldap.ModifyRequest) error {
-	return p.doRequest(ctx, func(conn *ldap.Conn) error {
+	return p.doRequest(ctx, "ldap.Modify", []trace.Attribute{
+		trace.StringAttribute("ldap.dn", modifyRequest.DN),
+	}, func(conn *ldap.Conn) error {
 		return conn.Modify(modifyRequest)
 	})
 }
 
 // Add issues an AddRequest to the LDAP server.
 func (p *ConnectionPool) Add(ctx context.Context, addRequest *ldap.AddRequest) error {
-	return p.doRequest(ctx, func(conn *ldap.Conn) error {
+	return p.doRequest(ctx, "ldap.Add", []trace.Attribute{
+		trace.StringAttribute("ldap.dn", addRequest.DN),
+	}, func(conn *ldap.Conn) error {
 		return conn.Add(addRequest)
 	})
 }
@@ -219,3 +241,16 @@ func isTemporaryLDAPError(err error) bool {
 		return false
 	}
 }
+
+func errorToTraceStatus(err error) trace.Status {
+	switch err {
+	case nil:
+		return trace.Status{Code: trace.StatusCodeOK, Message: "OK"}
+	case context.Canceled:
+		return trace.Status{Code: trace.StatusCodeCancelled, Message: "CANCELED"}
+	case context.DeadlineExceeded:
+		return trace.Status{Code: trace.StatusCodeDeadlineExceeded, Message: "DEADLINE_EXCEEDED"}
+	default:
+		return trace.Status{Code: trace.StatusCodeUnknown, Message: err.Error()}
+	}
+}