write_buffer.go 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. package proto
  2. import (
  3. "encoding"
  4. "fmt"
  5. "strconv"
  6. )
  7. const bufferSize = 4096
  8. type WriteBuffer struct {
  9. b []byte
  10. }
  11. func NewWriteBuffer() *WriteBuffer {
  12. return &WriteBuffer{
  13. b: make([]byte, 0, 4096),
  14. }
  15. }
  16. func (w *WriteBuffer) Len() int { return len(w.b) }
  17. func (w *WriteBuffer) Bytes() []byte { return w.b }
  18. func (w *WriteBuffer) Reset() { w.b = w.b[:0] }
  19. func (w *WriteBuffer) Append(args []interface{}) error {
  20. w.b = append(w.b, ArrayReply)
  21. w.b = strconv.AppendUint(w.b, uint64(len(args)), 10)
  22. w.b = append(w.b, '\r', '\n')
  23. for _, arg := range args {
  24. if err := w.append(arg); err != nil {
  25. return err
  26. }
  27. }
  28. return nil
  29. }
  30. func (w *WriteBuffer) append(val interface{}) error {
  31. switch v := val.(type) {
  32. case nil:
  33. w.AppendString("")
  34. case string:
  35. w.AppendString(v)
  36. case []byte:
  37. w.AppendBytes(v)
  38. case int:
  39. w.AppendString(formatInt(int64(v)))
  40. case int8:
  41. w.AppendString(formatInt(int64(v)))
  42. case int16:
  43. w.AppendString(formatInt(int64(v)))
  44. case int32:
  45. w.AppendString(formatInt(int64(v)))
  46. case int64:
  47. w.AppendString(formatInt(v))
  48. case uint:
  49. w.AppendString(formatUint(uint64(v)))
  50. case uint8:
  51. w.AppendString(formatUint(uint64(v)))
  52. case uint16:
  53. w.AppendString(formatUint(uint64(v)))
  54. case uint32:
  55. w.AppendString(formatUint(uint64(v)))
  56. case uint64:
  57. w.AppendString(formatUint(v))
  58. case float32:
  59. w.AppendString(formatFloat(float64(v)))
  60. case float64:
  61. w.AppendString(formatFloat(v))
  62. case bool:
  63. if v {
  64. w.AppendString("1")
  65. } else {
  66. w.AppendString("0")
  67. }
  68. default:
  69. if bm, ok := val.(encoding.BinaryMarshaler); ok {
  70. bb, err := bm.MarshalBinary()
  71. if err != nil {
  72. return err
  73. }
  74. w.AppendBytes(bb)
  75. } else {
  76. return fmt.Errorf(
  77. "redis: can't marshal %T (consider implementing encoding.BinaryMarshaler)", val)
  78. }
  79. }
  80. return nil
  81. }
  82. func (w *WriteBuffer) AppendString(s string) {
  83. w.b = append(w.b, StringReply)
  84. w.b = strconv.AppendUint(w.b, uint64(len(s)), 10)
  85. w.b = append(w.b, '\r', '\n')
  86. w.b = append(w.b, s...)
  87. w.b = append(w.b, '\r', '\n')
  88. }
  89. func (w *WriteBuffer) AppendBytes(p []byte) {
  90. w.b = append(w.b, StringReply)
  91. w.b = strconv.AppendUint(w.b, uint64(len(p)), 10)
  92. w.b = append(w.b, '\r', '\n')
  93. w.b = append(w.b, p...)
  94. w.b = append(w.b, '\r', '\n')
  95. }