context.go 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. // Copyright 2012 The Gorilla Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package context
  5. import (
  6. "net/http"
  7. "sync"
  8. "time"
  9. )
  10. var (
  11. mutex sync.Mutex
  12. data = make(map[*http.Request]map[interface{}]interface{})
  13. datat = make(map[*http.Request]int64)
  14. )
  15. // Set stores a value for a given key in a given request.
  16. func Set(r *http.Request, key, val interface{}) {
  17. mutex.Lock()
  18. defer mutex.Unlock()
  19. if data[r] == nil {
  20. data[r] = make(map[interface{}]interface{})
  21. datat[r] = time.Now().Unix()
  22. }
  23. data[r][key] = val
  24. }
  25. // Get returns a value stored for a given key in a given request.
  26. func Get(r *http.Request, key interface{}) interface{} {
  27. mutex.Lock()
  28. defer mutex.Unlock()
  29. if data[r] != nil {
  30. return data[r][key]
  31. }
  32. return nil
  33. }
  34. // GetOk returns stored value and presence state like multi-value return of map access.
  35. func GetOk(r *http.Request, key interface{}) (interface{}, bool) {
  36. mutex.Lock()
  37. defer mutex.Unlock()
  38. if _, ok := data[r]; ok {
  39. value, ok := data[r][key]
  40. return value, ok
  41. }
  42. return nil, false
  43. }
  44. // Delete removes a value stored for a given key in a given request.
  45. func Delete(r *http.Request, key interface{}) {
  46. mutex.Lock()
  47. defer mutex.Unlock()
  48. if data[r] != nil {
  49. delete(data[r], key)
  50. }
  51. }
  52. // Clear removes all values stored for a given request.
  53. //
  54. // This is usually called by a handler wrapper to clean up request
  55. // variables at the end of a request lifetime. See ClearHandler().
  56. func Clear(r *http.Request) {
  57. mutex.Lock()
  58. defer mutex.Unlock()
  59. clear(r)
  60. }
  61. // clear is Clear without the lock.
  62. func clear(r *http.Request) {
  63. delete(data, r)
  64. delete(datat, r)
  65. }
  66. // Purge removes request data stored for longer than maxAge, in seconds.
  67. // It returns the amount of requests removed.
  68. //
  69. // If maxAge <= 0, all request data is removed.
  70. //
  71. // This is only used for sanity check: in case context cleaning was not
  72. // properly set some request data can be kept forever, consuming an increasing
  73. // amount of memory. In case this is detected, Purge() must be called
  74. // periodically until the problem is fixed.
  75. func Purge(maxAge int) int {
  76. mutex.Lock()
  77. defer mutex.Unlock()
  78. count := 0
  79. if maxAge <= 0 {
  80. count = len(data)
  81. data = make(map[*http.Request]map[interface{}]interface{})
  82. datat = make(map[*http.Request]int64)
  83. } else {
  84. min := time.Now().Unix() - int64(maxAge)
  85. for r := range data {
  86. if datat[r] < min {
  87. clear(r)
  88. count++
  89. }
  90. }
  91. }
  92. return count
  93. }
  94. // ClearHandler wraps an http.Handler and clears request values at the end
  95. // of a request lifetime.
  96. func ClearHandler(h http.Handler) http.Handler {
  97. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  98. defer Clear(r)
  99. h.ServeHTTP(w, r)
  100. })
  101. }