GODRIVER-2045 Propose an example of custom dialer with DNS cache. (#1088)

Propose an example of custom dialer with DNS cache.

Author: qingyang-hu

examples/example_customdns_test.go

@@ -0,0 +1,138 @@
+// Copyright (C) MongoDB, Inc. 2022-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package examples
+
+import (
+	"context"
+	"log"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/miekg/dns"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+func resolve(ctx context.Context, cache *dnsCache, in *dns.Conn, out *dns.Conn) {
+	for ctx.Err() == nil {
+		q, err := in.ReadMsg()
+		if err != nil {
+			// TODO: Handle error.
+			log.Fatalf("Unhandled error in ReadMsg: %v", err)
+		}
+		if len(q.Question) != 1 {
+			// Multiple questions in a single query is not actually used in real life.
+			continue
+		}
+
+		a, err := func() (*dns.Msg, error) {
+			cache.lock.Lock()
+			defer cache.lock.Unlock()
+
+			now := time.Now()
+			if rr, ok := cache.records[q.Question[0]]; ok && rr.exp.After(now) {
+				a := new(dns.Msg)
+				a.SetReply(q)
+				a.Compress = false
+				a.Answer = append(a.Answer, rr.record)
+				return a, nil
+			}
+
+			err := out.WriteMsg(q)
+			if err != nil {
+				return nil, err
+			}
+
+			m, err := out.ReadMsg()
+			if err != nil {
+				return nil, err
+			}
+
+			l := len(m.Answer)
+			for i, q := range m.Question {
+				if i >= l {
+					break
+				}
+				a := m.Answer[i]
+				cache.records[q] = &RR{
+					a,
+					now.Add(time.Second * time.Duration(a.Header().Ttl)),
+				}
+			}
+			return m, nil
+		}()
+		if err != nil {
+			// TODO: Handle error.
+			log.Fatalf("Unhandled error in record retrieval: %v", err)
+		}
+
+		if err := in.WriteMsg(a); err != nil {
+			// TODO: Handle error.
+			log.Fatalf("Unhandled error in WriteMsg: %v", err)
+		}
+	}
+}
+
+type RR struct {
+	record dns.RR
+	exp    time.Time
+}
+
+type dnsCache struct {
+	records map[dns.Question]*RR
+	lock    sync.Mutex
+}
+
+type dialer struct {
+	*net.Dialer
+	cache *dnsCache
+}
+
+func NewDialer() dialer {
+	cache := &dnsCache{
+		records: make(map[dns.Question]*RR),
+		lock:    sync.Mutex{},
+	}
+	return dialer{
+		Dialer: &net.Dialer{
+			Resolver: &net.Resolver{
+				PreferGo: true,
+				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+					var d net.Dialer
+					outConn, err := d.DialContext(ctx, network, address)
+					conn, inConn := net.Pipe()
+					if err == nil {
+						go resolve(ctx, cache, &dns.Conn{Conn: inConn}, &dns.Conn{Conn: outConn})
+					}
+					return conn, err
+				},
+			},
+		},
+		cache: cache,
+	}
+}
+
+func TestCustomDialer(t *testing.T) {
+	client, err := mongo.NewClient(options.Client().ApplyURI("mongodb://testurl:27017").SetDialer(NewDialer()))
+	if err != nil {
+		t.Fatalf("error creating client: %v", err)
+	}
+	ctx := context.Background()
+	err = client.Connect(ctx)
+	if err != nil {
+		t.Fatalf("error connecting: %v", err)
+	}
+	defer client.Disconnect(context.Background())
+	coll := client.Database("test").Collection("test")
+	_, err = coll.InsertOne(context.Background(), bson.D{{"text", "text"}})
+	if err != nil {
+		t.Fatalf("error inserting: %v", err)
+	}
+}