Updating p4runtime solution (#171)

exercises/p4runtime/solution/mycontroller.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python2
 import argparse
+import grpc
 import os
 import sys
 from time import sleep
@@ -10,6 +11,7 @@ sys.path.append(
     os.path.join(os.path.dirname(os.path.abspath(__file__)),
                  '../../utils/'))
 import p4runtime_lib.bmv2
+from p4runtime_lib.switch import ShutdownAllSwitchConnections
 import p4runtime_lib.helper
 
 SWITCH_TO_HOST_PORT = 1
@@ -109,7 +111,7 @@ def readTableRules(p4info_helper, sw):
         for entity in response.entities:
             entry = entity.table_entry
             # TODO For extra credit, you can use the p4info_helper to translate
-            #      the IDs the entry to names
+            #      the IDs in the entry to names
             table_name = p4info_helper.get_tables_name(entry.table_id)
             print '%s: ' % table_name,
             for m in entry.match:
@@ -123,7 +125,6 @@ def readTableRules(p4info_helper, sw):
                 print '%r' % p.value,
             print
 
-
 def printCounter(p4info_helper, sw, counter_name, index):
     """
     Reads the specified counter at the specified index from the switch. In our
@@ -143,11 +144,18 @@ def printCounter(p4info_helper, sw, counter_name, index):
                 counter.data.packet_count, counter.data.byte_count
             )
 
+def printGrpcError(e):
+    print "gRPC Error:", e.details(),
+    status_code = e.code()
+    print "(%s)" % status_code.name,
+    traceback = sys.exc_info()[2]
+    print "[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno)
 
 def main(p4info_file_path, bmv2_file_path):
     # Instantiate a P4Runtime helper from the p4info file
     p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)
 
+    try:
         # Create a switch connection object for s1 and s2;
         # this is backed by a P4Runtime gRPC connection.
         # Also, dump all P4Runtime messages sent to switch to given txt files.
@@ -188,7 +196,6 @@ def main(p4info_file_path, bmv2_file_path):
         readTableRules(p4info_helper, s2)
 
         # Print the tunnel counters every 2 seconds
-    try:
         while True:
             sleep(2)
             print '\n----- Reading tunnel counters -----'
@@ -196,9 +203,13 @@ def main(p4info_file_path, bmv2_file_path):
             printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
             printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 200)
             printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 200)
+
     except KeyboardInterrupt:
         print " Shutting down."
+    except grpc.RpcError as e:
+        printGrpcError(e)
 
+    ShutdownAllSwitchConnections()
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='P4Runtime Controller')
@@ -218,5 +229,4 @@ if __name__ == '__main__':
         parser.print_help()
         print "\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json
         parser.exit(1)
-
     main(args.p4info, args.bmv2_json)