From d8960afcfe1ca8c4fe10a26491de6c2e7522f80d Mon Sep 17 00:00:00 2001 From: Irine Sistiana <49315432+IrineSistiana@users.noreply.github.com> Date: Wed, 1 Nov 2023 11:03:28 +0800 Subject: [PATCH] dual_selector: impl cache --- plugin/executable/dual_selector/cache.go | 11 +++ .../executable/dual_selector/dual_selector.go | 89 +++++++++++++------ .../dual_selector/dual_selector_test.go | 5 +- 3 files changed, 75 insertions(+), 30 deletions(-) create mode 100644 plugin/executable/dual_selector/cache.go diff --git a/plugin/executable/dual_selector/cache.go b/plugin/executable/dual_selector/cache.go new file mode 100644 index 000000000..f22cc76fb --- /dev/null +++ b/plugin/executable/dual_selector/cache.go @@ -0,0 +1,11 @@ +package dual_selector + +import "hash/maphash" + +type key string + +var seed = maphash.MakeSeed() + +func (k key) Sum() uint64 { + return maphash.String(seed, string(k)) +} diff --git a/plugin/executable/dual_selector/dual_selector.go b/plugin/executable/dual_selector/dual_selector.go index 29a34f2c4..86dff3319 100644 --- a/plugin/executable/dual_selector/dual_selector.go +++ b/plugin/executable/dual_selector/dual_selector.go @@ -21,8 +21,10 @@ package dual_selector import ( "context" + "io" "time" + "github.com/IrineSistiana/mosdns/v5/pkg/cache" "github.com/IrineSistiana/mosdns/v5/pkg/dnsutils" "github.com/IrineSistiana/mosdns/v5/pkg/pool" "github.com/IrineSistiana/mosdns/v5/pkg/query_context" @@ -34,6 +36,11 @@ import ( const ( referenceWaitTimeout = time.Millisecond * 500 defaultSubRoutineTimeout = time.Second * 5 + + // TODO: Make cache configurable? + cacheSize = 64 * 1024 + cacheTlt = time.Hour + cacheGcInterval = time.Minute ) func init() { @@ -46,10 +53,13 @@ func init() { } var _ sequence.RecursiveExecutable = (*Selector)(nil) +var _ io.Closer = (*Selector)(nil) type Selector struct { sequence.BQ prefer uint16 // dns.TypeA or dns.TypeAAAA + + preferTypOkCache *cache.Cache[key, bool] } // Exec implements handler.Executable. @@ -60,30 +70,47 @@ func (s *Selector) Exec(ctx context.Context, qCtx *query_context.Context, next s } qtype := q.Question[0].Qtype - // skip queries that have preferred type or have other unrelated types. - if qtype == s.prefer || (qtype != dns.TypeA && qtype != dns.TypeAAAA) { + // skip queries that have other unrelated types. + if qtype != dns.TypeA && qtype != dns.TypeAAAA { return next.ExecNext(ctx, qCtx) } - // start reference goroutine - qCtxRef := qCtx.Copy() - var refQtype uint16 - if qtype == dns.TypeA { - refQtype = dns.TypeAAAA - } else { - refQtype = dns.TypeA + qName := key(q.Question[0].Name) + if qtype == s.prefer { + err := next.ExecNext(ctx, qCtx) + if err != nil { + return err + } + + if r := qCtx.R(); r != nil && msgAnsHasRR(r, s.prefer) { + s.preferTypOkCache.Store(qName, true, time.Now().Add(cacheTlt)) + } + return nil + } + + // Qtype is not the preferred type. + preferredTypOk, _, _ := s.preferTypOkCache.Get(qName) + if preferredTypOk { + // We know that domain has preferred type so this qtype can be blocked + // right away. + r := dnsutils.GenEmptyReply(q, dns.RcodeSuccess) + qCtx.SetResponse(r) + return nil } - qCtxRef.Q().Question[0].Qtype = refQtype - ddl, ok := ctx.Deadline() - if !ok { + // async check whether domain has the preferred type + qCtxPreferred := qCtx.Copy() + qCtxPreferred.Q().Question[0].Qtype = s.prefer + + ddl, cacheOk := ctx.Deadline() + if !cacheOk { ddl = time.Now().Add(defaultSubRoutineTimeout) } - shouldBlock := make(chan struct{}, 0) - shouldPass := make(chan struct{}, 0) + shouldBlock := make(chan struct{}) + shouldPass := make(chan struct{}) go func() { - qCtx := qCtxRef + qCtx := qCtxPreferred ctx, cancel := context.WithDeadline(context.Background(), ddl) defer cancel() err := next.ExecNext(ctx, qCtx) @@ -92,13 +119,13 @@ func (s *Selector) Exec(ctx context.Context, qCtx *query_context.Context, next s close(shouldPass) return } - if r := qCtx.R(); r != nil && msgAnsHasRR(r, refQtype) { - // Target domain has reference type. + if r := qCtx.R(); r != nil && msgAnsHasRR(r, s.prefer) { + // Target domain has preferred type. + s.preferTypOkCache.Store(qName, true, time.Now().Add(cacheTlt)) close(shouldBlock) return } close(shouldPass) - return }() // start original query goroutine @@ -114,11 +141,11 @@ func (s *Selector) Exec(ctx context.Context, qCtx *query_context.Context, next s select { case <-ctx.Done(): return context.Cause(ctx) - case <-shouldBlock: // Reference indicates we should block this query before the original query finished. + case <-shouldBlock: // Domain has preferred type. Block this type now. r := dnsutils.GenEmptyReply(q, dns.RcodeSuccess) qCtx.SetResponse(r) return nil - case err := <-doneChan: // The original query finished. Waiting for reference. + case err := <-doneChan: // The original query finished. Waiting for preferred type check. waitTimeoutTimer := pool.GetTimer(referenceWaitTimeout) defer pool.ReleaseTimer(waitTimeoutTimer) select { @@ -140,17 +167,27 @@ func (s *Selector) Exec(ctx context.Context, qCtx *query_context.Context, next s } } +func (s *Selector) Close() error { + s.preferTypOkCache.Close() + return nil +} + func NewPreferIpv4(bq sequence.BQ) *Selector { - return &Selector{ - BQ: bq, - prefer: dns.TypeA, - } + return newSelector(bq, dns.TypeA) } func NewPreferIpv6(bq sequence.BQ) *Selector { + return newSelector(bq, dns.TypeAAAA) +} + +func newSelector(bq sequence.BQ, preferType uint16) *Selector { + if preferType != dns.TypeA && preferType != dns.TypeAAAA { + panic("dual_selector: invalid dns qtype") + } return &Selector{ - BQ: bq, - prefer: dns.TypeAAAA, + BQ: bq, + prefer: preferType, + preferTypOkCache: cache.New[key, bool](cache.Opts{Size: cacheSize, CleanerInterval: cacheGcInterval}), } } diff --git a/plugin/executable/dual_selector/dual_selector_test.go b/plugin/executable/dual_selector/dual_selector_test.go index 95196cb52..9616e7310 100644 --- a/plugin/executable/dual_selector/dual_selector_test.go +++ b/plugin/executable/dual_selector/dual_selector_test.go @@ -152,10 +152,7 @@ func TestSelector_Exec(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - s := &Selector{ - BQ: sequence.NewBQ(coremain.NewTestMosdnsWithPlugins(nil), zap.NewNop()), - prefer: tt.prefer, - } + s := newSelector(sequence.NewBQ(coremain.NewTestMosdnsWithPlugins(nil), zap.NewNop()), tt.prefer) q := new(dns.Msg) q.SetQuestion("example.", tt.qtype)