writer.go 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
  1. /*
  2. Copyright 2020 Docker Compose CLI 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 progress
  14. import (
  15. "context"
  16. "os"
  17. "sync"
  18. "github.com/containerd/console"
  19. "github.com/moby/term"
  20. "golang.org/x/sync/errgroup"
  21. )
  22. // Writer can write multiple progress events
  23. type Writer interface {
  24. Start(context.Context) error
  25. Stop()
  26. Event(Event)
  27. Events([]Event)
  28. TailMsgf(string, ...interface{})
  29. }
  30. type writerKey struct{}
  31. // WithContextWriter adds the writer to the context
  32. func WithContextWriter(ctx context.Context, writer Writer) context.Context {
  33. return context.WithValue(ctx, writerKey{}, writer)
  34. }
  35. // ContextWriter returns the writer from the context
  36. func ContextWriter(ctx context.Context) Writer {
  37. s, ok := ctx.Value(writerKey{}).(Writer)
  38. if !ok {
  39. return &noopWriter{}
  40. }
  41. return s
  42. }
  43. type progressFunc func(context.Context) error
  44. type progressFuncWithStatus func(context.Context) (string, error)
  45. // Run will run a writer and the progress function in parallel
  46. func Run(ctx context.Context, pf progressFunc) error {
  47. _, err := RunWithStatus(ctx, func(ctx context.Context) (string, error) {
  48. return "", pf(ctx)
  49. })
  50. return err
  51. }
  52. // RunWithStatus will run a writer and the progress function in parallel and return a status
  53. func RunWithStatus(ctx context.Context, pf progressFuncWithStatus) (string, error) {
  54. eg, _ := errgroup.WithContext(ctx)
  55. w, err := NewWriter(os.Stderr)
  56. var result string
  57. if err != nil {
  58. return "", err
  59. }
  60. eg.Go(func() error {
  61. return w.Start(context.Background())
  62. })
  63. ctx = WithContextWriter(ctx, w)
  64. eg.Go(func() error {
  65. defer w.Stop()
  66. s, err := pf(ctx)
  67. if err == nil {
  68. result = s
  69. }
  70. return err
  71. })
  72. err = eg.Wait()
  73. return result, err
  74. }
  75. const (
  76. // ModeAuto detect console capabilities
  77. ModeAuto = "auto"
  78. // ModeTTY use terminal capability for advanced rendering
  79. ModeTTY = "tty"
  80. // ModePlain dump raw events to output
  81. ModePlain = "plain"
  82. )
  83. // Mode define how progress should be rendered, either as ModePlain or ModeTTY
  84. var Mode = ModeAuto
  85. // NewWriter returns a new multi-progress writer
  86. func NewWriter(out console.File) (Writer, error) {
  87. _, isTerminal := term.GetFdInfo(out)
  88. if Mode == ModeAuto && isTerminal {
  89. return newTTYWriter(out)
  90. }
  91. if Mode == ModeTTY {
  92. return newTTYWriter(out)
  93. }
  94. return &plainWriter{
  95. out: out,
  96. done: make(chan bool),
  97. }, nil
  98. }
  99. func newTTYWriter(out console.File) (Writer, error) {
  100. con, err := console.ConsoleFromFile(out)
  101. if err != nil {
  102. return nil, err
  103. }
  104. return &ttyWriter{
  105. out: con,
  106. eventIDs: []string{},
  107. events: map[string]Event{},
  108. repeated: false,
  109. done: make(chan bool),
  110. mtx: &sync.Mutex{},
  111. }, nil
  112. }