github.com/livekit/protocol@v1.39.3/logger/zaputil/deferrer.go (about)

     1  // Copyright 2023 LiveKit, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package zaputil
    16  
    17  import (
    18  	"slices"
    19  	"sync"
    20  
    21  	"go.uber.org/atomic"
    22  	"go.uber.org/zap"
    23  	"go.uber.org/zap/zapcore"
    24  )
    25  
    26  type deferredWrite struct {
    27  	core   zapcore.Core
    28  	ent    zapcore.Entry
    29  	fields []zapcore.Field
    30  }
    31  
    32  // ---------------------------------
    33  
    34  type Deferrer struct {
    35  	mu     sync.Mutex
    36  	fields atomic.Pointer[[]zapcore.Field]
    37  	writes []*deferredWrite
    38  }
    39  
    40  func (b *Deferrer) buffer(core zapcore.Core, ent zapcore.Entry, fields []zapcore.Field) bool {
    41  	b.mu.Lock()
    42  	defer b.mu.Unlock()
    43  	if b.fields.Load() != nil {
    44  		return false
    45  	}
    46  	b.writes = append(b.writes, &deferredWrite{core, ent, fields})
    47  	return true
    48  }
    49  
    50  func (b *Deferrer) flush() {
    51  	b.mu.Lock()
    52  	writes := b.writes
    53  	b.writes = nil
    54  	b.mu.Unlock()
    55  
    56  	fields := slices.Clone(*b.fields.Load())
    57  	n := len(fields)
    58  
    59  	for _, w := range writes {
    60  		fields = append(fields[:n], w.fields...)
    61  		w.core.Write(w.ent, fields)
    62  	}
    63  }
    64  
    65  func (b *Deferrer) write(core zapcore.Core, ent zapcore.Entry, fields []zapcore.Field) error {
    66  	for {
    67  		if dfs := b.fields.Load(); dfs != nil {
    68  			return core.Write(ent, slices.Concat(fields, *dfs))
    69  		}
    70  		if b.buffer(core, ent, fields) {
    71  			return nil
    72  		}
    73  	}
    74  }
    75  
    76  func (b *Deferrer) Resolve(args ...any) {
    77  	fields := make([]zapcore.Field, len(args))
    78  	for {
    79  		fields = fields[:0]
    80  		for i := 0; i < len(args); i++ {
    81  			switch arg := args[i].(type) {
    82  			case zapcore.Field:
    83  				fields = append(fields, arg)
    84  			case string:
    85  				if i < len(args)-1 {
    86  					fields = append(fields, zap.Any(arg, args[i+1]))
    87  					i++
    88  				}
    89  			}
    90  		}
    91  
    92  		prev := b.fields.Load()
    93  		if prev != nil {
    94  			for _, pf := range *prev {
    95  				overwritten := slices.ContainsFunc(fields, func(f zapcore.Field) bool {
    96  					return f.Key == pf.Key
    97  				})
    98  				if !overwritten {
    99  					fields = append(fields, pf)
   100  				}
   101  			}
   102  		}
   103  
   104  		if b.fields.CompareAndSwap(prev, &fields) {
   105  			if prev == nil {
   106  				b.flush()
   107  			}
   108  			return
   109  		}
   110  	}
   111  }
   112  
   113  func (b *Deferrer) Reset() {
   114  	b.fields.Store(nil)
   115  }
   116  
   117  // ---------------------------------
   118  
   119  type NoOpDeferrer struct{}
   120  
   121  func (n NoOpDeferrer) Resolve(args ...any) {}
   122  
   123  func (n NoOpDeferrer) Reset() {}
   124  
   125  // ---------------------------------
   126  
   127  type DeferredFieldResolver interface {
   128  	Resolve(args ...any)
   129  	Reset()
   130  }
   131  
   132  type deferredValueCore struct {
   133  	zapcore.Core
   134  	def *Deferrer
   135  }
   136  
   137  func NewDeferredValueCore(core zapcore.Core, def *Deferrer) zapcore.Core {
   138  	return &deferredValueCore{core, def}
   139  }
   140  
   141  func (c *deferredValueCore) With(fields []zapcore.Field) zapcore.Core {
   142  	return &deferredValueCore{
   143  		Core: c.Core.With(fields),
   144  		def:  c.def,
   145  	}
   146  }
   147  
   148  func (c *deferredValueCore) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
   149  	if c.Enabled(ent.Level) {
   150  		return ce.AddCore(ent, c)
   151  	}
   152  	return ce
   153  }
   154  
   155  func (c *deferredValueCore) Write(ent zapcore.Entry, fields []zapcore.Field) error {
   156  	return c.def.write(c.Core, ent, fields)
   157  }