write_syncer.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. // Copyright (c) 2016 Uber Technologies, Inc.
  2. //
  3. // Permission is hereby granted, free of charge, to any person obtaining a copy
  4. // of this software and associated documentation files (the "Software"), to deal
  5. // in the Software without restriction, including without limitation the rights
  6. // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  7. // copies of the Software, and to permit persons to whom the Software is
  8. // furnished to do so, subject to the following conditions:
  9. //
  10. // The above copyright notice and this permission notice shall be included in
  11. // all copies or substantial portions of the Software.
  12. //
  13. // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  14. // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  15. // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  16. // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  17. // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  18. // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  19. // THE SOFTWARE.
  20. package zapcore
  21. import (
  22. "io"
  23. "sync"
  24. "go.uber.org/multierr"
  25. )
  26. // A WriteSyncer is an io.Writer that can also flush any buffered data. Note
  27. // that *os.File (and thus, os.Stderr and os.Stdout) implement WriteSyncer.
  28. type WriteSyncer interface {
  29. io.Writer
  30. Sync() error
  31. }
  32. // AddSync converts an io.Writer to a WriteSyncer. It attempts to be
  33. // intelligent: if the concrete type of the io.Writer implements WriteSyncer,
  34. // we'll use the existing Sync method. If it doesn't, we'll add a no-op Sync.
  35. func AddSync(w io.Writer) WriteSyncer {
  36. switch w := w.(type) {
  37. case WriteSyncer:
  38. return w
  39. default:
  40. return writerWrapper{w}
  41. }
  42. }
  43. type lockedWriteSyncer struct {
  44. sync.Mutex
  45. ws WriteSyncer
  46. }
  47. // Lock wraps a WriteSyncer in a mutex to make it safe for concurrent use. In
  48. // particular, *os.Files must be locked before use.
  49. func Lock(ws WriteSyncer) WriteSyncer {
  50. if _, ok := ws.(*lockedWriteSyncer); ok {
  51. // no need to layer on another lock
  52. return ws
  53. }
  54. return &lockedWriteSyncer{ws: ws}
  55. }
  56. func (s *lockedWriteSyncer) Write(bs []byte) (int, error) {
  57. s.Lock()
  58. n, err := s.ws.Write(bs)
  59. s.Unlock()
  60. return n, err
  61. }
  62. func (s *lockedWriteSyncer) Sync() error {
  63. s.Lock()
  64. err := s.ws.Sync()
  65. s.Unlock()
  66. return err
  67. }
  68. type writerWrapper struct {
  69. io.Writer
  70. }
  71. func (w writerWrapper) Sync() error {
  72. return nil
  73. }
  74. type multiWriteSyncer []WriteSyncer
  75. // NewMultiWriteSyncer creates a WriteSyncer that duplicates its writes
  76. // and sync calls, much like io.MultiWriter.
  77. func NewMultiWriteSyncer(ws ...WriteSyncer) WriteSyncer {
  78. if len(ws) == 1 {
  79. return ws[0]
  80. }
  81. // Copy to protect against https://github.com/golang/go/issues/7809
  82. return multiWriteSyncer(append([]WriteSyncer(nil), ws...))
  83. }
  84. // See https://golang.org/src/io/multi.go
  85. // When not all underlying syncers write the same number of bytes,
  86. // the smallest number is returned even though Write() is called on
  87. // all of them.
  88. func (ws multiWriteSyncer) Write(p []byte) (int, error) {
  89. var writeErr error
  90. nWritten := 0
  91. for _, w := range ws {
  92. n, err := w.Write(p)
  93. writeErr = multierr.Append(writeErr, err)
  94. if nWritten == 0 && n != 0 {
  95. nWritten = n
  96. } else if n < nWritten {
  97. nWritten = n
  98. }
  99. }
  100. return nWritten, writeErr
  101. }
  102. func (ws multiWriteSyncer) Sync() error {
  103. var err error
  104. for _, w := range ws {
  105. err = multierr.Append(err, w.Sync())
  106. }
  107. return err
  108. }