Improve with logging deep components

Propogate context with logger to ToxicCollection, ToxicLink.

api.go

@@ -16,6 +16,20 @@ import (
  "github.com/Shopify/toxiproxy/v2/toxics"
 )
 
+func stopBrowsersMiddleware(next http.Handler) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if strings.HasPrefix(r.UserAgent(), "Mozilla/") {
+ http.Error(w, "User agent not allowed", 403)
+ } else {
+ next.ServeHTTP(w, r)
+ }
+ })
+}
+
+func timeoutMiddleware(next http.Handler) http.Handler {
+ return http.TimeoutHandler(next, 30*time.Second, "")
+}
+
 type ApiServer struct {
  Collection *ProxyCollection
  Metrics *metricsContainer
@@ -46,20 +60,6 @@ func (server *ApiServer) PopulateConfig(filename string) {
  }
 }
 
-func stopBrowsersMiddleware(next http.Handler) http.Handler {
- return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- if strings.HasPrefix(r.UserAgent(), "Mozilla/") {
- http.Error(w, "User agent not allowed", 403)
- } else {
- next.ServeHTTP(w, r)
- }
- })
-}
-
-func timeoutMiddleware(next http.Handler) http.Handler {
- return http.TimeoutHandler(next, 5*time.Second, "")
-}
-
 func (server *ApiServer) Listen(host string, port string) {
  r := mux.NewRouter()
  r.Use(hlog.NewHandler(*server.Logger))
@@ -153,6 +153,7 @@ func (server *ApiServer) ProxyIndex(response http.ResponseWriter, request *http.
 }
 
 func (server *ApiServer) ResetState(response http.ResponseWriter, request *http.Request) {
+ ctx := request.Context()
  proxies := server.Collection.Proxies()
 
  for _, proxy := range proxies {
@@ -161,13 +162,13 @@ func (server *ApiServer) ResetState(response http.ResponseWriter, request *http.
  return
  }
 
- proxy.Toxics.ResetToxics()
+ proxy.Toxics.ResetToxics(ctx)
  }
 
  response.WriteHeader(http.StatusNoContent)
  _, err := response.Write(nil)
  if err != nil {
- log := zerolog.Ctx(request.Context())
+ log := zerolog.Ctx(ctx)
  log.Warn().Err(err).Msg("ResetState: Failed to write headers to client")
  }
 }
@@ -414,21 +415,22 @@ func (server *ApiServer) ToxicUpdate(response http.ResponseWriter, request *http
 
 func (server *ApiServer) ToxicDelete(response http.ResponseWriter, request *http.Request) {
  vars := mux.Vars(request)
+ ctx := request.Context()
+ log := zerolog.Ctx(ctx)
 
  proxy, err := server.Collection.Get(vars["proxy"])
  if server.apiError(response, err) {
  return
  }
 
- err = proxy.Toxics.RemoveToxic(vars["toxic"])
+ err = proxy.Toxics.RemoveToxic(ctx, vars["toxic"])
  if server.apiError(response, err) {
  return
  }
 
  response.WriteHeader(http.StatusNoContent)
  _, err = response.Write(nil)
  if err != nil {
- log := zerolog.Ctx(request.Context())
  log.Warn().Err(err).Msg("ToxicDelete: Failed to write headers to client")
  }
 }

link.go

@@ -1,6 +1,8 @@
 package toxiproxy
 
 import (
+ "context"
+ "fmt"
  "io"
  "net"
 
@@ -183,23 +185,36 @@ func (link *ToxicLink) UpdateToxic(toxic *toxics.ToxicWrapper) {
 }
 
 // Remove an existing toxic from the chain.
-func (link *ToxicLink) RemoveToxic(toxic *toxics.ToxicWrapper) {
- i := toxic.Index
+func (link *ToxicLink) RemoveToxic(ctx context.Context, toxic *toxics.ToxicWrapper) {
+ toxic_index := toxic.Index
+ log := zerolog.Ctx(ctx).
+ With().
+ Str("component", "ToxicLink").
+ Str("method", "RemoveToxic").
+ Str("toxic", toxic.Name).
+ Str("toxic_type", toxic.Type).
+ Int("toxic_index", toxic.Index).
+ Str("link_addr", fmt.Sprintf("%p", link)).
+ Str("toxic_stub_addr", fmt.Sprintf("%p", link.stubs[toxic_index])).
+ Str("prev_toxic_stub_addr", fmt.Sprintf("%p", link.stubs[toxic_index-1])).
+ Logger()
 
- if link.stubs[i].InterruptToxic() {
+ if link.stubs[toxic_index].InterruptToxic() {
  cleanup, ok := toxic.Toxic.(toxics.CleanupToxic)
  if ok {
- cleanup.Cleanup(link.stubs[i])
+ cleanup.Cleanup(link.stubs[toxic_index])
  // Cleanup could have closed the stub.
- if link.stubs[i].Closed() {
+ if link.stubs[toxic_index].Closed() {
+ log.Trace().Msg("Cleanup closed toxic and removed toxic")
+ // TODO: Check if cleanup happen would link.stubs recalculated?
  return
  }
  }
 
+ log.Trace().Msg("Interrupt the previous toxic to update its output")
  stop := make(chan bool)
- // Interrupt the previous toxic to update its output
  go func() {
- stop <- link.stubs[i-1].InterruptToxic()
+ stop <- link.stubs[toxic_index-1].InterruptToxic()
  }()
 
  // Unblock the previous toxic if it is trying to flush
@@ -210,32 +225,32 @@ func (link *ToxicLink) RemoveToxic(toxic *toxics.ToxicWrapper) {
  select {
  case interrupted = <-stop:
  stopped = true
- case tmp := <-link.stubs[i].Input:
+ case tmp := <-link.stubs[toxic_index].Input:
  if tmp == nil {
- link.stubs[i].Close()
+ link.stubs[toxic_index].Close()
  if !stopped {
  <-stop
  }
  return
  }
- link.stubs[i].Output <- tmp
+ link.stubs[toxic_index].Output <- tmp
  }
  }
 
  // Empty the toxic's buffer if necessary
- for len(link.stubs[i].Input) > 0 {
- tmp := <-link.stubs[i].Input
+ for len(link.stubs[toxic_index].Input) > 0 {
+ tmp := <-link.stubs[toxic_index].Input
  if tmp == nil {
- link.stubs[i].Close()
+ link.stubs[toxic_index].Close()
  return
  }
- link.stubs[i].Output <- tmp
+ link.stubs[toxic_index].Output <- tmp
  }
 
- link.stubs[i-1].Output = link.stubs[i].Output
- link.stubs = append(link.stubs[:i], link.stubs[i+1:]...)
+ link.stubs[toxic_index-1].Output = link.stubs[toxic_index].Output
+ link.stubs = append(link.stubs[:toxic_index], link.stubs[toxic_index+1:]...)
 
- go link.stubs[i-1].Run(link.toxics.chain[link.direction][i-1])
+ go link.stubs[toxic_index-1].Run(link.toxics.chain[link.direction][toxic_index-1])
  }
 }
 

link_test.go

@@ -1,8 +1,11 @@
 package toxiproxy
 
 import (
+ "context"
  "encoding/binary"
+ "flag"
  "io"
+ "os"
  "testing"
  "time"
 
@@ -81,6 +84,7 @@ func TestStubInitializaationWithToxics(t *testing.T) {
 }
 
 func TestAddRemoveStubs(t *testing.T) {
+ ctx := context.Background()
  collection := NewToxicCollection(nil)
  link := NewToxicLink(nil, collection, stream.Downstream, zerolog.Nop())
  go link.stubs[0].Run(collection.chain[stream.Downstream][0])
@@ -117,7 +121,7 @@ func TestAddRemoveStubs(t *testing.T) {
  }
 
  // Remove stubs
- collection.chainRemoveToxic(toxic)
+ collection.chainRemoveToxic(ctx, toxic)
  if cap(link.stubs[len(link.stubs)-1].Output) != 0 {
  t.Fatalf("Link output buffer was not initialized as 0: %d", cap(link.stubs[0].Output))
  }
@@ -134,6 +138,7 @@ func TestAddRemoveStubs(t *testing.T) {
 }
 
 func TestNoDataDropped(t *testing.T) {
+ ctx := context.Background()
  collection := NewToxicCollection(nil)
  link := NewToxicLink(nil, collection, stream.Downstream, zerolog.Nop())
  go link.stubs[0].Run(collection.chain[stream.Downstream][0])
@@ -159,17 +164,17 @@ func TestNoDataDropped(t *testing.T) {
  }
  link.input.Close()
  }()
- go func() {
+ go func(ctx context.Context) {
  for {
  select {
  case <-done:
  return
  default:
  collection.chainAddToxic(toxic)
- collection.chainRemoveToxic(toxic)
+ collection.chainRemoveToxic(ctx, toxic)
  }
  }
- }()
+ }(ctx)
 
  buf := make([]byte, 2)
  for i := 0; i < 64*1024; i++ {
@@ -238,7 +243,11 @@ func TestToxicity(t *testing.T) {
 
 func TestStateCreated(t *testing.T) {
  collection := NewToxicCollection(nil)
- link := NewToxicLink(nil, collection, stream.Downstream, zerolog.Nop())
+ log := zerolog.Nop()
+ if flag.Lookup("test.v").DefValue == "true" {
+ log = zerolog.New(os.Stdout).With().Caller().Timestamp().Logger()
+ }
+ link := NewToxicLink(nil, collection, stream.Downstream, log)
  go link.stubs[0].Run(collection.chain[stream.Downstream][0])
  collection.links["test"] = link
 

toxic_collection.go

@@ -2,6 +2,7 @@ package toxiproxy
 
 import (
  "bytes"
+ "context"
  "encoding/json"
  "fmt"
  "io"
@@ -43,14 +44,14 @@ func NewToxicCollection(proxy *Proxy) *ToxicCollection {
  return collection
 }
 
-func (c *ToxicCollection) ResetToxics() {
+func (c *ToxicCollection) ResetToxics(ctx context.Context) {
  c.Lock()
  defer c.Unlock()
 
  // Remove all but the first noop toxic
  for dir := range c.chain {
  for len(c.chain[dir]) > 1 {
- c.chainRemoveToxic(c.chain[dir][1])
+ c.chainRemoveToxic(ctx, c.chain[dir][1])
  }
  }
 }
@@ -158,16 +159,28 @@ func (c *ToxicCollection) UpdateToxicJson(
  return nil, ErrToxicNotFound
 }
 
-func (c *ToxicCollection) RemoveToxic(name string) error {
+func (c *ToxicCollection) RemoveToxic(ctx context.Context, name string) error {
+ log := zerolog.Ctx(ctx).
+ With().
+ Str("component", "ToxicCollection").
+ Str("method", "RemoveToxic").
+ Str("toxic", name).
+ Str("proxy", c.proxy.Name).
+ Logger()
+ log.Trace().Msg("Acquire locking...")
  c.Lock()
  defer c.Unlock()
 
+ log.Trace().Msg("Getting toxic by name...")
  toxic := c.findToxicByName(name)
- if toxic != nil {
- c.chainRemoveToxic(toxic)
- return nil
+ if toxic == nil {
+ log.Trace().Msg("Could not find toxic by name")
+ return ErrToxicNotFound
  }
- return ErrToxicNotFound
+
+ c.chainRemoveToxic(ctx, toxic)
+ log.Trace().Msg("Finished")
+ return nil
 }
 
 func (c *ToxicCollection) StartLink(
@@ -217,17 +230,17 @@ func (c *ToxicCollection) chainAddToxic(toxic *toxics.ToxicWrapper) {
  c.chain[dir] = append(c.chain[dir], toxic)
 
  // Asynchronously add the toxic to each link
- group := sync.WaitGroup{}
+ wg := sync.WaitGroup{}
  for _, link := range c.links {
  if link.direction == dir {
- group.Add(1)
- go func(link *ToxicLink) {
- defer group.Done()
+ wg.Add(1)
+ go func(link *ToxicLink, wg *sync.WaitGroup) {
+ defer wg.Done()
  link.AddToxic(toxic)
- }(link)
+ }(link, &wg)
  }
  }
- group.Wait()
+ wg.Wait()
 }
 
 func (c *ToxicCollection) chainUpdateToxic(toxic *toxics.ToxicWrapper) {
@@ -247,25 +260,40 @@ func (c *ToxicCollection) chainUpdateToxic(toxic *toxics.ToxicWrapper) {
  group.Wait()
 }
 
-func (c *ToxicCollection) chainRemoveToxic(toxic *toxics.ToxicWrapper) {
+func (c *ToxicCollection) chainRemoveToxic(ctx context.Context, toxic *toxics.ToxicWrapper) {
+ log := zerolog.Ctx(ctx).
+ With().
+ Str("component", "ToxicCollection").
+ Str("method", "chainRemoveToxic").
+ Str("toxic", toxic.Name).
+ Str("direction", toxic.Direction.String()).
+ Logger()
+
  dir := toxic.Direction
  c.chain[dir] = append(c.chain[dir][:toxic.Index], c.chain[dir][toxic.Index+1:]...)
  for i := toxic.Index; i < len(c.chain[dir]); i++ {
  c.chain[dir][i].Index = i
  }
 
  // Asynchronously remove the toxic from each link
- group := sync.WaitGroup{}
+ wg := sync.WaitGroup{}
+
+ event_array := zerolog.Arr()
  for _, link := range c.links {
  if link.direction == dir {
- group.Add(1)
- go func(link *ToxicLink) {
- defer group.Done()
- link.RemoveToxic(toxic)
- }(link)
+ event_array = event_array.Str(fmt.Sprintf("Link[%p] %s", link, link.Direction()))
+ wg.Add(1)
+ go func(ctx context.Context, link *ToxicLink, log zerolog.Logger) {
+ defer wg.Done()
+ link.RemoveToxic(ctx, toxic)
+ }(ctx, link, log)
  }
  }
- group.Wait()
+
+ log.Trace().
+ Array("links", event_array).
+ Msg("Waiting to update links")
+ wg.Wait()
 
  toxic.Index = -1
 }