warnings.go 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. /*
  2. Copyright 2020 The Kubernetes Authors.
  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. http://www.apache.org/licenses/LICENSE-2.0
  7. Unless required by applicable law or agreed to in writing, software
  8. distributed under the License is distributed on an "AS IS" BASIS,
  9. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. See the License for the specific language governing permissions and
  11. limitations under the License.
  12. */
  13. package rest
  14. import (
  15. "fmt"
  16. "io"
  17. "net/http"
  18. "sync"
  19. "k8s.io/klog/v2"
  20. "k8s.io/apimachinery/pkg/util/net"
  21. )
  22. // WarningHandler is an interface for handling warning headers
  23. type WarningHandler interface {
  24. // HandleWarningHeader is called with the warn code, agent, and text when a warning header is countered.
  25. HandleWarningHeader(code int, agent string, text string)
  26. }
  27. var (
  28. defaultWarningHandler WarningHandler = WarningLogger{}
  29. defaultWarningHandlerLock sync.RWMutex
  30. )
  31. // SetDefaultWarningHandler sets the default handler clients use when warning headers are encountered.
  32. // By default, warnings are logged. Several built-in implementations are provided:
  33. // - NoWarnings suppresses warnings.
  34. // - WarningLogger logs warnings.
  35. // - NewWarningWriter() outputs warnings to the provided writer.
  36. func SetDefaultWarningHandler(l WarningHandler) {
  37. defaultWarningHandlerLock.Lock()
  38. defer defaultWarningHandlerLock.Unlock()
  39. defaultWarningHandler = l
  40. }
  41. func getDefaultWarningHandler() WarningHandler {
  42. defaultWarningHandlerLock.RLock()
  43. defer defaultWarningHandlerLock.RUnlock()
  44. l := defaultWarningHandler
  45. return l
  46. }
  47. // NoWarnings is an implementation of WarningHandler that suppresses warnings.
  48. type NoWarnings struct{}
  49. func (NoWarnings) HandleWarningHeader(code int, agent string, message string) {}
  50. // WarningLogger is an implementation of WarningHandler that logs code 299 warnings
  51. type WarningLogger struct{}
  52. func (WarningLogger) HandleWarningHeader(code int, agent string, message string) {
  53. if code != 299 || len(message) == 0 {
  54. return
  55. }
  56. klog.Warning(message)
  57. }
  58. type warningWriter struct {
  59. // out is the writer to output warnings to
  60. out io.Writer
  61. // opts contains options controlling warning output
  62. opts WarningWriterOptions
  63. // writtenLock guards written and writtenCount
  64. writtenLock sync.Mutex
  65. writtenCount int
  66. written map[string]struct{}
  67. }
  68. // WarningWriterOptions controls the behavior of a WarningHandler constructed using NewWarningWriter()
  69. type WarningWriterOptions struct {
  70. // Deduplicate indicates a given warning message should only be written once.
  71. // Setting this to true in a long-running process handling many warnings can result in increased memory use.
  72. Deduplicate bool
  73. // Color indicates that warning output can include ANSI color codes
  74. Color bool
  75. }
  76. // NewWarningWriter returns an implementation of WarningHandler that outputs code 299 warnings to the specified writer.
  77. func NewWarningWriter(out io.Writer, opts WarningWriterOptions) *warningWriter {
  78. h := &warningWriter{out: out, opts: opts}
  79. if opts.Deduplicate {
  80. h.written = map[string]struct{}{}
  81. }
  82. return h
  83. }
  84. const (
  85. yellowColor = "\u001b[33;1m"
  86. resetColor = "\u001b[0m"
  87. )
  88. // HandleWarningHeader prints warnings with code=299 to the configured writer.
  89. func (w *warningWriter) HandleWarningHeader(code int, agent string, message string) {
  90. if code != 299 || len(message) == 0 {
  91. return
  92. }
  93. w.writtenLock.Lock()
  94. defer w.writtenLock.Unlock()
  95. if w.opts.Deduplicate {
  96. if _, alreadyWritten := w.written[message]; alreadyWritten {
  97. return
  98. }
  99. w.written[message] = struct{}{}
  100. }
  101. w.writtenCount++
  102. if w.opts.Color {
  103. fmt.Fprintf(w.out, "%sWarning:%s %s\n", yellowColor, resetColor, message)
  104. } else {
  105. fmt.Fprintf(w.out, "Warning: %s\n", message)
  106. }
  107. }
  108. func (w *warningWriter) WarningCount() int {
  109. w.writtenLock.Lock()
  110. defer w.writtenLock.Unlock()
  111. return w.writtenCount
  112. }
  113. func handleWarnings(headers http.Header, handler WarningHandler) []net.WarningHeader {
  114. if handler == nil {
  115. handler = getDefaultWarningHandler()
  116. }
  117. warnings, _ := net.ParseWarningHeaders(headers["Warning"])
  118. for _, warning := range warnings {
  119. handler.HandleWarningHeader(warning.Code, warning.Agent, warning.Text)
  120. }
  121. return warnings
  122. }